feat: support batch lookup

Author: linguoxuan

crates/fluss/src/client/connection.rs

@@ -17,6 +17,7 @@
 
 use crate::client::WriterClient;
 use crate::client::admin::FlussAdmin;
+use crate::client::lookup::LookupClient;
 use crate::client::metadata::Metadata;
 use crate::client::table::FlussTable;
 use crate::config::Config;
@@ -32,6 +33,7 @@ pub struct FlussConnection {
     network_connects: Arc<RpcClient>,
     args: Config,
     writer_client: RwLock<Option<Arc<WriterClient>>>,
+    lookup_client: RwLock<Option<Arc<LookupClient>>>,
 }
 
 impl FlussConnection {
@@ -48,6 +50,7 @@ impl FlussConnection {
             network_connects: connections.clone(),
             args: arg.clone(),
             writer_client: Default::default(),
+            lookup_client: Default::default(),
         })
     }
 
@@ -90,6 +93,30 @@ impl FlussConnection {
         Ok(new_client)
     }
 
+    /// Gets or creates a lookup client for batched lookup operations.
+    pub fn get_or_create_lookup_client(&self) -> Result<Arc<LookupClient>> {
+        // 1. Fast path: Attempt to acquire a read lock to check if the client already exists.
+        if let Some(client) = self.lookup_client.read().as_ref() {
+            return Ok(client.clone());
+        }
+
+        // 2. Slow path: Acquire the write lock.
+        let mut lookup_guard = self.lookup_client.write();
+
+        // 3. Double-check: Another thread might have initialized the client
+        // while this thread was waiting for the write lock.
+        if let Some(client) = lookup_guard.as_ref() {
+            return Ok(client.clone());
+        }
+
+        // 4. Initialize the client since we are certain it doesn't exist yet.
+        let new_client = Arc::new(LookupClient::new(&self.args, self.metadata.clone()));
+
+        // 5. Store and return the newly created client.
+        *lookup_guard = Some(new_client.clone());
+        Ok(new_client)
+    }
+
     pub async fn get_table(&self, table_path: &TablePath) -> Result<FlussTable<'_>> {
         self.metadata.update_table_metadata(table_path).await?;
         let table_info = self.metadata.get_cluster().get_table(table_path)?.clone();

crates/fluss/src/client/lookup/lookup_client.rs

@@ -0,0 +1,194 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Lookup client that batches multiple lookups together for improved throughput.
+//!
+//! This client achieves parity with the Java client by:
+//! - Queuing lookup operations instead of sending them immediately
+//! - Batching multiple lookups to the same server/bucket
+//! - Running a background sender task to process batches
+
+use super::{LookupQuery, LookupQueue};
+use crate::client::lookup::lookup_sender::LookupSender;
+use crate::client::metadata::Metadata;
+use crate::config::Config;
+use crate::error::{Error, Result};
+use crate::metadata::{TableBucket, TablePath};
+use log::{debug, error};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::time::Duration;
+use tokio::sync::mpsc;
+use tokio::task::JoinHandle;
+
+/// A client that lookups values from the server with batching support.
+///
+/// The lookup client uses a queue and background sender to batch multiple
+/// lookup operations together, reducing network round trips and improving
+/// throughput.
+///
+/// # Example
+///
+/// ```ignore
+/// let lookup_client = LookupClient::new(config, metadata);
+/// let result = lookup_client.lookup(table_path, table_bucket, key_bytes).await?;
+/// ```
+pub struct LookupClient {
+    /// Channel to send lookup requests to the queue
+    lookup_tx: mpsc::Sender<LookupQuery>,
+    /// Handle to the sender task
+    sender_handle: Option<JoinHandle<()>>,
+    /// Shutdown signal sender
+    shutdown_tx: mpsc::Sender<()>,
+    /// Whether the client is closed
+    closed: AtomicBool,
+}
+
+impl LookupClient {
+    /// Creates a new lookup client.
+    pub fn new(config: &Config, metadata: Arc<Metadata>) -> Self {
+        // Extract configuration values
+        let queue_size = config.lookup_queue_size;
+        let max_batch_size = config.lookup_max_batch_size;
+        let batch_timeout_ms = config.lookup_batch_timeout_ms;
+        let max_inflight = config.lookup_max_inflight_requests;
+        let max_retries = config.lookup_max_retries;
+
+        // Create queue and channels
+        let (queue, lookup_tx, re_enqueue_tx) =
+            LookupQueue::new(queue_size, max_batch_size, batch_timeout_ms);
+
+        // Create sender
+        let mut sender =
+            LookupSender::new(metadata, queue, re_enqueue_tx, max_inflight, max_retries);
+
+        // Create shutdown channel
+        let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
+
+        // Spawn sender task
+        let sender_handle = tokio::spawn(async move {
+            tokio::select! {
+                _ = sender.run() => {
+                    debug!("Lookup sender completed");
+                }
+                _ = shutdown_rx.recv() => {
+                    debug!("Lookup sender received shutdown signal");
+                    sender.initiate_close();
+                }
+            }
+        });
+
+        Self {
+            lookup_tx,
+            sender_handle: Some(sender_handle),
+            shutdown_tx,
+            closed: AtomicBool::new(false),
+        }
+    }
+
+    /// Looks up a value by its primary key.
+    ///
+    /// This method queues the lookup operation and returns a future that will
+    /// complete when the server responds. Multiple lookups may be batched together
+    /// for improved throughput.
+    ///
+    /// # Arguments
+    /// * `table_path` - The table path
+    /// * `table_bucket` - The table bucket
+    /// * `key_bytes` - The encoded primary key bytes
+    ///
+    /// # Returns
+    /// * `Ok(Some(bytes))` - The value bytes if found
+    /// * `Ok(None)` - If the key was not found
+    /// * `Err(Error)` - If the lookup fails
+    pub async fn lookup(
+        &self,
+        table_path: TablePath,
+        table_bucket: TableBucket,
+        key_bytes: Vec<u8>,
+    ) -> Result<Option<Vec<u8>>> {
+        // Check if the client is closed
+        if self.closed.load(Ordering::Acquire) {
+            return Err(Error::UnexpectedError {
+                message: "Lookup client is closed".to_string(),
+                source: None,
+            });
+        }
+
+        let (result_tx, result_rx) = tokio::sync::oneshot::channel();
+
+        let query = LookupQuery::new(table_path, table_bucket, key_bytes, result_tx);
+
+        // Send to queue
+        self.lookup_tx
+            .send(query)
+            .await
+            .map_err(|_| Error::UnexpectedError {
+                message: "Failed to queue lookup: channel closed".to_string(),
+                source: None,
+            })?;
+
+        // Wait for result
+        result_rx.await.map_err(|_| Error::UnexpectedError {
+            message: "Lookup result channel closed".to_string(),
+            source: None,
+        })?
+    }
+
+    /// Closes the lookup client gracefully.
+    pub async fn close(mut self, timeout: Duration) {
+        debug!("Closing lookup client");
+
+        // Mark as closed to reject new lookups
+        self.closed.store(true, Ordering::Release);
+
+        // Send shutdown signal
+        let _ = self.shutdown_tx.send(()).await;
+
+        // Wait for sender to complete with timeout
+        if let Some(handle) = self.sender_handle.take() {
+            debug!("Waiting for sender task to complete...");
+            let abort_handle = handle.abort_handle();
+
+            match tokio::time::timeout(timeout, handle).await {
+                Ok(Ok(())) => {
+                    debug!("Lookup sender task completed gracefully.");
+                }
+                Ok(Err(join_error)) => {
+                    error!("Lookup sender task panicked: {:?}", join_error);
+                }
+                Err(_elapsed) => {
+                    error!("Lookup sender task did not complete within timeout. Forcing shutdown.");
+                    abort_handle.abort();
+                }
+            }
+        } else {
+            debug!("Lookup client was already closed or never initialized properly.");
+        }
+
+        debug!("Lookup client closed");
+    }
+}
+
+impl Drop for LookupClient {
+    fn drop(&mut self) {
+        // Abort the sender task on drop if it wasn't already consumed by close()
+        if let Some(handle) = self.sender_handle.take() {
+            handle.abort();
+        }
+    }
+}

