future: spawn_blocking when waiting on cond var

Author: muzarski

scylla-rust-wrapper/src/future.rs

@@ -108,61 +108,96 @@ impl CassFuture {
         })
     }
 
-    pub fn with_waited_result<T>(&self, f: impl FnOnce(&mut CassFutureResult) -> T) -> T {
+    pub fn with_waited_result<T>(
+        self: &Arc<Self>,
+        f: impl FnOnce(&mut CassFutureResult) -> T,
+    ) -> T {
         self.with_waited_state(|s| f(s.value.as_mut().unwrap()))
     }
 
-    fn with_waited_state<T>(&self, f: impl FnOnce(&mut CassFutureState) -> T) -> T {
+    fn with_waited_state<T>(self: &Arc<Self>, f: impl FnOnce(&mut CassFutureState) -> T) -> T {
         let mut guard = self.state.lock().unwrap();
         let handle = guard.join_handle.take();
+        mem::drop(guard);
+
         if let Some(handle) = handle {
-            mem::drop(guard);
             // unwrap: JoinError appears only when future either panic'ed or canceled.
             RUNTIME.block_on(handle).unwrap();
             guard = self.state.lock().unwrap();
         } else {
-            guard = self
-                .wait_for_value
-                .wait_while(guard, |state| state.value.is_none())
-                // unwrap: Error appears only when mutex is poisoned.
-                .unwrap();
+            // We need to spawn a thread that will be responsible for waiting
+            // on a cond variable (this operation blocks the thread).
+            // If we don't do that, this won't work for `current-thread` tokio runtime..
+            let self_clone = Arc::clone(self);
+            let wait_for_value = async {
+                RUNTIME
+                    .spawn_blocking(move || {
+                        let mut guard = self_clone.state.lock().unwrap();
+                        guard = self_clone
+                            .wait_for_value
+                            .wait_while(guard, |state| state.value.is_none())
+                            // unwrap: Error appears only when mutex is poisoned.
+                            .unwrap();
+                        mem::drop(guard);
+                    })
+                    .await
+            };
+            RUNTIME.block_on(wait_for_value).unwrap();
+            guard = self.state.lock().unwrap();
         }
         f(&mut guard)
     }
 
     fn with_waited_result_timed<T>(
-        &self,
+        self: &Arc<Self>,
         f: impl FnOnce(&mut CassFutureResult) -> T,
         timeout_duration: Duration,
     ) -> Result<T, TimeoutError> {
         self.with_waited_state_timed(|s| f(s.value.as_mut().unwrap()), timeout_duration)
     }
 
     pub(self) fn with_waited_state_timed<T>(
-        &self,
+        self: &Arc<Self>,
         f: impl FnOnce(&mut CassFutureState) -> T,
         timeout_duration: Duration,
     ) -> Result<T, TimeoutError> {
         let mut guard = self.state.lock().unwrap();
         let handle = guard.join_handle.take();
+        mem::drop(guard);
+
         if let Some(handle) = handle {
-            mem::drop(guard);
             // Need to wrap it with async{} block, so the timeout is lazily executed inside the runtime.
             // See mention about panics: https://docs.rs/tokio/latest/tokio/time/fn.timeout.html.
             let timed = async { tokio::time::timeout(timeout_duration, handle).await };
             // unwrap: JoinError appears only when future either panic'ed or canceled.
             RUNTIME.block_on(timed).map_err(|_| TimeoutError)?.unwrap();
             guard = self.state.lock().unwrap();
         } else {
-            let (guard_result, timeout_result) = self
-                .wait_for_value
-                .wait_timeout_while(guard, timeout_duration, |state| state.value.is_none())
-                // unwrap: Error appears only when mutex is poisoned.
-                .unwrap();
+            // We need to spawn a thread that will be responsible for waiting
+            // on a cond variable (this operation blocks the thread).
+            // If we don't do that, this won't work for `current-thread` tokio runtime..
+            let self_clone = Arc::clone(self);
+            let wait_for_value = async {
+                RUNTIME
+                    .spawn_blocking(move || {
+                        let guard = self_clone.state.lock().unwrap();
+                        let (guard_result, timeout_result) = self_clone
+                            .wait_for_value
+                            .wait_timeout_while(guard, timeout_duration, |state| {
+                                state.value.is_none()
+                            })
+                            // unwrap: Error appears only when mutex is poisoned.
+                            .unwrap();
+                        mem::drop(guard_result);
+                        timeout_result
+                    })
+                    .await
+            };
+            let timeout_result = RUNTIME.block_on(wait_for_value).unwrap();
             if timeout_result.timed_out() {
                 return Err(TimeoutError);
             }
-            guard = guard_result;
+            guard = self.state.lock().unwrap();
         }
 
         Ok(f(&mut guard))
@@ -208,15 +243,15 @@ pub unsafe extern "C" fn cass_future_set_callback(
 
 #[no_mangle]
 pub unsafe extern "C" fn cass_future_wait(future_raw: *const CassFuture) {
-    ptr_to_ref(future_raw).with_waited_result(|_| ());
+    clone_arced(future_raw).with_waited_result(|_| ());
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn cass_future_wait_timed(
     future_raw: *const CassFuture,
     timeout_us: cass_duration_t,
 ) -> cass_bool_t {
-    ptr_to_ref(future_raw)
+    clone_arced(future_raw)
         .with_waited_result_timed(|_| (), Duration::from_micros(timeout_us))
         .is_ok() as cass_bool_t
 }
@@ -232,7 +267,7 @@ pub unsafe extern "C" fn cass_future_ready(future_raw: *const CassFuture) -> cas
 
 #[no_mangle]
 pub unsafe extern "C" fn cass_future_error_code(future_raw: *const CassFuture) -> CassError {
-    ptr_to_ref(future_raw).with_waited_result(|r: &mut CassFutureResult| match r {
+    clone_arced(future_raw).with_waited_result(|r: &mut CassFutureResult| match r {
         Ok(CassResultValue::QueryError(err)) => CassError::from(err.as_ref()),
         Err((err, _)) => *err,
         _ => CassError::CASS_OK,
@@ -245,7 +280,7 @@ pub unsafe extern "C" fn cass_future_error_message(
     message: *mut *const ::std::os::raw::c_char,
     message_length: *mut size_t,
 ) {
-    ptr_to_ref(future).with_waited_state(|state: &mut CassFutureState| {
+    clone_arced(future).with_waited_state(|state: &mut CassFutureState| {
         let value = &state.value;
         let msg = state
             .err_string
@@ -267,7 +302,7 @@ pub unsafe extern "C" fn cass_future_free(future_raw: *const CassFuture) {
 pub unsafe extern "C" fn cass_future_get_result(
     future_raw: *const CassFuture,
 ) -> *const CassResult {
-    ptr_to_ref(future_raw)
+    clone_arced(future_raw)
         .with_waited_result(|r: &mut CassFutureResult| -> Option<Arc<CassResult>> {
             match r.as_ref().ok()? {
                 CassResultValue::QueryResult(qr) => Some(qr.clone()),
@@ -281,7 +316,7 @@ pub unsafe extern "C" fn cass_future_get_result(
 pub unsafe extern "C" fn cass_future_get_error_result(
     future_raw: *const CassFuture,
 ) -> *const CassErrorResult {
-    ptr_to_ref(future_raw)
+    clone_arced(future_raw)
         .with_waited_result(|r: &mut CassFutureResult| -> Option<Arc<CassErrorResult>> {
             match r.as_ref().ok()? {
                 CassResultValue::QueryError(qr) => Some(qr.clone()),
@@ -295,7 +330,7 @@ pub unsafe extern "C" fn cass_future_get_error_result(
 pub unsafe extern "C" fn cass_future_get_prepared(
     future_raw: *mut CassFuture,
 ) -> *const CassPrepared {
-    ptr_to_ref(future_raw)
+    clone_arced(future_raw)
         .with_waited_result(|r: &mut CassFutureResult| -> Option<Arc<CassPrepared>> {
             match r.as_ref().ok()? {
                 CassResultValue::Prepared(p) => Some(p.clone()),
@@ -310,7 +345,7 @@ pub unsafe extern "C" fn cass_future_tracing_id(
     future: *const CassFuture,
     tracing_id: *mut CassUuid,
 ) -> CassError {
-    ptr_to_ref(future).with_waited_result(|r: &mut CassFutureResult| match r {
+    clone_arced(future).with_waited_result(|r: &mut CassFutureResult| match r {
         Ok(CassResultValue::QueryResult(result)) => match result.metadata.tracing_id {
             Some(id) => {
                 *tracing_id = CassUuid::from(id);

scylla-rust-wrapper/src/lib.rs

@@ -36,7 +36,12 @@ pub mod uuid;
 pub mod value;
 
 lazy_static! {
-    pub static ref RUNTIME: Runtime = Runtime::new().unwrap();
+    // TODO: Revert to multi-thread. It's just for the sake of review.
+    // To check that tests pass with current-thread runtime.
+    pub static ref RUNTIME: Runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
     pub static ref LOGGER: RwLock<Logger> = RwLock::new(Logger {
         cb: Some(stderr_log_callback),
         data: std::ptr::null_mut(),