Merge pull request #395 from havaker/keepalive

Keepalive queries

Author: psarna

scylla/src/transport/cluster.rs

@@ -14,6 +14,7 @@ use itertools::Itertools;
 use std::collections::{BTreeMap, HashMap};
 use std::net::SocketAddr;
 use std::sync::Arc;
+use std::time::Duration;
 use tracing::{debug, warn};
 
 /// Cluster manages up to date information and connections to database nodes.
@@ -101,6 +102,7 @@ impl Cluster {
             metadata_reader: MetadataReader::new(
                 initial_peers,
                 pool_config.connection_config.clone(),
+                pool_config.keepalive_interval,
                 server_events_sender,
                 fetch_schema_metadata,
             ),
@@ -310,7 +312,7 @@ impl ClusterData {
 
 impl ClusterWorker {
     pub async fn work(mut self) {
-        use tokio::time::{Duration, Instant};
+        use tokio::time::Instant;
 
         let refresh_duration = Duration::from_secs(60); // Refresh topology every 60 seconds
         let mut last_refresh_time = Instant::now();

scylla/src/transport/connection.rs

@@ -196,14 +196,6 @@ pub struct ConnectionConfig {
     // should be Some only in control connections,
     pub event_sender: Option<mpsc::Sender<Event>>,
     pub default_consistency: Consistency,
-    /*
-    These configuration options will be added in the future:
-
-    pub tcp_keepalive: bool,
-
-    pub load_balancing: Option<String>,
-    pub retry_policy: Option<String>,
-    */
 }
 
 impl Default for ConnectionConfig {

scylla/src/transport/connection_pool.rs

@@ -47,6 +47,7 @@ pub struct PoolConfig {
     pub connection_config: ConnectionConfig,
     pub pool_size: PoolSize,
     pub can_use_shard_aware_port: bool,
+    pub keepalive_interval: Option<Duration>,
 }
 
 impl Default for PoolConfig {
@@ -55,6 +56,7 @@ impl Default for PoolConfig {
             connection_config: Default::default(),
             pool_size: Default::default(),
             can_use_shard_aware_port: true,
+            keepalive_interval: None,
         }
     }
 }
@@ -137,6 +139,7 @@ pub struct NodeConnectionPool {
     conns: Arc<ArcSwap<MaybePoolConnections>>,
     use_keyspace_request_sender: mpsc::Sender<UseKeyspaceRequest>,
     _refiller_handle: RemoteHandle<()>,
+    _keepaliver_handle: Option<RemoteHandle<()>>,
     pool_updated_notify: Arc<Notify>,
 }
 
@@ -150,6 +153,8 @@ impl NodeConnectionPool {
         let (use_keyspace_request_sender, use_keyspace_request_receiver) = mpsc::channel(1);
         let pool_updated_notify = Arc::new(Notify::new());
 
+        let keepalive_interval = pool_config.keepalive_interval;
+
         let refiller = PoolRefiller::new(
             address,
             port,
@@ -159,13 +164,29 @@ impl NodeConnectionPool {
         );
 
         let conns = refiller.get_shared_connections();
-        let (fut, handle) = refiller.run(use_keyspace_request_receiver).remote_handle();
+        let (fut, refiller_handle) = refiller.run(use_keyspace_request_receiver).remote_handle();
         tokio::spawn(fut);
 
+        let keepaliver_handle = if let Some(interval) = keepalive_interval {
+            let keepaliver = Keepaliver {
+                connections: conns.clone(),
+                keepalive_interval: interval,
+                node_address: address,
+            };
+
+            let (fut, keepaliver_handle) = keepaliver.work().remote_handle();
+            tokio::spawn(fut);
+
+            Some(keepaliver_handle)
+        } else {
+            None
+        };
+
         Self {
             conns,
             use_keyspace_request_sender,
-            _refiller_handle: handle,
+            _refiller_handle: refiller_handle,
+            _keepaliver_handle: keepaliver_handle,
             pool_updated_notify,
         }
     }
@@ -315,6 +336,82 @@ impl NodeConnectionPool {
     }
 }
 
+struct Keepaliver {
+    connections: Arc<ArcSwap<MaybePoolConnections>>,
+    node_address: IpAddr, // This address is only used to enrich the log messages
+    keepalive_interval: Duration,
+}
+
+impl Keepaliver {
+    pub fn load_connections(&self) -> Vec<Arc<Connection>> {
+        use MaybePoolConnections::*;
+        use PoolConnections::*;
+
+        let pool = self.connections.load_full();
+        match &*pool {
+            Ready(NotSharded(conns)) => conns.clone(),
+            Ready(Sharded { connections, .. }) => connections.iter().flatten().cloned().collect(),
+            Initializing => vec![],
+            Broken => {
+                debug!("Cannot send connection keepalives for node {} as there are no alive connections in the pool", self.node_address);
+                vec![]
+            }
+        }
+    }
+
+    async fn work(self) {
+        let mut interval = tokio::time::interval(self.keepalive_interval);
+        interval.tick().await; // Use up the first, instant tick.
+
+        // Default behaviour (Burst) is not suitable for sending keepalives.
+        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+
+        // Wait for the second tick (so that `self.keepalive_interval` time passes since entering
+        // this function)
+        interval.tick().await;
+
+        loop {
+            let send_keepalives = self.send_keepalives();
+
+            tokio::select! {
+                _ = send_keepalives => {
+                    // Sending keepalives finished before receiving new tick.
+                    // Wait for the new tick and start over.
+                    interval.tick().await;
+                }
+                _ = interval.tick() => {
+                    // New tick arrived before `send_keepalives` was finished.
+                    // Stop polling `send_keepalives` and start over.
+                    //
+                    // `Interval::tick()` is cancellation safe, so it's ok to use it like that.
+                }
+            }
+        }
+    }
+
+    async fn send_keepalives(&self) {
+        let connections = self.load_connections();
+        let mut futures = connections
+            .into_iter()
+            .map(Self::send_keepalive_query)
+            .collect::<FuturesUnordered<_>>();
+
+        while futures.next().await.is_some() {}
+    }
+
+    async fn send_keepalive_query(connection: Arc<Connection>) {
+        if let Err(err) = connection
+            .query_single_page("select key from system.local where key = 'local'", &[])
+            .await
+        {
+            warn!(
+                "Failed to execute keepalive request on a connection {:p} - {}",
+                connection, err
+            );
+        }
+    }
+}
+
 const EXCESS_CONNECTION_BOUND_PER_SHARD_MULTIPLIER: usize = 10;
 
 // TODO: Make it configurable through a policy (issue #184)

scylla/src/transport/session.rs

@@ -107,12 +107,9 @@ pub struct SessionConfig {
 
     /// If true, full schema is fetched with every metadata refresh.
     pub fetch_schema_metadata: bool,
-    /*
-    These configuration options will be added in the future:
 
-
-    pub tcp_keepalive: bool,
-    */
+    /// Interval of sending keepalive requests
+    pub keepalive_interval: Option<Duration>,
 }
 
 /// Describes database server known on Session startup.
@@ -153,6 +150,7 @@ impl SessionConfig {
             disallow_shard_aware_port: false,
             default_consistency: Consistency::LocalQuorum,
             fetch_schema_metadata: true,
+            keepalive_interval: None,
         }
     }
 
@@ -220,6 +218,7 @@ impl SessionConfig {
             connection_config: self.get_connection_config(),
             pool_size: self.connection_pool_size.clone(),
             can_use_shard_aware_port: !self.disallow_shard_aware_port,
+            keepalive_interval: self.keepalive_interval,
         }
     }
 

scylla/src/transport/session_builder.rs

@@ -13,6 +13,7 @@ use std::time::Duration;
 use crate::statement::Consistency;
 #[cfg(feature = "ssl")]
 use openssl::ssl::SslContext;
+use tracing::warn;
 
 /// SessionBuilder is used to create new Session instances
 /// # Example
@@ -472,6 +473,33 @@ impl SessionBuilder {
         self.config.fetch_schema_metadata = fetch;
         self
     }
+
+    /// Set the keepalive interval.
+    /// The default is `None`, it corresponds to no keepalive messages being send.
+    ///
+    /// # Example
+    /// ```
+    /// # use scylla::{Session, SessionBuilder};
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let session: Session = SessionBuilder::new()
+    ///     .known_node("127.0.0.1:9042")
+    ///     .keepalive_interval(std::time::Duration::from_secs(42))
+    ///     .build()
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn keepalive_interval(mut self, interval: Duration) -> Self {
+        if interval <= Duration::from_secs(1) {
+            warn!(
+                "Setting the keepalive interval to low values ({:?}) is not recommended as it can have a negative impact on performance. Consider setting it above 1 second.",
+                interval
+            );
+        }
+
+        self.config.keepalive_interval = Some(interval);
+        self
+    }
 }
 
 /// Creates a [`SessionBuilder`] with default configuration, same as [`SessionBuilder::new`]

scylla/src/transport/topology.rs

@@ -15,13 +15,16 @@ use std::fmt::Formatter;
 use std::net::{IpAddr, SocketAddr};
 use std::num::NonZeroUsize;
 use std::str::FromStr;
+use std::time::Duration;
 use strum_macros::EnumString;
 use tokio::sync::mpsc;
 use tracing::{debug, error, trace, warn};
 
 /// Allows to read current metadata from the cluster
 pub(crate) struct MetadataReader {
     connection_config: ConnectionConfig,
+    keepalive_interval: Option<Duration>,
+
     control_connection_address: SocketAddr,
     control_connection: NodeConnectionPool,
 
@@ -153,6 +156,7 @@ impl MetadataReader {
     pub fn new(
         known_peers: &[SocketAddr],
         mut connection_config: ConnectionConfig,
+        keepalive_interval: Option<Duration>,
         server_event_sender: mpsc::Sender<Event>,
         fetch_schema: bool,
     ) -> Self {
@@ -168,11 +172,13 @@ impl MetadataReader {
         let control_connection = Self::make_control_connection_pool(
             control_connection_address,
             connection_config.clone(),
+            keepalive_interval,
         );
 
         MetadataReader {
             control_connection_address,
             control_connection,
+            keepalive_interval,
             connection_config,
             known_peers: known_peers.into(),
             fetch_schema,
@@ -222,6 +228,7 @@ impl MetadataReader {
             self.control_connection = Self::make_control_connection_pool(
                 self.control_connection_address,
                 self.connection_config.clone(),
+                self.keepalive_interval,
             );
 
             debug!(
@@ -263,9 +270,11 @@ impl MetadataReader {
     fn make_control_connection_pool(
         addr: SocketAddr,
         connection_config: ConnectionConfig,
+        keepalive_interval: Option<Duration>,
     ) -> NodeConnectionPool {
         let pool_config = PoolConfig {
             connection_config,
+            keepalive_interval,
 
             // We want to have only one connection to receive events from
             pool_size: PoolSize::PerHost(NonZeroUsize::new(1).unwrap()),