crates/fluss/src/client/lookup/lookup_query.rs

@@ -0,0 +1,92 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Lookup query representation for batching lookup operations.
+
+use crate::metadata::{TableBucket, TablePath};
+use std::sync::atomic::{AtomicI32, Ordering};
+use tokio::sync::oneshot;
+
+/// Represents a single lookup query that will be batched and sent to the server.
+pub struct LookupQuery {
+    /// The table path for this lookup
+    table_path: TablePath,
+    /// The table bucket for this lookup
+    table_bucket: TableBucket,
+    /// The encoded primary key bytes
+    key: Vec<u8>,
+    /// Channel to send the result back to the caller
+    result_tx: Option<oneshot::Sender<Result<Option<Vec<u8>>, crate::error::Error>>>,
+    /// Number of retry attempts
+    retries: AtomicI32,
+}
+
+impl LookupQuery {
+    /// Creates a new lookup query.
+    pub fn new(
+        table_path: TablePath,
+        table_bucket: TableBucket,
+        key: Vec<u8>,
+        result_tx: oneshot::Sender<Result<Option<Vec<u8>>, crate::error::Error>>,
+    ) -> Self {
+        Self {
+            table_path,
+            table_bucket,
+            key,
+            result_tx: Some(result_tx),
+            retries: AtomicI32::new(0),
+        }
+    }
+
+    /// Returns the table path.
+    #[allow(dead_code)]
+    pub fn table_path(&self) -> &TablePath {
+        &self.table_path
+    }
+
+    /// Returns the table bucket.
+    pub fn table_bucket(&self) -> &TableBucket {
+        &self.table_bucket
+    }
+
+    /// Returns the encoded key bytes.
+    pub fn key(&self) -> &[u8] {
+        &self.key
+    }
+
+    /// Returns the current retry count.
+    pub fn retries(&self) -> i32 {
+        self.retries.load(Ordering::Acquire)
+    }
+
+    /// Increments the retry counter.
+    pub fn increment_retries(&self) {
+        self.retries.fetch_add(1, Ordering::AcqRel);
+    }
+
+    /// Completes the lookup with a result.
+    pub fn complete(&mut self, result: Result<Option<Vec<u8>>, crate::error::Error>) {
+        if let Some(tx) = self.result_tx.take() {
+            let _ = tx.send(result);
+        }
+    }
+
+    /// Returns true if the result has already been sent.
+    pub fn is_done(&self) -> bool {
+        self.result_tx.is_none()
+    }
+}