feat: implement batch count limit, fix response handling (#12)

- Add max_requests_per_batch option (default: 1000) matching Go client
- Implement can_add() in ReadBatch and WriteBatch using request count
- Fix WriteBatch::flush() to properly error remaining inflight operations
  when server returns fewer responses than expected
- Add doc comments to OxiaClientOptions fields

Signed-off-by: mattisonchao <mattisonchao@gmail.com>

Author: mattisonchao

liboxia-native/src/batch.rs

@@ -52,18 +52,21 @@ pub(crate) struct ReadBatch {
     shard_id: i64,
     shard_manager: Arc<ShardManager>,
     provider_manager: Arc<ProviderManager>,
+    max_requests: u32,
 }
 impl ReadBatch {
     pub fn new(
         shard_id: i64,
         shard_manager: Arc<ShardManager>,
         provider_manager: Arc<ProviderManager>,
+        max_requests: u32,
     ) -> Self {
         ReadBatch {
             get_inflight: Vec::new(),
             shard_id,
             shard_manager,
             provider_manager,
+            max_requests,
         }
     }
 
@@ -72,8 +75,7 @@ impl ReadBatch {
     }
 
     fn can_add(&self, _: &Operation) -> bool {
-        //todo: support it
-        true
+        (self.get_inflight.len() as u32) < self.max_requests
     }
     fn add(&mut self, operation: Operation) {
         if let Operation::Get(get) = operation {
@@ -155,16 +157,22 @@ pub(crate) struct WriteBatch {
     delete_inflight: Vec<DeleteOperation>,
     delete_range_inflight: Vec<DeleteRangeOperation>,
     write_stream_manager: Arc<WriteStreamManager>,
+    max_requests: u32,
 }
 
 impl WriteBatch {
-    pub fn new(shard_id: i64, write_stream_manager: Arc<WriteStreamManager>) -> Self {
+    pub fn new(
+        shard_id: i64,
+        write_stream_manager: Arc<WriteStreamManager>,
+        max_requests: u32,
+    ) -> Self {
         WriteBatch {
             shard_id,
             put_inflight: Vec::new(),
             delete_inflight: Vec::new(),
             delete_range_inflight: Vec::new(),
             write_stream_manager,
+            max_requests,
         }
     }
 
@@ -174,9 +182,13 @@ impl WriteBatch {
             && self.delete_range_inflight.is_empty()
     }
 
+    fn total_count(&self) -> u32 {
+        (self.put_inflight.len() + self.delete_inflight.len() + self.delete_range_inflight.len())
+            as u32
+    }
+
     fn can_add(&self, _: &Operation) -> bool {
-        //todo: support it
-        true
+        self.total_count() < self.max_requests
     }
 
     fn add(&mut self, operation: Operation) {
@@ -212,24 +224,32 @@ impl WriteBatch {
         }
         match self.write_stream_manager.write(write_request).await {
             Ok(response) => {
-                for (mut operation, put_response) in
-                    self.put_inflight.drain(..).zip(response.puts.into_iter())
-                {
-                    operation.complete(put_response);
+                let mut put_responses = response.puts.into_iter();
+                for mut operation in self.put_inflight.drain(..) {
+                    match put_responses.next() {
+                        Some(put_response) => operation.complete(put_response),
+                        None => operation.complete_exception(UnexpectedStatus(
+                            "missing put response from server".to_string(),
+                        )),
+                    }
                 }
-                for (mut operation, delete_response) in self
-                    .delete_inflight
-                    .drain(..)
-                    .zip(response.deletes.into_iter())
-                {
-                    operation.complete(delete_response);
+                let mut delete_responses = response.deletes.into_iter();
+                for mut operation in self.delete_inflight.drain(..) {
+                    match delete_responses.next() {
+                        Some(delete_response) => operation.complete(delete_response),
+                        None => operation.complete_exception(UnexpectedStatus(
+                            "missing delete response from server".to_string(),
+                        )),
+                    }
                 }
-                for (mut operation, delete_range_response) in self
-                    .delete_range_inflight
-                    .drain(..)
-                    .zip(response.delete_ranges.into_iter())
-                {
-                    operation.complete(delete_range_response);
+                let mut delete_range_responses = response.delete_ranges.into_iter();
+                for mut operation in self.delete_range_inflight.drain(..) {
+                    match delete_range_responses.next() {
+                        Some(delete_range_response) => operation.complete(delete_range_response),
+                        None => operation.complete_exception(UnexpectedStatus(
+                            "missing delete_range response from server".to_string(),
+                        )),
+                    }
                 }
             }
             Err(err) => {

liboxia-native/src/batch_manager.rs

@@ -27,10 +27,20 @@ impl Batcher {
         shard_manager: Arc<ShardManager>,
         provider_manager: Arc<ProviderManager>,
         write_stream_manager: Arc<WriteStreamManager>,
+        max_requests_per_batch: u32,
     ) -> Batch {
         match self {
-            Batcher::Read => Batch::Read(ReadBatch::new(shard_id, shard_manager, provider_manager)),
-            Batcher::Write => Batch::Write(WriteBatch::new(shard_id, write_stream_manager)),
+            Batcher::Read => Batch::Read(ReadBatch::new(
+                shard_id,
+                shard_manager,
+                provider_manager,
+                max_requests_per_batch,
+            )),
+            Batcher::Write => Batch::Write(WriteBatch::new(
+                shard_id,
+                write_stream_manager,
+                max_requests_per_batch,
+            )),
         }
     }
 }
@@ -48,14 +58,16 @@ impl Drop for BatchManager {
 }
 
 impl BatchManager {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         shard_id: i64,
         batcher: Batcher,
         shard_manager: Arc<ShardManager>,
         provider_manager: Arc<ProviderManager>,
         write_stream_manager: Arc<WriteStreamManager>,
         batch_linger: Duration,
-        batch_max_size: u32,
+        _batch_max_size: u32,
+        max_requests_per_batch: u32,
     ) -> Self {
         let (tx, rx) = mpsc::unbounded_channel();
         let context = CancellationToken::new();
@@ -68,7 +80,7 @@ impl BatchManager {
             provider_manager,
             write_stream_manager,
             batch_linger,
-            batch_max_size,
+            max_requests_per_batch,
         ));
         BatchManager {
             context,
@@ -106,14 +118,15 @@ async fn start_batcher(
     provider_manager: Arc<ProviderManager>,
     write_stream_manager: Arc<WriteStreamManager>,
     batch_linger: Duration,
-    _batch_max_size: u32,
+    max_requests_per_batch: u32,
 ) {
     let mut buffer = Vec::new();
     let mut batch = batcher.create_batch(
         shard_id,
         shard_manager.clone(),
         provider_manager.clone(),
         write_stream_manager.clone(),
+        max_requests_per_batch,
     );
     let mut interval = interval(batch_linger);
     loop {
@@ -127,7 +140,7 @@ async fn start_batcher(
                    continue
                 }
                 batch.flush().await;
-                batch = batcher.create_batch( shard_id, shard_manager.clone(), provider_manager.clone(), write_stream_manager.clone());
+                batch = batcher.create_batch( shard_id, shard_manager.clone(), provider_manager.clone(), write_stream_manager.clone(), max_requests_per_batch);
             }
             size = rx.recv_many(&mut buffer, usize::MAX) => {
                 if size == 0 {
@@ -137,7 +150,7 @@ async fn start_batcher(
                 for operation  in buffer.drain(..) {
                     if !batch.can_add(&operation) {
                         batch.flush().await;
-                        batch = batcher.create_batch( shard_id, shard_manager.clone(), provider_manager.clone(), write_stream_manager.clone());
+                        batch = batcher.create_batch( shard_id, shard_manager.clone(), provider_manager.clone(), write_stream_manager.clone(), max_requests_per_batch);
                         interval.reset();
                     }
                     batch.add(operation);

liboxia-native/src/client.rs

@@ -516,6 +516,7 @@ impl OxiaClient {
                 self.inner.write_stream_manager.clone(),
                 self.inner.options.batch_linger,
                 self.inner.options.batch_max_size,
+                self.inner.options.max_requests_per_batch,
             ))
         };
         let ref_mut = match batcher {

liboxia-native/src/client_builder.rs

@@ -28,6 +28,7 @@ pub struct OxiaClientBuilder {
     identity: Option<String>,
     batch_linger: Option<Duration>,
     batch_max_size: Option<u32>,
+    max_requests_per_batch: Option<u32>,
     session_timeout: Option<Duration>,
     request_timeout: Option<Duration>,
 }
@@ -62,6 +63,11 @@ impl OxiaClientBuilder {
         self
     }
 
+    pub fn max_requests_per_batch(mut self, max_requests_per_batch: u32) -> Self {
+        self.max_requests_per_batch = Some(max_requests_per_batch);
+        self
+    }
+
     pub fn session_timeout(mut self, session_timeout: Duration) -> Self {
         self.session_timeout = Some(session_timeout);
         self
@@ -89,6 +95,9 @@ impl OxiaClientBuilder {
         if let Some(batch_max_size) = self.batch_max_size {
             options.batch_max_size = batch_max_size;
         }
+        if let Some(max_requests_per_batch) = self.max_requests_per_batch {
+            options.max_requests_per_batch = max_requests_per_batch;
+        }
         if let Some(session_timeout) = self.session_timeout {
             options.session_timeout = session_timeout;
         }

liboxia-native/src/client_options.rs

@@ -1,14 +1,24 @@
 use std::time::Duration;
 use uuid::Uuid;
 
+/// Configuration options for the Oxia client.
 #[derive(Debug, Clone)]
 pub struct OxiaClientOptions {
+    /// The address of the Oxia service (e.g., "http://localhost:6648").
     pub service_address: String,
+    /// The Oxia namespace to use (default: "default").
     pub namespace: String,
+    /// Client identity string, used for ephemeral record tracking.
     pub identity: String,
+    /// Wait time before sending a batch (default: 5ms).
     pub batch_linger: Duration,
+    /// Maximum batch size in bytes (default: 128KB).
     pub batch_max_size: u32,
+    /// Maximum number of requests per batch (default: 1000).
+    pub max_requests_per_batch: u32,
+    /// Session timeout for ephemeral records (default: 15s).
     pub session_timeout: Duration,
+    /// Timeout for individual requests (default: 30s).
     pub request_timeout: Duration,
 }
 
@@ -20,6 +30,7 @@ impl Default for OxiaClientOptions {
             identity: Uuid::new_v4().to_string(),
             batch_linger: Duration::from_millis(5),
             batch_max_size: 128 * 1024,
+            max_requests_per_batch: 1000,
             session_timeout: Duration::from_secs(15),
             request_timeout: Duration::from_secs(30),
         }