runtime: introduce cache for tokio runtimes

As @Lorak-mmk noted in the review, runtimes are currently not shared
between `CassCluster` instances, which leads to possible many tokio
runtimes being created in the application, with possibly a lot
of threads. This commit introduces a cache for tokio runtimes, which
is encapsulated in the global `Runtimes` struct.
1. `CassCluster` now does not store a `tokio::runtime::Runtime`
   directly, but rather an optional number of threads in the runtime.
2. The `Runtimes` struct is a global cache for tokio runtimes.
   It allows to get a default runtime or a runtime with a specified
   number of threads. Upon `cass_session_connect`, if a runtime is not
   created yet, it will create a new one and cache it for future use.
   The handling of the cache is fully transparent to the user of the
   abstraction. `CassSession` since then holds an `Arc<Runtime>`.
3. Once all `CassSession` instances that reference a runtime are
   dropped, the runtime is also dropped. This is done by storing weak
   pointers to runtimes in the `Runtimes` struct.

Interesting to note: as Weak pointers keep the Arc allocation alive,
a workflow that for consecutive `i`s connects a `CassSession` with a
runtime with `i` threads and then drops it, will lead to space leaks.
This is an artificial case, though.
Remember that while the allocation will be still kept alive, the
runtime itself will not be running, as it is dropped when the last
`CassSession` referencing it is dropped.

Author: wprzytula

scylla-rust-wrapper/src/cluster.rs

@@ -7,6 +7,7 @@ use crate::load_balancing::{
     CassHostFilter, DcRestriction, LoadBalancingConfig, LoadBalancingKind,
 };
 use crate::retry_policy::CassRetryPolicy;
+use crate::runtime::RUNTIMES;
 use crate::ssl::CassSsl;
 use crate::timestamp_generator::CassTimestampGen;
 use crate::types::*;
