refactor: Simplify dereferencing logic in MetaStore and update async methods for better clarity (#17938)

Author: drmingdrmer

Cargo.lock

@@ -20,6 +20,7 @@ databend-common-meta-kvapi = { workspace = true }
 databend-common-meta-semaphore = { workspace = true }
 databend-common-meta-types = { workspace = true }
 databend-meta = { workspace = true }
+futures = { workspace = true }
 log = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true }

src/meta/store/Cargo.toml

@@ -17,26 +17,19 @@ pub(crate) mod local;
 use std::ops::Deref;
 use std::pin::Pin;
 use std::sync::Arc;
-use std::task::Context;
-use std::task::Poll;
 use std::time::Duration;
 
 use databend_common_grpc::RpcClientConf;
 use databend_common_meta_client::errors::CreationError;
 use databend_common_meta_client::ClientHandle;
 use databend_common_meta_client::MetaGrpcClient;
-use databend_common_meta_kvapi::kvapi;
-use databend_common_meta_kvapi::kvapi::KVStream;
-use databend_common_meta_kvapi::kvapi::UpsertKVReply;
 use databend_common_meta_semaphore::acquirer::Permit;
 use databend_common_meta_semaphore::errors::AcquireError;
 use databend_common_meta_semaphore::Semaphore;
 use databend_common_meta_types::protobuf::WatchRequest;
 use databend_common_meta_types::protobuf::WatchResponse;
 use databend_common_meta_types::MetaError;
-use databend_common_meta_types::TxnReply;
-use databend_common_meta_types::TxnRequest;
-use databend_common_meta_types::UpsertKV;
+use futures::stream::TryStreamExt;
 pub use local::LocalMetaService;
 use log::info;
 use tokio_stream::Stream;
@@ -56,6 +49,17 @@ pub enum MetaStore {
     R(Arc<ClientHandle>),
 }
 
+impl Deref for MetaStore {
+    type Target = Arc<ClientHandle>;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            MetaStore::L(l) => l.deref(),
+            MetaStore::R(r) => r,
+        }
+    }
+}
+
 impl MetaStore {
     /// Create a local meta service for testing.
     ///
@@ -79,24 +83,16 @@ impl MetaStore {
         }
     }
 
-    pub async fn get_local_addr(&self) -> std::result::Result<Option<String>, MetaError> {
-        let client = match self {
-            MetaStore::L(l) => l.deref().deref(),
-            MetaStore::R(grpc_client) => grpc_client,
-        };
-
-        let client_info = client.get_client_info().await?;
-        Ok(Some(client_info.client_addr))
+    pub async fn get_local_addr(&self) -> Result<String, MetaError> {
+        let client_info = self.get_client_info().await?;
+        Ok(client_info.client_addr)
     }
 
     pub async fn watch(&self, request: WatchRequest) -> Result<WatchStream, MetaError> {
-        let client = match self {
-            MetaStore::L(l) => l.deref(),
-            MetaStore::R(grpc_client) => grpc_client,
-        };
+        let client = self.deref();
 
         let streaming = client.request(request).await?;
-        Ok(Box::pin(WatchResponseStream::create(streaming)))
+        Ok(Box::pin(streaming.map_err(MetaError::from)))
     }
 
     pub async fn new_acquired(
@@ -106,11 +102,7 @@ impl MetaStore {
         id: impl ToString,
         lease: Duration,
     ) -> Result<Permit, AcquireError> {
-        let client = match self {
-            MetaStore::L(l) => l.deref(),
-            MetaStore::R(grpc_client) => grpc_client,
-        };
-
+        let client = self.deref();
         Semaphore::new_acquired(client.clone(), prefix, capacity, id, lease).await
     }
 
@@ -121,48 +113,11 @@ impl MetaStore {
         id: impl ToString,
         lease: Duration,
     ) -> Result<Permit, AcquireError> {
-        let client = match self {
-            MetaStore::L(l) => l.deref(),
-            MetaStore::R(grpc_client) => grpc_client,
-        };
-
+        let client = self.deref();
         Semaphore::new_acquired_by_time(client.clone(), prefix, capacity, id, lease).await
     }
 }
 
-#[async_trait::async_trait]
-impl kvapi::KVApi for MetaStore {
-    type Error = MetaError;
-
-    async fn upsert_kv(&self, act: UpsertKV) -> Result<UpsertKVReply, Self::Error> {
-        match self {
-            MetaStore::L(x) => x.upsert_kv(act).await,
-            MetaStore::R(x) => x.upsert_kv(act).await,
-        }
-    }
-
-    async fn get_kv_stream(&self, keys: &[String]) -> Result<KVStream<Self::Error>, Self::Error> {
-        match self {
-            MetaStore::L(x) => x.get_kv_stream(keys).await,
-            MetaStore::R(x) => x.get_kv_stream(keys).await,
-        }
-    }
-
-    async fn list_kv(&self, prefix: &str) -> Result<KVStream<Self::Error>, Self::Error> {
-        match self {
-            MetaStore::L(x) => x.list_kv(prefix).await,
-            MetaStore::R(x) => x.list_kv(prefix).await,
-        }
-    }
-
-    async fn transaction(&self, txn: TxnRequest) -> Result<TxnReply, Self::Error> {
-        match self {
-            MetaStore::L(x) => x.transaction(txn).await,
-            MetaStore::R(x) => x.transaction(txn).await,
-        }
-    }
-}
-
 impl MetaStoreProvider {
     pub fn new(rpc_conf: RpcClientConf) -> Self {
         MetaStoreProvider { rpc_conf }
@@ -188,37 +143,3 @@ impl MetaStoreProvider {
         }
     }
 }
-
-pub struct WatchResponseStream<E, S>
-where
-    E: Into<MetaError> + Send + 'static,
-    S: Stream<Item = Result<WatchResponse, E>> + Send + Unpin + 'static,
-{
-    inner: S,
-}
-
-impl<E, S> WatchResponseStream<E, S>
-where
-    E: Into<MetaError> + Send + 'static,
-    S: Stream<Item = Result<WatchResponse, E>> + Send + Unpin + 'static,
-{
-    pub fn create(inner: S) -> WatchResponseStream<E, S> {
-        WatchResponseStream { inner }
-    }
-}
-
-impl<E, S> Stream for WatchResponseStream<E, S>
-where
-    E: Into<MetaError> + Send + 'static,
-    S: Stream<Item = Result<WatchResponse, E>> + Send + Unpin + 'static,
-{
-    type Item = Result<WatchResponse, MetaError>;
-
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        Pin::new(&mut self.inner).poll_next(cx).map(|x| match x {
-            None => None,
-            Some(Ok(resp)) => Some(Ok(resp)),
-            Some(Err(e)) => Some(Err(e.into())),
-        })
-    }
-}

