feat(cluster): auto discover ip when ip is unspecified or loop back

Author: zhang2014

Cargo.lock

@@ -55,6 +55,16 @@ impl MetaStore {
             MetaStore::R(_) => false,
         }
     }
+
+    pub async fn get_local_addr(&self) -> std::result::Result<Option<String>, MetaError> {
+        match self {
+            MetaStore::L(_) => Ok(None),
+            MetaStore::R(grpc_client) => {
+                let client_info = grpc_client.get_client_info().await?;
+                Ok(Some(client_info.client_addr))
+            }
+        }
+    }
 }
 
 #[async_trait::async_trait]

src/meta/store/src/lib.rs

@@ -22,6 +22,7 @@ common-meta-api = { path = "../../meta/api" }
 common-meta-types = { path = "../../meta/types" }
 common-proto-conv = { path = "../../meta/proto-conv" }
 common-protos = { path = "../../meta/protos" }
+common-meta-store = { path = "../../meta/store" }
 
 async-trait = "0.1.56"
 serde_json = "1.0.81"

src/query/management/Cargo.toml

@@ -28,4 +28,6 @@ pub trait ClusterApi: Sync + Send {
 
     // Keep the tenant's cluster node alive.
     async fn heartbeat(&self, node: &NodeInfo, seq: Option<u64>) -> Result<u64>;
+
+    async fn get_local_addr(&self) -> Result<Option<String>>;
 }

src/query/management/src/cluster/cluster_api.rs

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use std::ops::Add;
-use std::sync::Arc;
 use std::time::Duration;
 use std::time::UNIX_EPOCH;
 
@@ -22,6 +21,7 @@ use common_base::base::unescape_for_key;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_api::KVApi;
+use common_meta_store::MetaStore;
 use common_meta_types::KVMeta;
 use common_meta_types::MatchSeq;
 use common_meta_types::NodeInfo;
