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> {