multiplex channel && rpc service add authentication

Signed-off-by: junxiang Mu <1948535941@qq.com>

common/common/Cargo.toml

@@ -6,4 +6,5 @@ edition.workspace = true
 [dependencies]
 lazy_static.workspace = true
 tokio.workspace = true
+tonic.workspace = true
 tracing-error.workspace = true

common/common/src/globals.rs

@@ -1,8 +1,12 @@
+use std::collections::HashMap;
+
 use lazy_static::lazy_static;
 use tokio::sync::RwLock;
+use tonic::transport::Channel;
 
 lazy_static! {
     pub static ref GLOBAL_Local_Node_Name: RwLock<String> = RwLock::new("".to_string());
     pub static ref GLOBAL_Rustfs_Host: RwLock<String> = RwLock::new("".to_string());
     pub static ref GLOBAL_Rustfs_Port: RwLock<String> = RwLock::new("9000".to_string());
+    pub static ref GLOBAL_Conn_Map: RwLock<HashMap<String, Channel>> = RwLock::new(HashMap::new());
 }

common/lock/src/remote_client.rs

@@ -1,6 +1,6 @@
 use async_trait::async_trait;
 use common::error::{Error, Result};
-use protos::proto_gen::node_service::{node_service_client::NodeServiceClient, GenerallyLockRequest};
+use protos::{node_service_time_out_client, proto_gen::node_service::GenerallyLockRequest};
 use tonic::Request;
 use tracing::info;
 
@@ -8,18 +8,13 @@ use crate::{lock_args::LockArgs, Locker};
 
 #[derive(Debug, Clone)]
 pub struct RemoteClient {
-    url: url::Url,
+    addr: String,
 }
 
 impl RemoteClient {
     pub fn new(url: url::Url) -> Self {
-        Self { url }
-    }
-
-    async fn get_client_v2(&self) -> Result<NodeServiceClient<tonic::transport::Channel>> {
-        // Ok(NodeServiceClient::connect("http://220.181.1.138:9000").await?)
-        let addr = format!("{}://{}:{}", self.url.scheme(), self.url.host_str().unwrap(), self.url.port().unwrap());
-        Ok(NodeServiceClient::connect(addr).await?)
+        let addr = format!("{}://{}:{}", url.scheme(), url.host_str().unwrap(), url.port().unwrap());
+        Self { addr }
     }
 }
 
