feat: grpc client support multi peers (#380)

* feat: grpc client use channel manager

* cr

src/client/Cargo.toml

@@ -17,6 +17,9 @@ datafusion = { git = "https://github.com/apache/arrow-datafusion.git", branch =
     "simd",
 ] }
 datatypes = { path = "../datatypes" }
+enum_dispatch = "0.3"
+parking_lot = "0.12"
+rand = "0.8"
 snafu = { version = "0.7", features = ["backtraces"] }
 tonic = "0.8"
 

src/client/examples/insert.rs

@@ -12,7 +12,7 @@ fn main() {
 
 #[tokio::main]
 async fn run() {
-    let client = Client::connect("http://127.0.0.1:3001").await.unwrap();
+    let client = Client::with_urls(vec!["127.0.0.1:3001"]);
     let db = Database::new("greptime", client);
 
     let expr = InsertExpr {

src/client/examples/logical.rs

@@ -18,7 +18,7 @@ fn main() {
 
 #[tokio::main]
 async fn run() {
-    let client = Client::connect("http://127.0.0.1:3001").await.unwrap();
+    let client = Client::with_urls(vec!["127.0.0.1:3001"]);
 
     let create_table_expr = CreateExpr {
         catalog_name: Some("greptime".to_string()),

src/client/examples/physical.rs

@@ -16,7 +16,7 @@ fn main() {
 
 #[tokio::main]
 async fn run() {
-    let client = Client::connect("http://127.0.0.1:3001").await.unwrap();
+    let client = Client::with_urls(vec!["127.0.0.1:3001"]);
     let db = Database::new("greptime", client);
 
     let physical = mock_physical_plan();

src/client/examples/select.rs

@@ -10,7 +10,7 @@ fn main() {
 
 #[tokio::main]
 async fn run() {
-    let client = Client::connect("http://127.0.0.1:3001").await.unwrap();
+    let client = Client::with_urls(vec!["127.0.0.1:3001"]);
     let db = Database::new("greptime", client);
 
     let sql = Select::Sql("select * from demo".to_string());

src/client/src/admin.rs

@@ -22,10 +22,6 @@ impl Admin {
         }
     }
 
-    pub async fn start(&mut self, url: impl Into<String>) -> Result<()> {
-        self.client.start(url).await
-    }
-
     pub async fn create(&self, expr: CreateExpr) -> Result<AdminResult> {
         let header = ExprHeader {
             version: PROTOCOL_VERSION,

src/client/src/client.rs

@@ -1,47 +1,96 @@
-use api::v1::{greptime_client::GreptimeClient, *};
-use snafu::{OptionExt, ResultExt};
+use std::sync::Arc;
+
+use api::v1::greptime_client::GreptimeClient;
+use api::v1::*;
+use common_grpc::channel_manager::ChannelManager;
+use parking_lot::RwLock;
+use snafu::OptionExt;
+use snafu::ResultExt;
 use tonic::transport::Channel;
 
 use crate::error;
+use crate::load_balance::LoadBalance;
+use crate::load_balance::Loadbalancer;
 use crate::Result;
 
 #[derive(Clone, Debug, Default)]
 pub struct Client {
-    client: Option<GreptimeClient<Channel>>,
+    inner: Arc<Inner>,
+}
+
+#[derive(Debug, Default)]
+struct Inner {
+    channel_manager: ChannelManager,
+    peers: Arc<RwLock<Vec<String>>>,
+    load_balance: Loadbalancer,
+}
+
+impl Inner {
+    fn with_manager(channel_manager: ChannelManager) -> Self {
+        Self {
+            channel_manager,
+            ..Default::default()
+        }
+    }
+
+    fn set_peers(&self, peers: Vec<String>) {
+        let mut guard = self.peers.write();
+        *guard = peers;
+    }
+
+    fn get_peer(&self) -> Option<String> {
+        let guard = self.peers.read();
+        self.load_balance.get_peer(&guard).cloned()
+    }
 }
 
 impl Client {
-    pub async fn start(&mut self, url: impl Into<String>) -> Result<()> {
-        match self.client.as_ref() {
-            None => {
-                let url = url.into();
-                let client = GreptimeClient::connect(url.clone())
-                    .await
-                    .context(error::ConnectFailedSnafu { url })?;
-                self.client = Some(client);
-                Ok(())
-            }
-            Some(_) => error::IllegalGrpcClientStateSnafu {
-                err_msg: "already started",
-            }
-            .fail(),
-        }
+    pub fn new() -> Self {
+        Default::default()
     }
 
-    pub fn with_client(client: GreptimeClient<Channel>) -> Self {
+    pub fn with_manager(channel_manager: ChannelManager) -> Self {
+        let inner = Arc::new(Inner::with_manager(channel_manager));
+        Self { inner }
+    }
+
+    pub fn with_urls<U, A>(urls: A) -> Self
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
+        Self::with_manager_and_urls(ChannelManager::new(), urls)
+    }
+
+    pub fn with_manager_and_urls<U, A>(channel_manager: ChannelManager, urls: A) -> Self
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
+        let inner = Inner::with_manager(channel_manager);
+        let urls: Vec<String> = urls
+            .as_ref()
+            .iter()
+            .map(|peer| peer.as_ref().to_string())
+            .collect();
+        inner.set_peers(urls);
         Self {
-            client: Some(client),
+            inner: Arc::new(inner),
         }
     }
 
-    pub async fn connect(url: impl Into<String>) -> Result<Self> {
-        let url = url.into();
-        let client = GreptimeClient::connect(url.clone())
-            .await
-            .context(error::ConnectFailedSnafu { url })?;
-        Ok(Self {
-            client: Some(client),
-        })
+    pub fn start<U, A>(&self, urls: A)
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
+        let urls: Vec<String> = urls
+            .as_ref()
+            .iter()
+            .map(|peer| peer.as_ref().to_string())
+            .collect();
+
+        self.inner.set_peers(urls);
     }
 
     pub async fn admin(&self, req: AdminRequest) -> Result<AdminResponse> {
@@ -73,18 +122,59 @@ impl Client {
     }
 
     pub async fn batch(&self, req: BatchRequest) -> Result<BatchResponse> {
-        if let Some(client) = self.client.as_ref() {
-            let res = client
-                .clone()
-                .batch(req)
-                .await
-                .context(error::TonicStatusSnafu)?;
-            Ok(res.into_inner())
-        } else {
-            error::IllegalGrpcClientStateSnafu {
-                err_msg: "not started",
-            }
-            .fail()
+        let peer = self
+            .inner
+            .get_peer()
+            .context(error::IllegalGrpcClientStateSnafu {
+                err_msg: "No available peer found",
+            })?;
+        let mut client = self.make_client(peer)?;
+        let result = client.batch(req).await.context(error::TonicStatusSnafu)?;
+        Ok(result.into_inner())
+    }
+
+    fn make_client(&self, addr: impl AsRef<str>) -> Result<GreptimeClient<Channel>> {
+        let addr = addr.as_ref();
+        let channel = self
+            .inner
+            .channel_manager
+            .get(addr)
+            .context(error::CreateChannelSnafu { addr })?;
+        Ok(GreptimeClient::new(channel))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use super::Inner;
+    use crate::load_balance::Loadbalancer;
+
+    fn mock_peers() -> Vec<String> {
+        vec![
+            "127.0.0.1:3001".to_string(),
+            "127.0.0.1:3002".to_string(),
+            "127.0.0.1:3003".to_string(),
+        ]
+    }
+
+    #[test]
+    fn test_inner() {
+        let inner = Inner::default();
+
+        assert!(matches!(
+            inner.load_balance,
+            Loadbalancer::Random(crate::load_balance::Random)
+        ));
+        assert!(inner.get_peer().is_none());
+
+        let peers = mock_peers();
+        inner.set_peers(peers.clone());
+        let all: HashSet<String> = peers.into_iter().collect();
+
+        for _ in 0..20 {
+            assert!(all.contains(&inner.get_peer().unwrap()));
         }
     }
 }

src/client/src/database.rs

@@ -43,10 +43,6 @@ impl Database {
         }
     }
 
-    pub async fn start(&mut self, url: impl Into<String>) -> Result<()> {
-        self.client.start(url).await
-    }
-
     pub fn name(&self) -> &str {
         &self.name
     }

src/client/src/error.rs

@@ -85,6 +85,17 @@ pub enum Error {
         #[snafu(backtrace)]
         source: datatypes::error::Error,
     },
+
+    #[snafu(display(
+        "Failed to create gRPC channel, peer address: {}, source: {}",
+        addr,
+        source
+    ))]
+    CreateChannel {
+        addr: String,
+        #[snafu(backtrace)]
+        source: common_grpc::error::Error,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -107,6 +118,7 @@ impl ErrorExt for Error {
                 source.status_code()
             }
             Error::CreateRecordBatches { source } => source.status_code(),
+            Error::CreateChannel { source, .. } => source.status_code(),
             Error::IllegalGrpcClientState { .. } => StatusCode::Unexpected,
         }
     }

src/client/src/lib.rs

@@ -2,6 +2,7 @@ pub mod admin;
 mod client;
 mod database;
 mod error;
+pub mod load_balance;
 
 pub use self::{
     client::Client,

src/client/src/load_balance.rs

@@ -0,0 +1,52 @@
+use enum_dispatch::enum_dispatch;
+use rand::seq::SliceRandom;
+
+#[enum_dispatch]
+pub trait LoadBalance {
+    fn get_peer<'a>(&self, peers: &'a [String]) -> Option<&'a String>;
+}
+
+#[enum_dispatch(LoadBalance)]
+#[derive(Debug)]
+pub enum Loadbalancer {
+    Random,
+}
+
+impl Default for Loadbalancer {
+    fn default() -> Self {
+        Loadbalancer::from(Random)
+    }
+}
+
+#[derive(Debug)]
+pub struct Random;
+
+impl LoadBalance for Random {
+    fn get_peer<'a>(&self, peers: &'a [String]) -> Option<&'a String> {
+        peers.choose(&mut rand::thread_rng())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use super::{LoadBalance, Random};
+
+    #[test]
+    fn test_random_lb() {
+        let peers = vec![
+            "127.0.0.1:3001".to_string(),
+            "127.0.0.1:3002".to_string(),
+            "127.0.0.1:3003".to_string(),
+            "127.0.0.1:3004".to_string(),
+        ];
+        let all: HashSet<String> = peers.clone().into_iter().collect();
+
+        let random = Random;
+        for _ in 0..100 {
+            let peer = random.get_peer(&peers).unwrap();
+            all.contains(peer);
+        }
+    }
+}