Loosen locking around RPC calls for better multi-threading

rust/benches/rpc.rs

@@ -1,6 +1,7 @@
 use std::collections::HashSet;
 
 use criterion::*;
+use futures::future::join_all;
 use hdfs_native::{minidfs::MiniDfs, Client, WriteOptions};
 
 fn bench(c: &mut Criterion) {
@@ -34,6 +35,23 @@ fn bench(c: &mut Criterion) {
         let fs = hdfs::hdfs::get_hdfs().unwrap();
         b.iter(|| fs.get_file_status("/bench").unwrap())
     });
+
+    group.sampling_mode(SamplingMode::Flat);
+    group.bench_function("getFileInfo-parallel", |b| {
+        b.to_async(&rt).iter_batched(
+            || {
+                (0..100)
+                    .map(|_| client.get_file_info("/bench"))
+                    .collect::<Vec<_>>()
+            },
+            |futures| async {
+                for result in join_all(futures).await {
+                    result.unwrap();
+                }
+            },
+            BatchSize::SmallInput,
+        )
+    });
 }
 
 criterion_group!(benches, bench);

rust/src/hdfs/connection.rs

@@ -260,7 +260,11 @@ impl RpcConnection {
         Ok(())
     }
 
-    pub(crate) async fn call(&self, method_name: &str, message: &[u8]) -> Result<Bytes> {
+    pub(crate) async fn call(
+        &self,
+        method_name: &str,
+        message: &[u8],
+    ) -> Result<oneshot::Receiver<Result<Bytes>>> {
         let call_id = self.get_next_call_id();
         let conn_header = self.get_connection_header(call_id, 0);
 
@@ -284,7 +288,7 @@ impl RpcConnection {
         self.write_messages(&[&conn_header_buf, &header_buf, message])
             .await?;
 
-        receiver.await.unwrap()
+        Ok(receiver)
     }
 }
 

rust/src/hdfs/proxy.rs

@@ -24,7 +24,7 @@ const OBSERVER_RETRY_EXCEPTION: &str = "org.apache.hadoop.ipc.ObserverRetryOnAct
 #[derive(Debug)]
 struct ProxyConnection {
     url: String,
-    inner: Option<RpcConnection>,
+    inner: Arc<tokio::sync::Mutex<Option<RpcConnection>>>,
     alignment_context: Arc<Mutex<AlignmentContext>>,
     nameservice: Option<String>,
 }
@@ -37,37 +37,42 @@ impl ProxyConnection {
     ) -> Self {
         ProxyConnection {
             url,
-            inner: None,
+            inner: Arc::new(tokio::sync::Mutex::new(None)),
             alignment_context,
             nameservice,
         }
     }
 
-    async fn get_connection(&mut self) -> Result<&RpcConnection> {
-        if self.inner.is_none() || !self.inner.as_ref().unwrap().is_alive() {
-            self.inner = Some(
-                RpcConnection::connect(
-                    &self.url,
-                    self.alignment_context.clone(),
-                    self.nameservice.as_deref(),
-                )
-                .await?,
-            );
-        }
-        Ok(self.inner.as_ref().unwrap())
-    }
+    async fn call(&self, method_name: &str, message: &[u8]) -> Result<Bytes> {
+        let receiver = {
+            let mut connection = self.inner.lock().await;
+            match &mut *connection {
+                Some(c) if c.is_alive() => (),
+                c => {
+                    *c = Some(
+                        RpcConnection::connect(
+                            &self.url,
+                            self.alignment_context.clone(),
+                            self.nameservice.as_deref(),
+                        )
+                        .await?,
+                    );
+                }
+            }
 
-    async fn call(&mut self, method_name: &str, message: &[u8]) -> Result<Bytes> {
-        self.get_connection()
-            .await?
-            .call(method_name, message)
-            .await
+            connection
+                .as_ref()
+                .unwrap()
+                .call(method_name, message)
+                .await?
+        };
+        receiver.await.unwrap()
     }
 }
 
 #[derive(Debug)]
 pub(crate) struct NameServiceProxy {
-    proxy_connections: Vec<Arc<tokio::sync::Mutex<ProxyConnection>>>,
+    proxy_connections: Vec<ProxyConnection>,
     current_index: AtomicUsize,
     msycned: AtomicBool,
 }
@@ -80,22 +85,14 @@ impl NameServiceProxy {
 
         let proxy_connections = if let Some(port) = nameservice.port() {
             let url = format!("{}:{}", nameservice.host_str().unwrap(), port);
-            vec![Arc::new(tokio::sync::Mutex::new(ProxyConnection::new(
-                url,
-                alignment_context.clone(),
-                None,
-            )))]
+            vec![ProxyConnection::new(url, alignment_context.clone(), None)]
         } else if let Some(host) = nameservice.host_str() {
             // TODO: Add check for no configured namenodes
             config
                 .get_urls_for_nameservice(host)?
                 .into_iter()
                 .map(|url| {
-                    Arc::new(tokio::sync::Mutex::new(ProxyConnection::new(
-                        url,
-                        alignment_context.clone(),
-                        Some(host.to_string()),
-                    )))
+                    ProxyConnection::new(url, alignment_context.clone(), Some(host.to_string()))
                 })
                 .collect()
         } else {
@@ -142,8 +139,6 @@ impl NameServiceProxy {
         let mut attempts = 0;
         loop {
             let result = self.proxy_connections[proxy_index]
-                .lock()
-                .await
                 .call(method_name, &message)
                 .await;