[ArrowFlight] Align test mock server with actual server ack behavior. (#537)

Author: nikolaobradovic-db

rust/tests/src/arrow_tests.rs

@@ -825,6 +825,84 @@ mod arrow_flight_tests {
     mod recovery_tests {
         use super::*;
 
+        /// Regression test for the mock contract: `ack_up_to_records` must be
+        /// connection-relative. A retriable error forces a second DoPut connection
+        /// for the same table; the auto-ack on that connection must count only the
+        /// rows replayed on it, NOT the cumulative rows across both connections.
+        /// Under the old global row counter the auto-ack would be 6 (3 + 3) and this
+        /// test would fail, masking slicing/replay bugs.
+        #[tokio::test]
+        async fn test_auto_ack_is_connection_relative() -> Result<(), Box<dyn std::error::Error>> {
+            setup_tracing();
+            info!("Starting test_auto_ack_is_connection_relative");
+
+            let (mock_server, server_url) = start_mock_flight_server().await?;
+            let schema = create_test_arrow_schema();
+
+            // First connection errors (retriable) after the batch is decoded; the
+            // batch stays unacked and is replayed on the recovered connection, where
+            // responses are exhausted so it hits the auto-ack path.
+            mock_server
+                .inject_responses(
+                    TABLE_NAME,
+                    vec![MockFlightResponse::Error {
+                        status: tonic::Status::unavailable("Temporary network issue"),
+                        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()))
+                .arrow(schema.clone())
+                .recovery(true)
+                .recovery_timeout_ms(5000)
+                .recovery_backoff_ms(100)
+                .recovery_retries(3)
+                .build_arrow()
+                .await?;
+
+            // One 3-row batch. It is decoded on connection 1 (errors), then replayed
+            // and decoded again on connection 2 (auto-acked).
+            let batch = create_test_record_batch(
+                schema.clone(),
+                vec![1, 2, 3],
+                vec![Some("a"), Some("b"), Some("c")],
+            );
+            let offset = stream.ingest_batch(batch).await?;
+            stream.wait_for_offset(offset).await?;
+
+            // Both connections decoded the 3-row batch, so the GLOBAL observation is 6.
+            assert_eq!(
+                mock_server.get_total_records_received().await,
+                6,
+                "both connections should have decoded the batch (global observation)"
+            );
+
+            // But every auto-ack must be connection-relative: exactly the 3 rows
+            // replayed on the recovered connection, never the cumulative 6.
+            let acks = mock_server.get_auto_ack_records().await;
+            assert!(
+                !acks.is_empty(),
+                "expected an auto-ack on the recovered connection"
+            );
+            assert!(
+                acks.iter().all(|&r| r == 3),
+                "auto-ack must exclude rows from the first connection (expected all == 3), got {:?}",
+                acks
+            );
+
+            Ok(())
+        }
+
         #[tokio::test]
         async fn test_supervisor_recovery_after_retriable_error(
         ) -> Result<(), Box<dyn std::error::Error>> {

rust/tests/src/mock_arrow_flight.rs

@@ -88,8 +88,9 @@ pub struct MockFlightServer {
     row_count: Arc<Mutex<u64>>,
     /// Track response index across connection attempts
     response_indices: Arc<Mutex<HashMap<String, usize>>>,
-    /// Track expected offset per table (must be strictly sequential starting from 0)
-    expected_offsets: Arc<Mutex<HashMap<String, i64>>>,
+    /// Observation of every `ack_up_to_records` value emitted on the auto-ack path,
+    /// in emission order. Used by tests to assert acks are connection-relative.
+    auto_ack_records: Arc<Mutex<Vec<u64>>>,
 }
 
 impl MockFlightServer {
@@ -100,7 +101,7 @@ impl MockFlightServer {
             batch_count: Arc::new(Mutex::new(0)),
             row_count: Arc::new(Mutex::new(0)),
             response_indices: Arc::new(Mutex::new(HashMap::new())),
-            expected_offsets: Arc::new(Mutex::new(HashMap::new())),
+            auto_ack_records: Arc::new(Mutex::new(Vec::new())),
         }
     }
 
@@ -128,6 +129,12 @@ impl MockFlightServer {
         *self.row_count.lock().await
     }
 
+    /// Get every `ack_up_to_records` value emitted on the auto-ack path, in order.
+    #[allow(dead_code)]
+    pub async fn get_auto_ack_records(&self) -> Vec<u64> {
+        self.auto_ack_records.lock().await.clone()
+    }
+
     /// Reset the server state
     #[allow(dead_code)]
     pub async fn reset(&self) {
@@ -138,8 +145,7 @@ impl MockFlightServer {
         *self.max_offset_received.lock().await = -1;
         *self.batch_count.lock().await = 0;
         *self.row_count.lock().await = 0;
-        let mut expected_offsets = self.expected_offsets.lock().await;
-        expected_offsets.clear();
+        self.auto_ack_records.lock().await.clear();
     }
 }
 
@@ -219,11 +225,18 @@ impl FlightService for MockFlightServer {
         let batch_count = Arc::clone(&self.batch_count);
         let row_count = Arc::clone(&self.row_count);
         let response_indices = Arc::clone(&self.response_indices);
-        let expected_offsets = Arc::clone(&self.expected_offsets);
+        let auto_ack_records = Arc::clone(&self.auto_ack_records);
 
         tokio::spawn(async move {
             let mut stream_responses: Vec<MockFlightResponse> = Vec::new();
             let mut is_first_message = true;
+            // Per-connection state, owned as locals like the real server (fresh per
+            // DoPut connection). These must not leak across reconnects or concurrent
+            // same-table streams: `expected_offset` validates wire ordering, and
+            // `connection_record_count` is the cumulative-record tracker that drives
+            // auto-acks (the server derives ack_up_to_records per connection).
+            let mut expected_offset: i64 = 0;
+            let mut connection_record_count: u64 = 0;
 
             // Load configured responses
             {
@@ -240,12 +253,6 @@ impl FlightService for MockFlightServer {
                 *indices.get(&table_name).unwrap_or(&0)
             };
 
-            // Reset expected offset to 0 for each new connection
-            {
-                let mut offsets = expected_offsets.lock().await;
-                offsets.insert(table_name.clone(), 0);
-            }
-
             while let Ok(Some(flight_data)) = stream.message().await {
                 // Handle schema message (first message has no app_metadata or empty app_metadata)
                 if is_first_message {
@@ -278,30 +285,23 @@ impl FlightService for MockFlightServer {
                 if let Some(metadata) = &metadata {
                     debug!("Received batch with offset_id: {}", metadata.offset_id);
 
-                    // Validate offset is strictly sequential
-                    let expected = {
-                        let offsets = expected_offsets.lock().await;
-                        *offsets.get(&table_name).unwrap_or(&0)
-                    };
-                    if metadata.offset_id != expected {
+                    // Validate offset is strictly sequential (connection-local).
+                    if metadata.offset_id != expected_offset {
                         error!(
                             "Non-incremental offset: expected {}, got {}",
-                            expected, metadata.offset_id
+                            expected_offset, metadata.offset_id
                         );
                         let _ = tx
                             .send(Err(Status::invalid_argument(format!(
                                 "Non-incremental offset: expected {}, actual {}",
-                                expected, metadata.offset_id
+                                expected_offset, metadata.offset_id
                             ))))
                             .await;
                         return;
                     }
 
-                    // Update expected offset for next batch
-                    {
-                        let mut offsets = expected_offsets.lock().await;
-                        offsets.insert(table_name.clone(), metadata.offset_id + 1);
-                    }
+                    // Update expected offset for next batch.
+                    expected_offset = metadata.offset_id + 1;
 
                     // Update max offset
                     {
@@ -324,6 +324,10 @@ impl FlightService for MockFlightServer {
                         .map(|rb| rb.length() as u64)
                         .unwrap_or(0);
                     if rows > 0 {
+                        // Connection-local counter drives acks (mirrors the server's
+                        // per-connection cumulative tracker); the global row_count is
+                        // kept only as a cross-connection observation metric.
+                        connection_record_count += rows;
                         let mut count = row_count.lock().await;
                         *count += rows;
                     }
@@ -446,10 +450,10 @@ impl FlightService for MockFlightServer {
                 } else {
                     // Auto-ack if no more configured responses
                     if let Some(metadata) = metadata {
-                        let records = {
-                            let count = row_count.lock().await;
-                            *count
-                        };
+                        // Use the connection-local record count so acks are
+                        // connection-relative, matching the real server.
+                        let records = connection_record_count;
+                        auto_ack_records.lock().await.push(records);
                         let ack_metadata = FlightAckMetadata {
                             ack_up_to_offset: metadata.offset_id,
                             ack_up_to_records: records,
@@ -511,7 +515,7 @@ pub async fn start_mock_flight_server(
         batch_count: Arc::clone(&mock_server.batch_count),
         row_count: Arc::clone(&mock_server.row_count),
         response_indices: Arc::clone(&mock_server.response_indices),
-        expected_offsets: Arc::clone(&mock_server.expected_offsets),
+        auto_ack_records: Arc::clone(&mock_server.auto_ack_records),
     };
 
     let addr: std::net::SocketAddr = "127.0.0.1:0".parse()?;