@@ -35,14 +35,14 @@ use crate::cluster::ClusterApi;
 pub static CLUSTER_API_KEY_PREFIX: &str = "__fd_clusters";
 
 pub struct ClusterMgr {
-    kv_api: Arc<dyn KVApi>,
+    metastore: MetaStore,
     lift_time: Duration,
     cluster_prefix: String,
 }
 
 impl ClusterMgr {
     pub fn create(
-        kv_api: Arc<dyn KVApi>,
+        metastore: MetaStore,
         tenant: &str,
         cluster_id: &str,
         lift_time: Duration,
@@ -54,7 +54,7 @@ impl ClusterMgr {
         }
 
         Ok(ClusterMgr {
-            kv_api,
+            metastore,
             lift_time,
             cluster_prefix: format!(
                 "{}/{}/{}/databend_query",
@@ -87,7 +87,7 @@ impl ClusterApi for ClusterMgr {
         let value = Operation::Update(serde_json::to_vec(&node)?);
         let node_key = format!("{}/{}", self.cluster_prefix, escape_for_key(&node.id)?);
         let upsert_node = self
-            .kv_api
+            .metastore
             .upsert_kv(UpsertKVReq::new(&node_key, seq, value, meta));
 
         let res = upsert_node.await?.added_or_else(|v| {
@@ -101,7 +101,7 @@ impl ClusterApi for ClusterMgr {
     }
 
     async fn get_nodes(&self) -> Result<Vec<NodeInfo>> {
-        let values = self.kv_api.prefix_list_kv(&self.cluster_prefix).await?;
+        let values = self.metastore.prefix_list_kv(&self.cluster_prefix).await?;
 
         let mut nodes_info = Vec::with_capacity(values.len());
         for (node_key, value) in values {
@@ -116,7 +116,7 @@ impl ClusterApi for ClusterMgr {
 
     async fn drop_node(&self, node_id: String, seq: Option<u64>) -> Result<()> {
         let node_key = format!("{}/{}", self.cluster_prefix, escape_for_key(&node_id)?);
-        let upsert_node = self.kv_api.upsert_kv(UpsertKVReq::new(
+        let upsert_node = self.metastore.upsert_kv(UpsertKVReq::new(
             &node_key,
             seq.into(),
             Operation::Delete,
@@ -145,7 +145,7 @@ impl ClusterApi for ClusterMgr {
         };
 
         let upsert_meta =
-            self.kv_api
+            self.metastore
                 .upsert_kv(UpsertKVReq::new(&node_key, seq, Operation::AsIs, meta));
 
         match upsert_meta.await? {
@@ -157,4 +157,8 @@ impl ClusterApi for ClusterMgr {
             UpsertKVReply { .. } => self.add_node(node.clone()).await,
         }
     }
+
+    async fn get_local_addr(&self) -> Result<Option<String>> {
+        Ok(self.metastore.get_local_addr().await?)
+    }
 }

src/query/management/src/cluster/cluster_mgr.rs

@@ -21,6 +21,7 @@ use common_exception::Result;
 use common_management::*;
 use common_meta_api::KVApi;
 use common_meta_embedded::MetaEmbedded;
+use common_meta_store::MetaStore;
 use common_meta_types::NodeInfo;
 use common_meta_types::SeqV;
 
@@ -153,8 +154,8 @@ fn create_test_node_info() -> NodeInfo {
     }
 }
 
-async fn new_cluster_api() -> Result<(Arc<MetaEmbedded>, ClusterMgr)> {
-    let test_api = Arc::new(MetaEmbedded::new_temp().await?);
+async fn new_cluster_api() -> Result<(MetaStore, ClusterMgr)> {
+    let test_api = MetaStore::L(Arc::new(MetaEmbedded::new_temp().await?));
     let cluster_manager = ClusterMgr::create(
         test_api.clone(),
         "test-tenant-id",

src/query/management/tests/it/cluster.rs

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::net::SocketAddr;
 use std::ops::RangeInclusive;
+use std::str::FromStr;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
@@ -35,7 +37,7 @@ use common_exception::Result;
 use common_grpc::ConnectionFactory;
 use common_management::ClusterApi;
 use common_management::ClusterMgr;
-use common_meta_api::KVApi;
+use common_meta_store::MetaStore;
 use common_meta_store::MetaStoreProvider;
 use common_meta_types::NodeInfo;
 use common_metrics::label_counter_with_val_and_labels;
@@ -139,24 +141,24 @@ static CLUSTER_DISCOVERY: OnceCell<Singleton<Arc<ClusterDiscovery>>> = OnceCell:
 impl ClusterDiscovery {
     const METRIC_LABEL_FUNCTION: &'static str = "function";
 
-    pub async fn create_meta_client(cfg: &Config) -> Result<Arc<dyn KVApi>> {
+    pub async fn create_meta_client(cfg: &Config) -> Result<MetaStore> {
         let meta_api_provider = MetaStoreProvider::new(cfg.meta.to_meta_grpc_client_conf());
         match meta_api_provider.try_get_meta_store().await {
-            Ok(client) => Ok(client.arc()),
+            Ok(meta_store) => Ok(meta_store),
             Err(cause) => Err(cause.add_message_back("(while create cluster api).")),
         }
     }
 
     pub async fn init(cfg: Config, v: Singleton<Arc<ClusterDiscovery>>) -> Result<()> {
-        let meta_client = ClusterDiscovery::create_meta_client(&cfg).await?;
-        v.init(Self::try_create(&cfg, meta_client).await?)?;
+        let metastore = ClusterDiscovery::create_meta_client(&cfg).await?;
+        v.init(Self::try_create(&cfg, metastore).await?)?;
 
         CLUSTER_DISCOVERY.set(v).ok();
         Ok(())
     }
 
-    pub async fn try_create(cfg: &Config, api: Arc<dyn KVApi>) -> Result<Arc<ClusterDiscovery>> {
-        let (lift_time, provider) = Self::create_provider(cfg, api)?;
+    pub async fn try_create(cfg: &Config, metastore: MetaStore) -> Result<Arc<ClusterDiscovery>> {
+        let (lift_time, provider) = Self::create_provider(cfg, metastore)?;
 
         Ok(Arc::new(ClusterDiscovery {
             local_id: GlobalUniqName::unique(),
@@ -182,13 +184,13 @@ impl ClusterDiscovery {
 
     fn create_provider(
         cfg: &Config,
-        api: Arc<dyn KVApi>,
+        metastore: MetaStore,
     ) -> Result<(Duration, Arc<dyn ClusterApi>)> {
         // TODO: generate if tenant or cluster id is empty
         let tenant_id = &cfg.query.tenant_id;
         let cluster_id = &cfg.query.cluster_id;
         let lift_time = Duration::from_secs(60);
-        let cluster_manager = ClusterMgr::create(api, tenant_id, cluster_id, lift_time)?;
+        let cluster_manager = ClusterMgr::create(metastore, tenant_id, cluster_id, lift_time)?;
 
         Ok((lift_time, Arc::new(cluster_manager)))
     }
@@ -350,8 +352,27 @@ impl ClusterDiscovery {
 
     pub async fn register_to_metastore(self: &Arc<Self>, cfg: &Config) -> Result<()> {
         let cpus = cfg.query.num_cpus;
-        // TODO: 127.0.0.1 || ::0
-        let address = cfg.query.flight_api_address.clone();
+        let mut address = cfg.query.flight_api_address.clone();
+
+        if let Ok(socket_addr) = SocketAddr::from_str(&address) {
+            let ip_addr = socket_addr.ip();
+            if ip_addr.is_loopback() || ip_addr.is_unspecified() {
+                if let Some(local_addr) = self.api_provider.get_local_addr().await? {
+                    let local_socket_addr = SocketAddr::from_str(&local_addr)?;
+                    let new_addr = format!("{}:{}", local_socket_addr.ip(), socket_addr.port());
+                    tracing::warn!(
+                        "Used loopback or unspecified address as cluster flight address. \
+                        we rewrite it(\"{}\" -> \"{}\") for other nodes can connect it.\
+                        If your has proxy between nodes, you can specify the node's IP address in the configuration file.",
+                        address,
+                        new_addr
+                    );
+
+                    address = new_addr;
+                }
+            }
+        }
+
         let node_info = NodeInfo::create(self.local_id.clone(), cpus, address);
 
         self.drop_invalid_nodes(&node_info).await?;

src/query/service/src/clusters/cluster.rs

@@ -44,9 +44,9 @@ async fn test_remove_invalid_nodes() -> Result<()> {
         .query_flight_address("invalid_address_2")
         .build();
 
-    let meta_client = ClusterDiscovery::create_meta_client(&config_1).await?;
-    let cluster_discovery_1 = ClusterDiscovery::try_create(&config_1, meta_client.clone()).await?;
-    let cluster_discovery_2 = ClusterDiscovery::try_create(&config_2, meta_client.clone()).await?;
+    let metastore = ClusterDiscovery::create_meta_client(&config_1).await?;
+    let cluster_discovery_1 = ClusterDiscovery::try_create(&config_1, metastore.clone()).await?;
+    let cluster_discovery_2 = ClusterDiscovery::try_create(&config_2, metastore.clone()).await?;
 
     cluster_discovery_1.register_to_metastore(&config_1).await?;
     cluster_discovery_2.register_to_metastore(&config_2).await?;

src/query/service/tests/it/clusters.rs

@@ -0,0 +1,4 @@
+statement query I
+select * from system.clusters where host = '0.0.0.0' or host = '::';
+
+----