@@ -28,7 +23,9 @@ impl Locker for RemoteClient {
     async fn lock(&mut self, args: &LockArgs) -> Result<bool> {
         info!("remote lock");
         let args = serde_json::to_string(args)?;
-        let mut client = self.get_client_v2().await?;
+        let mut client = node_service_time_out_client(&self.addr)
+            .await
+            .map_err(|err| Error::from_string(format!("can not get client, err: {}", err.to_string())))?;
         let request = Request::new(GenerallyLockRequest { args });
 
         let response = client.lock(request).await?.into_inner();
@@ -43,7 +40,9 @@ impl Locker for RemoteClient {
     async fn unlock(&mut self, args: &LockArgs) -> Result<bool> {
         info!("remote unlock");
         let args = serde_json::to_string(args)?;
-        let mut client = self.get_client_v2().await?;
+        let mut client = node_service_time_out_client(&self.addr)
+            .await
+            .map_err(|err| Error::from_string(format!("can not get client, err: {}", err.to_string())))?;
         let request = Request::new(GenerallyLockRequest { args });
 
         let response = client.un_lock(request).await?.into_inner();
@@ -58,7 +57,9 @@ impl Locker for RemoteClient {
     async fn rlock(&mut self, args: &LockArgs) -> Result<bool> {
         info!("remote rlock");
         let args = serde_json::to_string(args)?;
-        let mut client = self.get_client_v2().await?;
+        let mut client = node_service_time_out_client(&self.addr)
+            .await
+            .map_err(|err| Error::from_string(format!("can not get client, err: {}", err.to_string())))?;
         let request = Request::new(GenerallyLockRequest { args });
 
         let response = client.r_lock(request).await?.into_inner();
@@ -73,7 +74,9 @@ impl Locker for RemoteClient {
     async fn runlock(&mut self, args: &LockArgs) -> Result<bool> {
         info!("remote runlock");
         let args = serde_json::to_string(args)?;
-        let mut client = self.get_client_v2().await?;
+        let mut client = node_service_time_out_client(&self.addr)
+            .await
+            .map_err(|err| Error::from_string(format!("can not get client, err: {}", err.to_string())))?;
         let request = Request::new(GenerallyLockRequest { args });
 
         let response = client.r_un_lock(request).await?.into_inner();
@@ -88,7 +91,9 @@ impl Locker for RemoteClient {
     async fn force_unlock(&mut self, args: &LockArgs) -> Result<bool> {
         info!("remote force_unlock");
         let args = serde_json::to_string(args)?;
-        let mut client = self.get_client_v2().await?;
+        let mut client = node_service_time_out_client(&self.addr)
+            .await
+            .map_err(|err| Error::from_string(format!("can not get client, err: {}", err.to_string())))?;
         let request = Request::new(GenerallyLockRequest { args });
 
         let response = client.force_un_lock(request).await?.into_inner();
@@ -103,7 +108,9 @@ impl Locker for RemoteClient {
     async fn refresh(&mut self, args: &LockArgs) -> Result<bool> {
         info!("remote refresh");
         let args = serde_json::to_string(args)?;
-        let mut client = self.get_client_v2().await?;
+        let mut client = node_service_time_out_client(&self.addr)
+            .await
+            .map_err(|err| Error::from_string(format!("can not get client, err: {}", err.to_string())))?;
         let request = Request::new(GenerallyLockRequest { args });
 
         let response = client.refresh(request).await?.into_inner();

common/protos/Cargo.toml

@@ -5,6 +5,7 @@ edition.workspace = true
 
 [dependencies]
 #async-backtrace = { workspace = true, optional = true }
+common.workspace = true
 flatbuffers = { workspace = true }
 prost = { workspace = true }
 protobuf = { workspace = true }

common/protos/src/generated/flatbuffers_generated/models.rs

@@ -1,9 +1,10 @@
 // automatically generated by the FlatBuffers compiler, do not modify
 
+
 // @generated
 
-use core::cmp::Ordering;
 use core::mem;
+use core::cmp::Ordering;
 
 extern crate flatbuffers;
 use self::flatbuffers::{EndianScalar, Follow};
@@ -11,114 +12,112 @@ use self::flatbuffers::{EndianScalar, Follow};
 #[allow(unused_imports, dead_code)]
 pub mod models {
 
-    use core::cmp::Ordering;
-    use core::mem;
+  use core::mem;
+  use core::cmp::Ordering;
 
-    extern crate flatbuffers;
-    use self::flatbuffers::{EndianScalar, Follow};
+  extern crate flatbuffers;
+  use self::flatbuffers::{EndianScalar, Follow};
 
-    pub enum PingBodyOffset {}
-    #[derive(Copy, Clone, PartialEq)]
+pub enum PingBodyOffset {}
+#[derive(Copy, Clone, PartialEq)]
 
-    pub struct PingBody<'a> {
-        pub _tab: flatbuffers::Table<'a>,
+pub struct PingBody<'a> {
+  pub _tab: flatbuffers::Table<'a>,
+}
+
+impl<'a> flatbuffers::Follow<'a> for PingBody<'a> {
+  type Inner = PingBody<'a>;
+  #[inline]
+  unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
+    Self { _tab: flatbuffers::Table::new(buf, loc) }
+  }
+}
+
+impl<'a> PingBody<'a> {
+  pub const VT_PAYLOAD: flatbuffers::VOffsetT = 4;
+
+  pub const fn get_fully_qualified_name() -> &'static str {
+    "models.PingBody"
+  }
+
+  #[inline]
+  pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
+    PingBody { _tab: table }
+  }
+  #[allow(unused_mut)]
+  pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>(
+    _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>,
+    args: &'args PingBodyArgs<'args>
+  ) -> flatbuffers::WIPOffset<PingBody<'bldr>> {
+    let mut builder = PingBodyBuilder::new(_fbb);
+    if let Some(x) = args.payload { builder.add_payload(x); }
+    builder.finish()
+  }
+
+
+  #[inline]
+  pub fn payload(&self) -> Option<flatbuffers::Vector<'a, u8>> {
+    // Safety:
+    // Created from valid Table for this object
+    // which contains a valid value in this slot
+    unsafe { self._tab.get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, u8>>>(PingBody::VT_PAYLOAD, None)}
+  }
+}
+
+impl flatbuffers::Verifiable for PingBody<'_> {
+  #[inline]
+  fn run_verifier(
+    v: &mut flatbuffers::Verifier, pos: usize
+  ) -> Result<(), flatbuffers::InvalidFlatbuffer> {
+    use self::flatbuffers::Verifiable;
+    v.visit_table(pos)?
+     .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, u8>>>("payload", Self::VT_PAYLOAD, false)?
+     .finish();
+    Ok(())
+  }
+}
+pub struct PingBodyArgs<'a> {
+    pub payload: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, u8>>>,
+}
+impl<'a> Default for PingBodyArgs<'a> {
+  #[inline]
+  fn default() -> Self {
+    PingBodyArgs {
+      payload: None,
     }
+  }
+}
 
-    impl<'a> flatbuffers::Follow<'a> for PingBody<'a> {
-        type Inner = PingBody<'a>;
-        #[inline]
-        unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner {
-            Self {
-                _tab: flatbuffers::Table::new(buf, loc),
-            }
-        }
+pub struct PingBodyBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> {
+  fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>,
+  start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
+}
+impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> PingBodyBuilder<'a, 'b, A> {
+  #[inline]
+  pub fn add_payload(&mut self, payload: flatbuffers::WIPOffset<flatbuffers::Vector<'b , u8>>) {
+    self.fbb_.push_slot_always::<flatbuffers::WIPOffset<_>>(PingBody::VT_PAYLOAD, payload);
+  }
+  #[inline]
+  pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>) -> PingBodyBuilder<'a, 'b, A> {
+    let start = _fbb.start_table();
+    PingBodyBuilder {
+      fbb_: _fbb,
+      start_: start,
     }
+  }
+  #[inline]
+  pub fn finish(self) -> flatbuffers::WIPOffset<PingBody<'a>> {
+    let o = self.fbb_.end_table(self.start_);
+    flatbuffers::WIPOffset::new(o.value())
+  }
+}
 
-    impl<'a> PingBody<'a> {
-        pub const VT_PAYLOAD: flatbuffers::VOffsetT = 4;
+impl core::fmt::Debug for PingBody<'_> {
+  fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+    let mut ds = f.debug_struct("PingBody");
+      ds.field("payload", &self.payload());
+      ds.finish()
+  }
+}
+}  // pub mod models
 
