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. 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.
3. Once all `CassCluster` 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 creates a `CassCluster` with a
runtime with `i` threads and then drops it, will lead to space leaks.
This is an artificial case, though.

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()
+        }
+        .expect("Failed to create an async runtime")
     }
 
     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 `CassCluster` instances that reference a runtime are dropped,
+/// the runtime is also dropped.
+pub(crate) struct Runtimes {
+    // Weak pointers are used to make runtimes dropped once all CassCluster 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,