From 2ef51f9bd864ae8b35c72a556b7daa8e85ac04bd Mon Sep 17 00:00:00 2001 From: sama Pyb Date: Fri, 31 Jul 2026 20:33:53 +0800 Subject: [PATCH] [Rust] Reject server over-acks Track the highest physical offset handed to each gRPC connection and reject malformed cumulative acknowledgements before they can advance durability state. Preserve pending single records and batches on protocol violations. Fixes #638 Signed-off-by: sama Pyb --- rust/NEXT_CHANGELOG.md | 5 + rust/sdk/src/stream/grpc/receiver.rs | 123 +++++++++++++++++++-- rust/sdk/src/stream/grpc/sender.rs | 130 +++++++++++++++++++--- rust/sdk/src/stream/grpc/supervisor.rs | 5 +- rust/sdk/src/stream/grpc/types.rs | 7 ++ rust/tests/src/mock_grpc.rs | 38 +++++++ rust/tests/src/rust_tests.rs | 145 +++++++++++++++++++++++++ 7 files changed, 426 insertions(+), 27 deletions(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 37c22529..f0b2ffdd 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -21,6 +21,11 @@ ### Bug Fixes +- **Rust gRPC — reject server over-acks**: cumulative durability + acknowledgements that exceed the highest request sent on the active + connection now fail the stream before any record is reported durable or the + acknowledgement watermark advances. + - **Arrow Flight — invalid acknowledgment watermarks are rejected** (Beta): ack progress is now monotonic, so delayed or duplicate responses cannot move the durable watermark backward. A response claiming more records than were actually submitted on the active connection is rejected without making buffered, unsent records appear durable. ### Documentation diff --git a/rust/sdk/src/stream/grpc/receiver.rs b/rust/sdk/src/stream/grpc/receiver.rs index 206f99cc..e8da159a 100644 --- a/rust/sdk/src/stream/grpc/receiver.rs +++ b/rust/sdk/src/stream/grpc/receiver.rs @@ -12,7 +12,7 @@ use tokio::time::Duration; use tokio_util::sync::CancellationToken; use tracing::{error, info, instrument, span, Level}; -use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; +use super::types::{CallbackMessage, OneshotMap, RecordLandingZone, SentOffsetWatermark}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; use crate::databricks::zerobus::ephemeral_stream_response::Payload as ResponsePayload; use crate::databricks::zerobus::{ @@ -20,6 +20,27 @@ use crate::databricks::zerobus::{ }; use crate::{OffsetId, StreamConfigurationOptions, ZerobusError, ZerobusResult}; +fn validate_ack_offset( + ack_offset: OffsetId, + last_acked_offset: OffsetId, + highest_sent_offset: OffsetId, +) -> ZerobusResult { + if ack_offset < 0 { + return Err(ZerobusError::InvalidStateError(format!( + "Server ack offset {ack_offset} is negative" + ))); + } + if ack_offset <= last_acked_offset { + return Ok(false); + } + if ack_offset > highest_sent_offset { + return Err(ZerobusError::InvalidStateError(format!( + "Server ack offset {ack_offset} exceeds highest sent offset {highest_sent_offset}" + ))); + } + Ok(true) +} + impl ZerobusStream { /// Spawns a task that continuously reads from `response_grpc_stream` /// and propagates the received durability acknowledgements to the @@ -36,11 +57,12 @@ impl ZerobusStream { server_error_tx: tokio::sync::watch::Sender>, recv_drain_token: CancellationToken, callback_tx: Option>, + highest_sent_offset: SentOffsetWatermark, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { let span = span!(Level::DEBUG, "inbound_stream_processor"); let _guard = span.enter(); - let mut last_acked_offset = -1; + let mut last_acked_offset: OffsetId = -1; let mut pause_deadline: Option = None; // Set when we exit because the supervisor signalled close (`recv_drain_token`). // On that path we drain the response stream inline so the server sees END_STREAM @@ -108,22 +130,48 @@ impl ZerobusStream { return Err(error); } }; + let sent_offset = *highest_sent_offset + .lock() + .expect("Sent offset watermark lock poisoned"); + match validate_ack_offset( + durability_ack_up_to_offset, + last_acked_offset, + sent_offset, + ) { + Ok(true) => {} + Ok(false) => continue, + Err(error) => { + error!("{error}"); + let _ = server_error_tx.send(Some(error.clone())); + return Err(error); + } + } let mut last_logical_acked_offset = -2; let mut map = oneshot_map.lock().await; - for _offset_to_ack in + for offset_to_ack in (last_acked_offset + 1)..=durability_ack_up_to_offset { - if let Ok(record) = landing_zone.remove_observed() { - let logical_offset = record.offset_id; - last_logical_acked_offset = logical_offset; - - if let Some(sender) = map.remove(&logical_offset) { - let _ = sender.send(Ok(logical_offset)); + let record = match landing_zone.remove_observed() { + Ok(record) => record, + Err(_) => { + let message = format!( + "Server ack offset {durability_ack_up_to_offset} could not be applied at physical offset {offset_to_ack}" + ); + error!("{message}"); + let error = ZerobusError::InvalidStateError(message); + let _ = server_error_tx.send(Some(error.clone())); + return Err(error); } + }; + let logical_offset = record.offset_id; + last_logical_acked_offset = logical_offset; - if let Some(ref tx) = callback_tx { - let _ = tx.send(CallbackMessage::Ack(logical_offset)); - } + if let Some(sender) = map.remove(&logical_offset) { + let _ = sender.send(Ok(logical_offset)); + } + + if let Some(ref tx) = callback_tx { + let _ = tx.send(CallbackMessage::Ack(logical_offset)); } } drop(map); @@ -247,3 +295,54 @@ impl ZerobusStream { }) } } + +#[cfg(test)] +mod tests { + use super::validate_ack_offset; + use crate::ZerobusError; + + #[test] + fn negative_ack_is_rejected() { + let error = validate_ack_offset(-1, -1, 0).expect_err("negative ack must fail"); + assert!(matches!( + error, + ZerobusError::InvalidStateError(message) + if message == "Server ack offset -1 is negative" + )); + } + + #[test] + fn duplicate_or_regressive_ack_is_ignored() { + assert!(!validate_ack_offset(3, 3, 5).expect("duplicate ack is valid")); + assert!(!validate_ack_offset(2, 3, 5).expect("regressive ack is valid")); + } + + #[test] + fn regressive_ack_does_not_lower_watermark() { + let mut last_acked_offset = -1; + for ack_offset in [1, 0, 2] { + if validate_ack_offset(ack_offset, last_acked_offset, 2) + .expect("ack sequence must be valid") + { + last_acked_offset = ack_offset; + } + } + assert_eq!(last_acked_offset, 2); + } + + #[test] + fn ack_beyond_highest_sent_offset_is_rejected() { + let error = validate_ack_offset(4, 2, 3).expect_err("over-ack must fail"); + assert!(matches!( + error, + ZerobusError::InvalidStateError(message) + if message + == "Server ack offset 4 exceeds highest sent offset 3" + )); + } + + #[test] + fn advancing_ack_within_sent_range_is_applied() { + assert!(validate_ack_offset(4, 2, 4).expect("valid ack must advance")); + } +} diff --git a/rust/sdk/src/stream/grpc/sender.rs b/rust/sdk/src/stream/grpc/sender.rs index 540d1491..00c0bc49 100644 --- a/rust/sdk/src/stream/grpc/sender.rs +++ b/rust/sdk/src/stream/grpc/sender.rs @@ -9,7 +9,7 @@ use std::sync::Arc; use tokio_util::sync::CancellationToken; use tracing::error; -use super::types::RecordLandingZone; +use super::types::{RecordLandingZone, SentOffsetWatermark}; use super::ZerobusStream; use crate::databricks::zerobus::EphemeralStreamRequest; use crate::offset_generator::OffsetIdGenerator; @@ -24,6 +24,7 @@ impl ZerobusStream { is_paused: Arc, server_error_tx: tokio::sync::watch::Sender>, cancellation_token: CancellationToken, + highest_sent_offset: SentOffsetWatermark, ) -> tokio::task::JoinHandle> { tokio::spawn(async move { let physical_offset_id_generator = OffsetIdGenerator::default(); @@ -39,24 +40,125 @@ impl ZerobusStream { } } => item.clone(), }; + + let permit = tokio::select! { + biased; + _ = cancellation_token.cancelled() => return Ok(()), + permit = outbound_stream.reserve() => permit, + }; + let permit = match permit { + Ok(permit) => permit, + Err(err) => { + error!("Failed to reserve outbound stream capacity: {}", err); + let error = ZerobusError::StreamClosedError(tonic::Status::internal( + "Failed to send record", + )); + let _ = server_error_tx.send(Some(error.clone())); + return Err(error); + } + }; + let offset_id = physical_offset_id_generator.next(); let request_payload = item.payload.into_request_payload(offset_id); + let request = EphemeralStreamRequest { + payload: Some(request_payload), + }; - let send_result = outbound_stream - .send(EphemeralStreamRequest { - payload: Some(request_payload), - }) - .await; - - if let Err(err) = send_result { - error!("Failed to send record: {}", err); - let error = ZerobusError::StreamClosedError(tonic::Status::internal( - "Failed to send record", - )); - let _ = server_error_tx.send(Some(error.clone())); - return Err(error); + { + let mut watermark = highest_sent_offset + .lock() + .expect("Sent offset watermark lock poisoned"); + permit.send(request); + *watermark = offset_id; } } }) } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicBool; + use std::sync::{Arc, Mutex}; + + use tokio::sync::{mpsc, watch}; + use tokio::time::{timeout, Duration}; + use tokio_util::sync::CancellationToken; + + use super::ZerobusStream; + use crate::databricks::zerobus::RecordType; + use crate::landing_zone::LandingZone; + use crate::stream::grpc::types::IngestRequest; + use crate::EncodedBatch; + + async fn wait_for_watermark(watermark: &Arc>, expected: i64) { + timeout(Duration::from_secs(1), async { + loop { + if *watermark.lock().expect("watermark lock poisoned") == expected { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("sender did not publish the expected watermark"); + } + + #[tokio::test] + async fn blocked_channel_does_not_publish_unsent_offset() { + let landing_zone = Arc::new(LandingZone::new(2)); + for logical_offset in 0..2 { + let payload = + EncodedBatch::try_from_record(vec![logical_offset as u8], RecordType::Proto) + .expect("record type must match"); + landing_zone + .add(Box::new(IngestRequest { + payload, + offset_id: logical_offset, + })) + .await; + } + + let (outbound_tx, mut outbound_rx) = mpsc::channel(1); + let (server_error_tx, _server_error_rx) = watch::channel(None); + let cancellation_token = CancellationToken::new(); + let highest_sent_offset = Arc::new(Mutex::new(-1)); + let task = ZerobusStream::spawn_sender_task( + outbound_tx, + Arc::clone(&landing_zone), + Arc::new(AtomicBool::new(false)), + server_error_tx, + cancellation_token.clone(), + Arc::clone(&highest_sent_offset), + ); + + wait_for_watermark(&highest_sent_offset, 0).await; + timeout(Duration::from_secs(1), async { + while landing_zone.observed_count() != 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the second item was not observed"); + assert_eq!( + *highest_sent_offset.lock().expect("watermark lock poisoned"), + 0, + "the second item is observed but blocked on channel capacity" + ); + + outbound_rx + .recv() + .await + .expect("first request must be sent"); + wait_for_watermark(&highest_sent_offset, 1).await; + outbound_rx + .recv() + .await + .expect("second request must be sent"); + + cancellation_token.cancel(); + task.await + .expect("sender task must not panic") + .expect("sender task must stop cleanly"); + } +} diff --git a/rust/sdk/src/stream/grpc/supervisor.rs b/rust/sdk/src/stream/grpc/supervisor.rs index 74b03cad..f87546d8 100644 --- a/rust/sdk/src/stream/grpc/supervisor.rs +++ b/rust/sdk/src/stream/grpc/supervisor.rs @@ -18,7 +18,7 @@ use tokio_util::sync::CancellationToken; use tonic::transport::Channel; use tracing::{debug, error, info, instrument, warn}; -use super::types::{CallbackMessage, OneshotMap, RecordLandingZone}; +use super::types::{CallbackMessage, OneshotMap, RecordLandingZone, SentOffsetWatermark}; use super::{ZerobusStream, STREAM_TEARDOWN_DRAIN_TIMEOUT_MS}; use crate::databricks::zerobus::zerobus_client::ZerobusClient; use crate::errors::should_retry_initial_connection; @@ -203,6 +203,7 @@ impl ZerobusStream { // 3. Spawn receiver and sender task. let is_paused = Arc::new(AtomicBool::new(false)); + let highest_sent_offset: SentOffsetWatermark = Arc::new(std::sync::Mutex::new(-1)); // Per-stream child token let per_stream_token = cancellation_token.child_token(); @@ -219,6 +220,7 @@ impl ZerobusStream { server_error_tx.clone(), recv_drain_token.clone(), callback_tx.clone(), + Arc::clone(&highest_sent_offset), ); let mut send_task = Self::spawn_sender_task( tx, @@ -226,6 +228,7 @@ impl ZerobusStream { Arc::clone(&is_paused), server_error_tx.clone(), per_stream_token.clone(), + highest_sent_offset, ); // 4. Wait for any of the two tasks to end. diff --git a/rust/sdk/src/stream/grpc/types.rs b/rust/sdk/src/stream/grpc/types.rs index 329cbbff..b3d29229 100644 --- a/rust/sdk/src/stream/grpc/types.rs +++ b/rust/sdk/src/stream/grpc/types.rs @@ -30,6 +30,13 @@ pub(super) type OneshotMap = /// Landing zone for ingest records. pub(super) type RecordLandingZone = Arc>>; +/// Highest physical offset handed to the active gRPC connection. +/// +/// The sender updates this watermark while holding the same lock that the +/// receiver uses for validation. This prevents a fast acknowledgement from +/// racing the sender between channel handoff and watermark publication. +pub(super) type SentOffsetWatermark = Arc>; + /// Messages sent to the callback handler task. #[derive(Debug, Clone)] pub(super) enum CallbackMessage { diff --git a/rust/tests/src/mock_grpc.rs b/rust/tests/src/mock_grpc.rs index be365957..dc1007f4 100644 --- a/rust/tests/src/mock_grpc.rs +++ b/rust/tests/src/mock_grpc.rs @@ -34,6 +34,16 @@ pub enum MockResponse { ack_up_to_offset: i64, delay_ms: u64, }, + /// Record acknowledgment sent when a different request offset is observed. + /// + /// This allows tests to model a server that acknowledges beyond what the + /// client has sent. + #[allow(dead_code)] + RecordAckOnOffset { + trigger_offset: i64, + ack_up_to_offset: i64, + delay_ms: u64, + }, /// Close stream signal #[allow(dead_code)] CloseStreamSignal { @@ -470,6 +480,34 @@ async fn handle_mock_response( (true, current_index) } } + MockResponse::RecordAckOnOffset { + trigger_offset, + ack_up_to_offset, + delay_ms, + } => { + if offset == Some(*trigger_offset) { + if *delay_ms > 0 { + sleep(Duration::from_millis(*delay_ms)).await; + } + info!( + "Sending RecordAck response for {} with ack_up_to_offset: {}", + request_type, ack_up_to_offset + ); + let response = EphemeralStreamResponse { + payload: Some(ResponsePayload::IngestRecordResponse( + IngestRecordResponse { + durability_ack_up_to_offset: Some(*ack_up_to_offset), + }, + )), + }; + if tx.send(Ok(response)).await.is_err() { + return (false, current_index); + } + (true, current_index + 1) + } else { + (true, current_index) + } + } MockResponse::CloseStreamSignal { duration_seconds, delay_ms, diff --git a/rust/tests/src/rust_tests.rs b/rust/tests/src/rust_tests.rs index 48fcee85..54e47225 100644 --- a/rust/tests/src/rust_tests.rs +++ b/rust/tests/src/rust_tests.rs @@ -2873,6 +2873,151 @@ mod failure_scenarios_tests { Ok(()) } + #[tokio::test] + async fn test_server_over_ack_fails_stream() -> Result<(), Box> { + setup_tracing(); + info!("Starting test_server_over_ack_fails_stream"); + + let (mock_server, server_url) = start_mock_server().await?; + + mock_server + .inject_responses( + TABLE_NAME, + vec![ + MockResponse::CreateStream { + stream_id: "test_stream_over_ack".to_string(), + delay_ms: 0, + }, + MockResponse::RecordAckOnOffset { + trigger_offset: 0, + ack_up_to_offset: 1, + delay_ms: 0, + }, + ], + ) + .await; + + let sdk = ZerobusSdk::builder() + .endpoint(server_url.clone()) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .compiled_proto(create_test_descriptor_proto().unwrap_or_default()) + .recovery(false) + .build() + .await?; + + let offset = stream + .ingest_record_offset(b"pending data".to_vec()) + .await?; + let error = stream + .wait_for_offset(offset) + .await + .expect_err("an acknowledgement beyond the sent offset must fail the stream"); + assert!( + !error.is_retryable(), + "a protocol violation must be terminal" + ); + assert!( + error + .to_string() + .contains("Server ack offset 1 exceeds highest sent offset 0"), + "unexpected error: {error}" + ); + assert_eq!(mock_server.get_write_count().await, 1); + + let unacked = stream.get_unacked_records().await?.collect::>(); + assert_eq!(unacked.len(), 1); + assert!(matches!( + &unacked[0], + databricks_zerobus_ingest_sdk::EncodedRecord::Proto(payload) + if payload == b"pending data" + )); + + let later_ingest = stream.ingest_record_offset(b"more data".to_vec()).await; + assert!(matches!( + later_ingest, + Err(ZerobusError::StreamClosedError(_)) + )); + + Ok(()) + } + + #[tokio::test] + async fn test_batch_over_ack_uses_physical_offset() -> Result<(), Box> + { + setup_tracing(); + info!("Starting test_batch_over_ack_uses_physical_offset"); + + let (mock_server, server_url) = start_mock_server().await?; + + mock_server + .inject_responses( + TABLE_NAME, + vec![ + MockResponse::CreateStream { + stream_id: "test_stream_batch_over_ack".to_string(), + delay_ms: 0, + }, + MockResponse::RecordAckOnOffset { + trigger_offset: 0, + ack_up_to_offset: 1, + delay_ms: 0, + }, + ], + ) + .await; + + let sdk = ZerobusSdk::builder() + .endpoint(server_url.clone()) + .unity_catalog_url("https://mock-uc.com") + .tls_config(Arc::new(NoTlsConfig)) + .build()?; + let stream = sdk + .stream_builder() + .table(TABLE_NAME) + .headers_provider(Arc::new(TestHeadersProvider::default())) + .compiled_proto(create_test_descriptor_proto().unwrap_or_default()) + .recovery(false) + .build() + .await?; + + let batch: Vec> = (0..5) + .map(|index| format!("batch-{index}").into_bytes()) + .collect(); + let offset = stream + .ingest_records_offset(batch.clone()) + .await? + .expect("a non-empty batch has an offset"); + let error = stream + .wait_for_offset(offset) + .await + .expect_err("a batch is one physical request, so ack 1 must fail"); + assert!( + error + .to_string() + .contains("Server ack offset 1 exceeds highest sent offset 0"), + "unexpected error: {error}" + ); + assert_eq!(mock_server.get_write_count().await, 5); + + let unacked = stream.get_unacked_records().await?.collect::>(); + assert_eq!(unacked.len(), batch.len()); + for (record, expected) in unacked.iter().zip(batch) { + assert!(matches!( + record, + databricks_zerobus_ingest_sdk::EncodedRecord::Proto(payload) + if payload == &expected + )); + } + + Ok(()) + } + #[tokio::test] async fn test_server_unresponsiveness_fails_stream( ) -> Result<(), Box> {