-        pub const fn get_fully_qualified_name() -> &'static str {
-            "models.PingBody"
-        }
-
-        #[inline]
-        pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self {
-            PingBody { _tab: table }
-        }
-        #[allow(unused_mut)]
-        pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>(
-            _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>,
-            args: &'args PingBodyArgs<'args>,
-        ) -> flatbuffers::WIPOffset<PingBody<'bldr>> {
-            let mut builder = PingBodyBuilder::new(_fbb);
-            if let Some(x) = args.payload {
-                builder.add_payload(x);
-            }
-            builder.finish()
-        }
-
-        #[inline]
-        pub fn payload(&self) -> Option<flatbuffers::Vector<'a, u8>> {
-            // Safety:
-            // Created from valid Table for this object
-            // which contains a valid value in this slot
-            unsafe {
-                self._tab
-                    .get::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'a, u8>>>(PingBody::VT_PAYLOAD, None)
-            }
-        }
-    }
-
-    impl flatbuffers::Verifiable for PingBody<'_> {
-        #[inline]
-        fn run_verifier(v: &mut flatbuffers::Verifier, pos: usize) -> Result<(), flatbuffers::InvalidFlatbuffer> {
-            use self::flatbuffers::Verifiable;
-            v.visit_table(pos)?
-                .visit_field::<flatbuffers::ForwardsUOffset<flatbuffers::Vector<'_, u8>>>("payload", Self::VT_PAYLOAD, false)?
-                .finish();
-            Ok(())
-        }
-    }
-    pub struct PingBodyArgs<'a> {
-        pub payload: Option<flatbuffers::WIPOffset<flatbuffers::Vector<'a, u8>>>,
-    }
-    impl<'a> Default for PingBodyArgs<'a> {
-        #[inline]
-        fn default() -> Self {
-            PingBodyArgs { payload: None }
-        }
-    }
-
-    pub struct PingBodyBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> {
-        fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>,
-        start_: flatbuffers::WIPOffset<flatbuffers::TableUnfinishedWIPOffset>,
-    }
-    impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> PingBodyBuilder<'a, 'b, A> {
-        #[inline]
-        pub fn add_payload(&mut self, payload: flatbuffers::WIPOffset<flatbuffers::Vector<'b, u8>>) {
-            self.fbb_
-                .push_slot_always::<flatbuffers::WIPOffset<_>>(PingBody::VT_PAYLOAD, payload);
-        }
-        #[inline]
-        pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>) -> PingBodyBuilder<'a, 'b, A> {
-            let start = _fbb.start_table();
-            PingBodyBuilder {
-                fbb_: _fbb,
-                start_: start,
-            }
-        }
-        #[inline]
-        pub fn finish(self) -> flatbuffers::WIPOffset<PingBody<'a>> {
-            let o = self.fbb_.end_table(self.start_);
-            flatbuffers::WIPOffset::new(o.value())
-        }
-    }
-
-    impl core::fmt::Debug for PingBody<'_> {
-        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-            let mut ds = f.debug_struct("PingBody");
-            ds.field("payload", &self.payload());
-            ds.finish()
-        }
-    }
-} // pub mod models