src/meta/store/src/lib.rs

@@ -126,7 +126,7 @@ pub trait WarehouseApi: Sync + Send {
         cluster: &str,
     ) -> Result<Vec<NodeInfo>>;
 
-    async fn get_local_addr(&self) -> Result<Option<String>>;
+    async fn get_local_addr(&self) -> Result<String>;
 
     async fn list_online_nodes(&self) -> Result<Vec<NodeInfo>>;
 

src/query/management/src/warehouse/warehouse_api.rs

@@ -2121,7 +2121,7 @@ impl WarehouseApi for WarehouseMgr {
 
     #[async_backtrace::framed]
     #[fastrace::trace]
-    async fn get_local_addr(&self) -> Result<Option<String>> {
+    async fn get_local_addr(&self) -> Result<String> {
         Ok(self.metastore.get_local_addr().await?)
     }
 

src/query/management/src/warehouse/warehouse_mgr.rs

@@ -512,18 +512,18 @@ impl ClusterDiscovery {
             if let Ok(socket_addr) = SocketAddr::from_str(lookup_ip) {
                 let ip_addr = socket_addr.ip();
                 if ip_addr.is_loopback() || ip_addr.is_unspecified() {
-                    if let Some(local_addr) = self.warehouse_manager.get_local_addr().await? {
-                        let local_socket_addr = SocketAddr::from_str(&local_addr)?;
-                        let new_addr = format!("{}:{}", local_socket_addr.ip(), socket_addr.port());
-                        warn!(
-                            "Detected loopback or unspecified address as {} endpoint. \
-                            We rewrite it(\"{}\" -> \"{}\") for advertising to other nodes. \
-                            If there are proxies between nodes, you can specify endpoint with --{}.",
-                            typ, lookup_ip, new_addr, typ
-                        );
-
-                        *lookup_ip = new_addr;
-                    }
+                    let local_addr = self.warehouse_manager.get_local_addr().await?;
+
+                    let local_socket_addr = SocketAddr::from_str(&local_addr)?;
+                    let new_addr = format!("{}:{}", local_socket_addr.ip(), socket_addr.port());
+                    warn!(
+                        "Detected loopback or unspecified address as {} endpoint. \
+                        We rewrite it(\"{}\" -> \"{}\") for advertising to other nodes. \
+                        If there are proxies between nodes, you can specify endpoint with --{}.",
+                        typ, lookup_ip, new_addr, typ
+                    );
+
+                    *lookup_ip = new_addr;
                 }
             }
         }