optimized conn pool usage

Author: jonas32

aerospike-core/src/commands/single_command.rs

@@ -13,7 +13,6 @@
 // limitations under the License.
 
 use std::sync::Arc;
-use std::time::Instant;
 
 use crate::cluster::partition::Partition;
 use crate::cluster::{Cluster, Node};
@@ -23,6 +22,7 @@ use crate::net::Connection;
 use crate::policy::Policy;
 use crate::Key;
 use aerospike_rt::sleep;
+use aerospike_rt::time::Instant;
 
 pub struct SingleCommand<'a> {
     cluster: Arc<Cluster>,
@@ -76,20 +76,13 @@ impl<'a> SingleCommand<'a> {
         loop {
             iterations += 1;
 
-            // too many retries
-            if let Some(max_retries) = policy.max_retries() {
-                if iterations > max_retries + 1 {
-                    bail!(ErrorKind::Connection(format!(
-                        "Timeout after {} tries",
-                        iterations
-                    )));
-                }
-            }
-
             // Sleep before trying again, after the first iteration
             if iterations > 1 {
                 if let Some(sleep_between_retries) = policy.sleep_between_retries() {
                     sleep(sleep_between_retries).await;
+                } else {
+                    // yield to free space for the runtime to execute other futures between runs because the loop would block the thread
+                    aerospike_rt::task::yield_now().await;
                 }
             }
 

aerospike-core/src/net/connection.rs

@@ -13,11 +13,9 @@
 // License for the specific language governing permissions and limitations under
 // the License.
 
-use std::ops::Add;
-
 use crate::commands::admin_command::AdminCommand;
 use crate::commands::buffer::Buffer;
-use crate::errors::Result;
+use crate::errors::{ErrorKind, Result};
 use crate::policy::ClientPolicy;
 #[cfg(all(any(feature = "rt-async-std"), not(feature = "rt-tokio")))]
 use aerospike_rt::async_std::net::Shutdown;
@@ -27,6 +25,7 @@ use aerospike_rt::net::TcpStream;
 use aerospike_rt::time::{Duration, Instant};
 #[cfg(all(any(feature = "rt-async-std"), not(feature = "rt-tokio")))]
 use futures::{AsyncReadExt, AsyncWriteExt};
+use std::ops::Add;
 
 #[derive(Debug)]
 pub struct Connection {
@@ -46,12 +45,17 @@ pub struct Connection {
 
 impl Connection {
     pub async fn new(addr: &str, policy: &ClientPolicy) -> Result<Self> {
-        let stream = TcpStream::connect(addr).await?;
+        let stream = aerospike_rt::timeout(Duration::from_secs(10), TcpStream::connect(addr)).await;
+        if stream.is_err() {
+            bail!(ErrorKind::Connection(
+                "Could not open network connection".to_string()
+            ));
+        }
         let mut conn = Connection {
             buffer: Buffer::new(policy.buffer_reclaim_threshold),
             bytes_read: 0,
             timeout: policy.timeout,
-            conn: stream,
+            conn: stream.unwrap()?,
             idle_timeout: policy.idle_timeout,
             idle_deadline: policy.idle_timeout.map(|timeout| Instant::now() + timeout),
         };
@@ -60,11 +64,11 @@ impl Connection {
         Ok(conn)
     }
 
-    pub fn close(&mut self) {
+    pub async fn close(&mut self) {
         #[cfg(all(any(feature = "rt-async-std"), not(feature = "rt-tokio")))]
         let _s = self.conn.shutdown(Shutdown::Both);
         #[cfg(all(any(feature = "rt-tokio"), not(feature = "rt-async-std")))]
-        let _s = self.conn.shutdown();
+        let _s = self.conn.shutdown().await;
     }
 
     pub async fn flush(&mut self) -> Result<()> {
@@ -77,7 +81,7 @@ impl Connection {
         self.buffer.resize_buffer(size)?;
         self.conn.read_exact(&mut self.buffer.data_buffer).await?;
         self.bytes_read += size;
-        self.buffer.reset_offset()?;
+        self.buffer.reset_offset();
         self.refresh();
         Ok(())
     }
@@ -103,7 +107,7 @@ impl Connection {
     fn refresh(&mut self) {
         self.idle_deadline = None;
         if let Some(idle_to) = self.idle_timeout {
-            self.idle_deadline = Some(Instant::now().add(idle_to))
+            self.idle_deadline = Some(Instant::now().add(idle_to));
         };
     }
 
@@ -112,7 +116,7 @@ impl Connection {
             return match AdminCommand::authenticate(self, user, password).await {
                 Ok(()) => Ok(()),
                 Err(err) => {
-                    self.close();
+                    self.close().await;
                     Err(err)
                 }
             };

aerospike-core/src/net/connection_pool.rs

@@ -13,7 +13,6 @@
 // License for the specific language governing permissions and limitations under
 // the License.
 
-use std::collections::VecDeque;
 use std::ops::{Deref, DerefMut, Drop};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
@@ -23,6 +22,8 @@ use crate::net::{Connection, Host};
 use crate::policy::ClientPolicy;
 use futures::executor::block_on;
 use futures::lock::Mutex;
+use std::collections::VecDeque;
+use std::time::Duration;
 
 #[derive(Debug)]
 struct IdleConnection(Connection);
@@ -66,17 +67,39 @@ impl Queue {
             if let Some(IdleConnection(mut conn)) = internals.connections.pop_front() {
                 if conn.is_idle() {
                     internals.num_conns -= 1;
-                    conn.close();
+                    conn.close().await;
                     continue;
                 }
                 connection = conn;
                 break;
             }
+
             if internals.num_conns >= self.0.capacity {
                 bail!(ErrorKind::NoMoreConnections);
             }
-            let conn = Connection::new(&self.0.host.address(), &self.0.policy).await?;
+
             internals.num_conns += 1;
+
+            // Free the lock to prevent deadlocking
+            drop(internals);
+
+            let conn = aerospike_rt::timeout(
+                Duration::from_secs(5),
+                Connection::new(&self.0.host.address(), &self.0.policy),
+            )
+            .await;
+
+            if conn.is_err() {
+                let mut internals = self.0.internals.lock().await;
+                internals.num_conns -= 1;
+                drop(internals);
+                bail!(ErrorKind::Connection(
+                    "Could not open network connection".to_string()
+                ));
+            }
+
+            let conn = conn.unwrap()?;
+
             connection = conn;
             break;
         }
@@ -92,7 +115,7 @@ impl Queue {
         if internals.num_conns < self.0.capacity {
             internals.connections.push_back(IdleConnection(conn));
         } else {
-            conn.close();
+            conn.close().await;
             internals.num_conns -= 1;
         }
     }
@@ -102,13 +125,13 @@ impl Queue {
             let mut internals = self.0.internals.lock().await;
             internals.num_conns -= 1;
         }
-        conn.close();
+        conn.close().await;
     }
 
     pub async fn clear(&mut self) {
         let mut internals = self.0.internals.lock().await;
         for mut conn in internals.connections.drain(..) {
-            conn.0.close();
+            conn.0.close().await;
         }
         internals.num_conns = 0;
     }

aerospike-rt/src/lib.rs

@@ -5,9 +5,11 @@ compile_error!("Please select a runtime from ['rt-tokio', 'rt-async-std']");
 compile_error!("Please select only one runtime");
 
 #[cfg(all(any(feature = "rt-async-std"), not(feature = "rt-tokio")))]
-pub use async_std::{self, fs, io, net, sync::RwLock, task, task::sleep, task::spawn};
+pub use async_std::{
+    self, fs, future::timeout, io, net, sync::RwLock, task, task::sleep, task::spawn,
+};
 #[cfg(all(any(feature = "rt-tokio"), not(feature = "rt-async-std")))]
-pub use tokio::{self, fs, io, net, spawn, sync::RwLock, task, time, time::sleep};
+pub use tokio::{self, fs, io, net, spawn, sync::RwLock, task, time, time::sleep, time::timeout};
 
 #[cfg(all(any(feature = "rt-async-std"), not(feature = "rt-tokio")))]
 pub use std::time;