common/protos/src/lib.rs

@@ -1,28 +1,44 @@
 mod generated;
-use std::time::Duration;
+use std::{error::Error, time::Duration};
 
+use common::globals::GLOBAL_Conn_Map;
 pub use generated::*;
 use proto_gen::node_service::node_service_client::NodeServiceClient;
-use tonic::{codec::CompressionEncoding, transport::Channel};
-use tower::timeout::Timeout;
+use tonic::{
+    metadata::MetadataValue,
+    service::interceptor::InterceptedService,
+    transport::{Channel, Endpoint},
+    Request, Status,
+};
 
 // Default 100 MB
 pub const DEFAULT_GRPC_SERVER_MESSAGE_LEN: usize = 100 * 1024 * 1024;
 
-pub fn node_service_time_out_client(
-    channel: Channel,
-    time_out: Duration,
-    max_message_size: usize,
-    grpc_enable_gzip: bool,
-) -> NodeServiceClient<Timeout<Channel>> {
-    let timeout_channel = Timeout::new(channel, time_out);
-    let client = NodeServiceClient::<Timeout<Channel>>::new(timeout_channel);
-    let client = NodeServiceClient::max_decoding_message_size(client, max_message_size);
-    if grpc_enable_gzip {
-        NodeServiceClient::max_encoding_message_size(client, max_message_size)
-            .accept_compressed(CompressionEncoding::Gzip)
-            .send_compressed(CompressionEncoding::Gzip)
-    } else {
-        NodeServiceClient::max_encoding_message_size(client, max_message_size)
-    }
+pub async fn node_service_time_out_client(
+    addr: &String,
+) -> Result<
+    NodeServiceClient<
+        InterceptedService<Channel, Box<dyn Fn(Request<()>) -> Result<Request<()>, Status> + Send + Sync + 'static>>,
+    >,
+    Box<dyn Error>,
+> {
+    let token: MetadataValue<_> = "rustfs rpc".parse()?;
+    let channel = match GLOBAL_Conn_Map.read().await.get(addr) {
+        Some(channel) => channel.clone(),
+        None => {
+            let connector = Endpoint::from_shared(addr.to_string())?.connect_timeout(Duration::from_secs(60));
+            let channel = connector.connect().await?;
+            channel
+        }
+    };
+    GLOBAL_Conn_Map.write().await.insert(addr.to_string(), channel.clone());
+
+    // let timeout_channel = Timeout::new(channel, Duration::from_secs(60));
+    Ok(NodeServiceClient::with_interceptor(
+        channel,
+        Box::new(move |mut req: Request<()>| {
+            req.metadata_mut().insert("authorization", token.clone());
+            Ok(req)
+        }),
+    ))
 }