[ISSUE #290]🎨Optimize RocketmqDefaultClient code

rocketmq-broker/src/out_api/broker_outer_api.rs

@@ -36,7 +36,7 @@
 pub struct BrokerOuterAPI {
     remoting_client: RocketmqDefaultClient,
     name_server_address: Option<String>,
-    broker_outer_executor: TokioExecutorService,
+    broker_outer_executor: Option<TokioExecutorService>,
 }
 
 impl BrokerOuterAPI {
@@ -66,13 +66,14 @@
 }
 
 impl BrokerOuterAPI {
-    pub fn update_name_server_address_list(&mut self, addrs: String) {
+    pub async fn update_name_server_address_list(&mut self, addrs: String) {
         let addr_vec = addrs
             .split("';'")
             .map(|s| s.to_string())
             .collect::<Vec<String>>();
         self.remoting_client
             .update_name_server_address_list(addr_vec)
+            .await
     }
 
     pub async fn register_broker_all(

rocketmq-remoting/src/clients.rs

@@ -81,7 +81,7 @@ impl RemoteClient {
 
 #[allow(async_fn_in_trait)]
 pub trait RemotingClient: RemotingService {
-    fn update_name_server_address_list(&mut self, addrs: Vec<String>);
+    async fn update_name_server_address_list(&mut self, addrs: Vec<String>);
 
     fn get_name_server_address_list(&self) -> Vec<String>;
 

rocketmq-remoting/src/clients/rocketmq_default_impl.rs

@@ -33,11 +33,10 @@
     service_bridge: ServiceBridge,
     tokio_client_config: TokioClientConfig,
     //cache connection
-    connection_tables: HashMap<String /* ip:port */, Client>,
-    connection_tables_lock: std::sync::RwLock<()>,
-    lock: tokio::sync::RwLock<()>,
-    namesrv_addr_list: Arc<std::sync::Mutex<Vec<String>>>,
-    namesrv_addr_choosed: Arc<std::sync::Mutex<Option<String>>>,
+    connection_tables:
+        tokio::sync::Mutex<HashMap<String /* ip:port */, Arc<tokio::sync::Mutex<Client>>>>,
+    namesrv_addr_list: Arc<tokio::sync::Mutex<Vec<String>>>,
+    namesrv_addr_choosed: Arc<tokio::sync::Mutex<Option<String>>>,
 }
 
 impl RocketmqDefaultClient {
@@ -46,42 +45,23 @@
             service_bridge: ServiceBridge::new(),
             tokio_client_config,
             connection_tables: Default::default(),
-            connection_tables_lock: Default::default(),
-            lock: Default::default(),
             namesrv_addr_list: Arc::new(Default::default()),
             namesrv_addr_choosed: Arc::new(Default::default()),
         }
     }
 }
 
 impl RocketmqDefaultClient {
-    async fn get_and_create_client(&mut self, addr: String) -> &Client {
-        let lc = self.lock.write().await;
-        if self.connection_tables.contains_key(&addr) {
-            return self.connection_tables.get(&addr).unwrap();
+    async fn get_and_create_client(&mut self, addr: String) -> Arc<tokio::sync::Mutex<Client>> {
+        let mut mutex_guard = self.connection_tables.lock().await;
+        if mutex_guard.contains_key(&addr) {
+            return mutex_guard.get(&addr).unwrap().clone();
         }
 
         let addr_inner = addr.clone();
         let client = Client::connect(addr_inner).await.unwrap();
-
-        self.connection_tables.insert(addr.clone(), client);
-        drop(lc);
-        self.connection_tables.get(&addr).unwrap()
-    }
-
-    async fn get_and_create_client_mut(&mut self, addr: String) -> Option<&mut Client> {
-        let lc = self.lock.write().await;
-
-        if self.connection_tables.contains_key(&addr) {
-            return self.connection_tables.get_mut(&addr);
-        }
-
-        let addr_inner = addr.clone();
-        let client = Client::connect(addr_inner).await.unwrap();
-
-        self.connection_tables.insert(addr.clone(), client);
-        drop(lc);
-        self.connection_tables.get_mut(&addr)
+        mutex_guard.insert(addr.clone(), Arc::new(tokio::sync::Mutex::new(client)));
+        mutex_guard.get(&addr).unwrap().clone()
     }
 }
 
@@ -106,8 +86,8 @@
 
 #[allow(unused_variables)]
 impl RemotingClient for RocketmqDefaultClient {
-    fn update_name_server_address_list(&mut self, addrs: Vec<String>) {
-        let mut old = self.namesrv_addr_list.lock().unwrap();
+    async fn update_name_server_address_list(&mut self, addrs: Vec<String>) {
+        let mut old = self.namesrv_addr_list.lock().await;
         let mut update = false;
 
         if !addrs.is_empty() {
@@ -134,19 +114,18 @@
                 old.clone_from(&addrs);
 
                 // should close the channel if choosed addr is not exist.
-                if let Some(namesrv_addr) = self.namesrv_addr_choosed.lock().unwrap().as_ref() {
+                if let Some(namesrv_addr) = self.namesrv_addr_choosed.lock().await.as_ref() {
                     if !addrs.contains(namesrv_addr) {
-                        let write_guard = self.connection_tables_lock.write().unwrap();
                         let mut remove_vec = Vec::new();
-                        for (addr, client) in self.connection_tables.iter() {
+                        let mut result = self.connection_tables.lock().await;
+                        for (addr, client) in result.iter() {
                             if addr.contains(namesrv_addr) {
                                 remove_vec.push(addr.clone());
                             }
                         }
                         for addr in &remove_vec {
-                            self.connection_tables.remove(addr);
+                            result.remove(addr);
                         }
-                        drop(write_guard);
                     }
                 }
             }
@@ -167,12 +146,9 @@
         request: RemotingCommand,
         timeout_millis: u64,
     ) -> RemotingCommand {
-        let client = self
-            .get_and_create_client_mut(addr.clone())
-            .await
-            .take()
-            .unwrap();
-        ServiceBridge::invoke_sync(client, request, timeout_millis)
+        let client = self.get_and_create_client(addr.clone()).await;
+        let client_ref = &mut *client.lock().await;
+        ServiceBridge::invoke_sync(client_ref, request, timeout_millis)
             .await
             .unwrap()
     }
@@ -184,12 +160,9 @@
         timeout_millis: u64,
         invoke_callback: impl InvokeCallback,
     ) -> Result<(), Box<dyn Error>> {
-        let client = self
-            .get_and_create_client_mut(addr.clone())
-            .await
-            .take()
-            .unwrap();
-        ServiceBridge::invoke_async(client, request, timeout_millis, invoke_callback).await;
+        let client = self.get_and_create_client(addr.clone()).await;
+        let client_ref = &mut *client.lock().await;
+        ServiceBridge::invoke_async(client_ref, request, timeout_millis, invoke_callback).await;
         Ok(())
     }