@@ -82,7 +83,11 @@ const DRIVER_NAME: &str = "ScyllaDB Cpp-Rust Driver";
 const DRIVER_VERSION: &str = env!("CARGO_PKG_VERSION");
 
 pub struct CassCluster {
-    runtime: Arc<tokio::runtime::Runtime>,
+    /// Number of threads in the tokio runtime thread pool.
+    ///
+    /// Specified with `cass_cluster_set_num_threads_io`.
+    /// If not set, the default tokio runtime is used.
+    num_threads_io: Option<usize>,
 
     session_builder: SessionBuilder,
     default_execution_profile_builder: ExecutionProfileBuilder,
@@ -101,8 +106,20 @@ pub struct CassCluster {
 }
 
 impl CassCluster {
-    pub(crate) fn get_runtime(&self) -> &Arc<tokio::runtime::Runtime> {
-        &self.runtime
+    /// Gets the runtime that has been set for the cluster.
+    /// If no runtime has been set yet, it creates a default runtime
+    /// and makes it cached in the global `Runtimes` instance.
+    pub(crate) fn get_runtime(&self) -> Arc<tokio::runtime::Runtime> {
+        let mut runtimes = RUNTIMES.lock().unwrap();
+
+        if let Some(num_threads_io) = self.num_threads_io {
+            // If the number of threads is set, we create a runtime with that number of threads.
+            runtimes.n_thread_runtime(num_threads_io)
+        } else {
+            // Otherwise, we use the default runtime.
+            runtimes.default_runtime()
+        }
+        .unwrap_or_else(|err| panic!("Failed to create an async runtime: {err}"))
     }
 
     pub(crate) fn execution_profile_map(&self) -> &HashMap<ExecProfileName, CassExecProfile> {
@@ -184,12 +201,6 @@ impl CassCluster {
 
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn cass_cluster_new() -> CassOwnedExclusivePtr<CassCluster, CMut> {
-    let Ok(default_runtime) = tokio::runtime::Runtime::new()
-        .inspect_err(|e| tracing::error!("Failed to create async runtime: {}", e))
-    else {
-        return CassPtr::null_mut();
-    };
-
     let default_execution_profile_builder = ExecutionProfileBuilder::default()
         .consistency(DEFAULT_CONSISTENCY)
         .serial_consistency(DEFAULT_SERIAL_CONSISTENCY)
@@ -321,7 +332,7 @@ pub unsafe extern "C" fn cass_cluster_new() -> CassOwnedExclusivePtr<CassCluster
     };
 
     BoxFFI::into_ptr(Box::new(CassCluster {
-        runtime: Arc::new(default_runtime),
+        num_threads_io: None,
 
         session_builder: default_session_builder,
         port: 9042,
@@ -1554,25 +1565,7 @@ pub unsafe extern "C" fn cass_cluster_set_num_threads_io(
         return CassError::CASS_ERROR_LIB_BAD_PARAMS;
     };
 
-    let runtime_res = match num_threads {
-        0 => tokio::runtime::Builder::new_current_thread()
-            .enable_all()
-            .build(),
-        n => tokio::runtime::Builder::new_multi_thread()
-            .worker_threads(n as usize)
-            .enable_all()
-            .build(),
-    };
-
-    let runtime = match runtime_res {
-        Ok(runtime) => runtime,
-        Err(err) => {
-            tracing::error!("Failed to create async runtime: {}", err);
-            return CassError::CASS_ERROR_LIB_BAD_PARAMS;
-        }
-    };
-
-    cluster.runtime = Arc::new(runtime);
+    cluster.num_threads_io = Some(num_threads as usize);
 
     CassError::CASS_OK
 }

scylla-rust-wrapper/src/lib.rs

@@ -30,6 +30,7 @@ pub(crate) mod misc;
 pub(crate) mod prepared;
 pub(crate) mod query_result;
 pub(crate) mod retry_policy;
+pub(crate) mod runtime;
 #[cfg(test)]
 mod ser_de_tests;
 pub(crate) mod session;

scylla-rust-wrapper/src/runtime.rs

@@ -0,0 +1,81 @@
+//! Manages tokio runtimes for the application.
+//!
+//! Runtime is per-cluster and can be changed with `cass_cluster_set_num_threads_io`.
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, Weak},
+};
+
+use tokio::runtime::Runtime;
+
+/// Manages tokio runtimes for the application.
+///
+/// Runtime is per-cluster and can be changed with `cass_cluster_set_num_threads_io`.
+/// Once a runtime is created, it is cached for future use.
+/// Once all `CassSession` instances that reference the runtime are dropped,
+/// the runtime is also dropped.
+pub(crate) struct Runtimes {
+    // Weak pointers are used to make runtimes dropped once all `CassSession` instances
+    // that reference them are freed.
+    default_runtime: Option<Weak<Runtime>>,
+    // This is Option to allow creating a static instance of Runtimes.
+    // (`HashMap::new` is not `const`).
+    n_thread_runtimes: Option<HashMap<usize, Weak<Runtime>>>,
+}
+
+pub(crate) static RUNTIMES: std::sync::Mutex<Runtimes> = {
+    std::sync::Mutex::new(Runtimes {
+        default_runtime: None,
+        n_thread_runtimes: None,
+    })
+};
+
+impl Runtimes {
+    fn cached_or_new_runtime(
+        weak_runtime: &mut Weak<Runtime>,
+        create_runtime: impl FnOnce() -> Result<Arc<Runtime>, std::io::Error>,
+    ) -> Result<Arc<Runtime>, std::io::Error> {
+        match weak_runtime.upgrade() {
+            Some(cached_runtime) => Ok(cached_runtime),
+            None => {
+                let runtime = create_runtime()?;
+                *weak_runtime = Arc::downgrade(&runtime);
+                Ok(runtime)
+            }
+        }
+    }
+
+    /// Returns a default tokio runtime.
+    ///
+    /// If it's not created yet, it will create a new one with the default configuration
+    /// and cache it for future use.
+    pub(crate) fn default_runtime(&mut self) -> Result<Arc<Runtime>, std::io::Error> {
+        let default_runtime_slot = self.default_runtime.get_or_insert_with(Weak::new);
+        Self::cached_or_new_runtime(default_runtime_slot, || Runtime::new().map(Arc::new))
+    }
+
+    /// Returns a tokio runtime with `n_threads` worker threads.
+    ///
+    /// If it's not created yet, it will create a new one and cache it for future use.
+    pub(crate) fn n_thread_runtime(
+        &mut self,
+        n_threads: usize,
+    ) -> Result<Arc<Runtime>, std::io::Error> {
+        let n_thread_runtimes = self.n_thread_runtimes.get_or_insert_with(HashMap::new);
+        let n_thread_runtime_slot = n_thread_runtimes.entry(n_threads).or_default();
+
+        Self::cached_or_new_runtime(n_thread_runtime_slot, || {
+            match n_threads {
+                0 => tokio::runtime::Builder::new_current_thread()
+                    .enable_all()
+                    .build(),
+                n => tokio::runtime::Builder::new_multi_thread()
+                    .worker_threads(n)
+                    .enable_all()
+                    .build(),
+            }
+            .map(Arc::new)
+        })
+    }
+}

scylla-rust-wrapper/src/session.rs

@@ -79,7 +79,7 @@ impl CassConnectedSession {
         let cluster_client_id = cluster.get_client_id();
 
         let fut = Self::connect_fut(
-            Arc::clone(cluster.get_runtime()),
+            cluster.get_runtime(),
             session,
             session_builder,
             cluster_client_id,
@@ -89,7 +89,7 @@ impl CassConnectedSession {
         );
 
         CassFuture::make_raw(
-            Arc::clone(cluster.get_runtime()),
+            cluster.get_runtime(),
             fut,
             #[cfg(cpp_integration_testing)]
             None,