session,future: store & use cluster's tokio runtime

This commit finishes the changes required to support configuration
of the tokio runtime used by the Scylla Rust wrapper.
`cass_cluster_set_num_threads_io` now properly configures the number
of threads in the runtime, and the runtime is passed to all futures
created by the wrapper.
This allows the user to configure the runtime used by the wrapper
for each `CassCluster` independently, granting full flexibility.

Author: wprzytula

scylla-rust-wrapper/src/future.rs

@@ -1,4 +1,3 @@
-use crate::RUNTIME;
 use crate::argconv::*;
 use crate::cass_error::CassError;
 use crate::cass_error::CassErrorMessage;
@@ -68,6 +67,9 @@ enum FutureKind {
 }
 
 struct ResolvableFuture {
+    /// Runtime used to spawn and execute the future.
+    runtime: Arc<tokio::runtime::Runtime>,
+
     /// Mutable state of the future that requires synchronized exclusive access
     /// in order to ensure thread safety of the future execution.
     state: Mutex<CassFutureState>,
@@ -111,12 +113,14 @@ impl CassFuture {
     }
 
     pub(crate) fn make_raw(
+        runtime: Arc<tokio::runtime::Runtime>,
         fut: impl Future<Output = CassFutureResult> + Send + 'static,
         #[cfg(cpp_integration_testing)] recording_listener: Option<
             Arc<crate::integration_testing::RecordingHistoryListener>,
         >,
     ) -> CassOwnedSharedPtr<CassFuture, CMut> {
         Self::new_from_future(
+            runtime,
             fut,
             #[cfg(cpp_integration_testing)]
             recording_listener,
@@ -125,6 +129,7 @@ impl CassFuture {
     }
 
     pub(crate) fn new_from_future(
+        runtime: Arc<tokio::runtime::Runtime>,
         fut: impl Future<Output = CassFutureResult> + Send + 'static,
         #[cfg(cpp_integration_testing)] recording_listener: Option<
             Arc<crate::integration_testing::RecordingHistoryListener>,
@@ -133,6 +138,7 @@ impl CassFuture {
         let cass_fut = Arc::new(CassFuture {
             err_string: OnceLock::new(),
             kind: FutureKind::Resolvable(ResolvableFuture {
+                runtime: Arc::clone(&runtime),
                 state: Mutex::new(Default::default()),
                 result: OnceLock::new(),
                 wait_for_value: Condvar::new(),
@@ -141,7 +147,7 @@ impl CassFuture {
             }),
         });
         let cass_fut_clone = Arc::clone(&cass_fut);
-        let join_handle = RUNTIME.spawn(async move {
+        let join_handle = runtime.spawn(async move {
             let resolvable_fut = match cass_fut_clone.kind {
                 FutureKind::Resolvable(ref resolvable) => resolvable,
                 _ => unreachable!("CassFuture has been created as Resolvable"),
@@ -223,7 +229,7 @@ impl CassFuture {
                 // the future.
                 mem::drop(guard);
                 // unwrap: JoinError appears only when future either panic'ed or canceled.
-                RUNTIME.block_on(handle).unwrap();
+                resolvable_fut.runtime.block_on(handle).unwrap();
 
                 // Once we are here, the future is resolved.
                 // The result is guaranteed to be set.
@@ -313,7 +319,7 @@ impl CassFuture {
                         future::Either::Right((_, handle)) => Err(JoinHandleTimeout(handle)),
                     }
                 };
-                match RUNTIME.block_on(timed) {
+                match resolvable_fut.runtime.block_on(timed) {
                     Err(JoinHandleTimeout(returned_handle)) => {
                         // We timed out. so we can't finish waiting for the future.
                         // The problem is that if current thread executor is used,
@@ -638,6 +644,15 @@ mod tests {
         time::Duration,
     };
 
+    fn runtime_for_test() -> Arc<tokio::runtime::Runtime> {
+        Arc::new(
+            tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap(),
+        )
+    }
+
     // This is not a particularly smart test, but if some thread is granted access the value
     // before it is truly computed, then weird things should happen, even a segfault.
     // In the incorrect implementation that inspired this test to be written, this test
@@ -646,11 +661,13 @@ mod tests {
     #[ntest::timeout(100)]
     fn cass_future_thread_safety() {
         const ERROR_MSG: &str = "NOBODY EXPECTED SPANISH INQUISITION";
+        let runtime = runtime_for_test();
         let fut = async {
             tokio::time::sleep(Duration::from_millis(10)).await;
             Err((CassError::CASS_OK, ERROR_MSG.into()))
         };
         let cass_fut = CassFuture::make_raw(
+            runtime,
             fut,
             #[cfg(cpp_integration_testing)]
             None,
@@ -685,11 +702,13 @@ mod tests {
     fn cass_future_resolves_after_timeout() {
         const ERROR_MSG: &str = "NOBODY EXPECTED SPANISH INQUISITION";
         const HUNDRED_MILLIS_IN_MICROS: u64 = 100 * 1000;
+        let runtime = runtime_for_test();
         let fut = async move {
             tokio::time::sleep(Duration::from_micros(HUNDRED_MILLIS_IN_MICROS)).await;
             Err((CassError::CASS_OK, ERROR_MSG.into()))
         };
         let cass_fut = CassFuture::make_raw(
+            runtime,
             fut,
             #[cfg(cpp_integration_testing)]
             None,
@@ -725,6 +744,8 @@ mod tests {
         const ERROR_MSG: &str = "NOBODY EXPECTED SPANISH INQUISITION";
         const HUNDRED_MILLIS_IN_MICROS: u64 = 100 * 1000;
 
+        let runtime = runtime_for_test();
+
         let create_future_and_flag = || {
             unsafe extern "C" fn mark_flag_cb(
                 _fut: CassBorrowedSharedPtr<CassFuture, CMut>,
@@ -741,6 +762,7 @@ mod tests {
                 Err((CassError::CASS_OK, ERROR_MSG.into()))
             };
             let cass_fut = CassFuture::make_raw(
+                Arc::clone(&runtime),
                 fut,
                 #[cfg(cpp_integration_testing)]
                 None,
@@ -816,7 +838,7 @@ mod tests {
         {
             let (cass_fut, flag_ptr) = create_future_and_flag();
 
-            RUNTIME.block_on(async {
+            runtime.block_on(async {
                 tokio::time::sleep(Duration::from_micros(HUNDRED_MILLIS_IN_MICROS + 10 * 1000))
                     .await
             });

scylla-rust-wrapper/src/lib.rs

@@ -4,7 +4,6 @@ use crate::logging::Logger;
 use crate::logging::stderr_log_callback;
 use std::sync::LazyLock;
 use std::sync::RwLock;
-use tokio::runtime::Runtime;
 
 #[macro_use]
 mod binding;
@@ -191,7 +190,6 @@ pub(crate) mod cass_version_types {
     include_bindgen_generated!("cppdriver_version_types.rs");
 }
 
-pub(crate) static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
 pub(crate) static LOGGER: LazyLock<RwLock<Logger>> = LazyLock::new(|| {
     RwLock::new(Logger {
         cb: Some(stderr_log_callback),

scylla-rust-wrapper/src/session.rs

@@ -31,6 +31,7 @@ use std::sync::Arc;
 use tokio::sync::RwLock;
 
 pub(crate) struct CassConnectedSession {
+    runtime: Arc<tokio::runtime::Runtime>,
     session: Session,
     exec_profile_map: HashMap<ExecProfileName, ExecutionProfileHandle>,
 }
@@ -78,6 +79,7 @@ impl CassConnectedSession {
         let cluster_client_id = cluster.get_client_id();
 
         let fut = Self::connect_fut(
+            Arc::clone(cluster.get_runtime()),
             session,
             session_builder,
             cluster_client_id,
@@ -87,13 +89,15 @@ impl CassConnectedSession {
         );
 
         CassFuture::make_raw(
+            Arc::clone(cluster.get_runtime()),
             fut,
             #[cfg(cpp_integration_testing)]
             None,
         )
     }
 
     async fn connect_fut(
+        runtime: Arc<tokio::runtime::Runtime>,
         session: Arc<CassSession>,
         session_builder_fut: impl Future<Output = SessionBuilder>,
         cluster_client_id: Option<uuid::Uuid>,
@@ -152,20 +156,40 @@ impl CassConnectedSession {
             .map_err(|err| (err.to_cass_error(), err.msg()))?;
 
         session_guard.connected = Some(CassConnectedSession {
+            runtime,
             session,
             exec_profile_map,
         });
         Ok(CassResultValue::Empty)
     }
 
     fn close_fut(session_opt: Arc<RwLock<CassSessionInner>>) -> Arc<CassFuture> {
+        let runtime = {
+            let Ok(session_guard) = session_opt.try_read() else {
+                return CassFuture::new_ready(Err((
+                    CassError::CASS_ERROR_LIB_NO_HOSTS_AVAILABLE,
+                    "Still connecting or already closing".msg(),
+                )));
+            };
+
+            let Some(connected_session) = session_guard.connected.as_ref() else {
+                return CassFuture::new_ready(Err((
+                    CassError::CASS_ERROR_LIB_NO_HOSTS_AVAILABLE,
+                    "Session is not connected".msg(),
+                )));
+            };
+
+            Arc::clone(&connected_session.runtime)
+        };
+
         CassFuture::new_from_future(
+            runtime,
             async move {
                 let mut session_guard = session_opt.write().await;
                 if session_guard.connected.is_none() {
                     return Err((
                         CassError::CASS_ERROR_LIB_UNABLE_TO_CLOSE,
-                        "Already closing or closed".msg(),
+                        "Session is not connected".msg(),
                     ));
                 }
 
@@ -286,6 +310,8 @@ pub unsafe extern "C" fn cass_session_execute_batch(
         )));
     };
 
+    let runtime = Arc::clone(&connected_session.runtime);
+
     let future = async move {
         let connected_session = session_guard
             .connected
@@ -315,6 +341,7 @@ pub unsafe extern "C" fn cass_session_execute_batch(
     };
 
     CassFuture::make_raw(
+        runtime,
         future,
         #[cfg(cpp_integration_testing)]
         None,
@@ -350,6 +377,8 @@ pub unsafe extern "C" fn cass_session_execute(
         )));
     };
 
+    let runtime = Arc::clone(&connected_session.runtime);
+
     let paging_state = statement_opt.paging_state.clone();
     let paging_enabled = statement_opt.paging_enabled;
     let mut statement = statement_opt.statement.clone();
@@ -483,6 +512,7 @@ pub unsafe extern "C" fn cass_session_execute(
     };
 
     CassFuture::make_raw(
+        runtime,
         future,
         #[cfg(cpp_integration_testing)]
         recording_listener,
@@ -517,9 +547,12 @@ pub unsafe extern "C" fn cass_session_prepare_from_existing(
         )));
     };
 
+    let runtime = Arc::clone(&connected_session.runtime);
+
     let statement = cass_statement.statement.clone();
 
     CassFuture::make_raw(
+        runtime,
         async move {
             let query = match &statement {
                 BoundStatement::Simple(q) => q,
@@ -588,6 +621,8 @@ pub unsafe extern "C" fn cass_session_prepare_n(
         )));
     };
 
+    let runtime = Arc::clone(&connected_session.runtime);
+
     let query = Statement::new(query_str.to_string());
 
     let fut = async move {
@@ -608,6 +643,7 @@ pub unsafe extern "C" fn cass_session_prepare_n(
     };
 
     CassFuture::make_raw(
+        runtime,
         fut,
         #[cfg(cpp_integration_testing)]
         None,