[Rust] Reject server over-acks

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

rust/sdk/src/stream/grpc/receiver.rs

@@ -12,14 +12,35 @@ 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::{
     CloseStreamSignal, EphemeralStreamResponse, IngestRecordResponse,
 };
 use crate::{OffsetId, StreamConfigurationOptions, ZerobusError, ZerobusResult};
 
+fn validate_ack_offset(
+    ack_offset: OffsetId,
+    last_acked_offset: OffsetId,
+    highest_sent_offset: OffsetId,
+) -> ZerobusResult<bool> {
+    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<Option<ZerobusError>>,
         recv_drain_token: CancellationToken,
         callback_tx: Option<tokio::sync::mpsc::UnboundedSender<CallbackMessage>>,
+        highest_sent_offset: SentOffsetWatermark,
     ) -> tokio::task::JoinHandle<ZerobusResult<()>> {
         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<tokio::time::Instant> = 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"));
+    }
+}

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<AtomicBool>,
         server_error_tx: tokio::sync::watch::Sender<Option<ZerobusError>>,
         cancellation_token: CancellationToken,
+        highest_sent_offset: SentOffsetWatermark,
     ) -> tokio::task::JoinHandle<ZerobusResult<()>> {
         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<Mutex<i64>>, 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");
+    }
+}

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,13 +220,15 @@ 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,
                 landing_zone_sender,
                 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.

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<LandingZone<Box<IngestRequest>>>;
 
+/// 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<std::sync::Mutex<OffsetId>>;
+
 /// Messages sent to the callback handler task.
 #[derive(Debug, Clone)]
 pub(super) enum CallbackMessage {