future: store result in OnceLock

muzarski · muzarski · commit 3925fa02475e · 2025-05-12T17:19:24.000+02:00
The result is going to be initialized only once, thus we do not need to
store it behind a mutex. We can use OnceLock instead. Thanks to that,
we can remove the unsafe logic which extends the lifetime of `coordinator`
reference in cass_future_coordinator. We now guarantee that
the result will be immutable once future is resolved - the guarantee
is provided on the type-level.
diff --git a/scylla-rust-wrapper/src/future.rs b/scylla-rust-wrapper/src/future.rs
@@ -9,14 +9,14 @@ use crate::query_result::{CassNode, CassResult};
 use crate::types::*;
 use crate::uuid::CassUuid;
 use futures::future;
-use scylla::response::Coordinator;
 use std::future::Future;
 use std::mem;
 use std::os::raw::c_void;
-use std::sync::{Arc, Condvar, Mutex};
+use std::sync::{Arc, Condvar, Mutex, OnceLock};
 use tokio::task::JoinHandle;
 use tokio::time::Duration;
 
+#[derive(Debug)]
 pub enum CassResultValue {
     Empty,
     QueryResult(Arc<CassResult>),
@@ -51,14 +51,14 @@ impl BoundCallback {
 
 #[derive(Default)]
 struct CassFutureState {
-    value: Option<CassFutureResult>,
     err_string: Option<String>,
     callback: Option<BoundCallback>,
     join_handle: Option<JoinHandle<()>>,
 }
 
 pub struct CassFuture {
     state: Mutex<CassFutureState>,
+    result: OnceLock<CassFutureResult>,
     wait_for_value: Condvar,
 }
 
@@ -88,14 +88,18 @@ impl CassFuture {
     ) -> Arc<CassFuture> {
         let cass_fut = Arc::new(CassFuture {
             state: Mutex::new(Default::default()),
+            result: OnceLock::new(),
             wait_for_value: Condvar::new(),
         });
         let cass_fut_clone = Arc::clone(&cass_fut);
         let join_handle = RUNTIME.spawn(async move {
             let r = fut.await;
             let maybe_cb = {
                 let mut guard = cass_fut_clone.state.lock().unwrap();
-                guard.value = Some(r);
+                cass_fut_clone
+                    .result
+                    .set(r)
+                    .expect("Tried to resolve future result twice!");
                 // Take the callback and call it after releasing the lock
                 guard.callback.take()
             };
@@ -116,16 +120,17 @@ impl CassFuture {
 
     pub fn new_ready(r: CassFutureResult) -> Arc<Self> {
         Arc::new(CassFuture {
-            state: Mutex::new(CassFutureState {
-                value: Some(r),
-                ..Default::default()
-            }),
+            state: Mutex::new(CassFutureState::default()),
+            result: OnceLock::from(r),
             wait_for_value: Condvar::new(),
         })
     }
 
-    pub fn with_waited_result<T>(&self, f: impl FnOnce(&mut CassFutureResult) -> T) -> T {
-        self.with_waited_state(|s| f(s.value.as_mut().unwrap()))
+    pub fn with_waited_result<'s, T>(&'s self, f: impl FnOnce(&'s CassFutureResult) -> T) -> T
+    where
+        T: 's,
+    {
+        self.with_waited_state(|_| f(self.result.get().unwrap()))
     }
 
     /// Awaits the future until completion.
@@ -154,7 +159,7 @@ impl CassFuture {
                 guard = self
                     .wait_for_value
                     .wait_while(guard, |state| {
-                        state.value.is_none() && state.join_handle.is_none()
+                        self.result.get().is_none() && state.join_handle.is_none()
                     })
                     // unwrap: Error appears only when mutex is poisoned.
                     .unwrap();
@@ -172,10 +177,10 @@ impl CassFuture {
 
     fn with_waited_result_timed<T>(
         &self,
-        f: impl FnOnce(&mut CassFutureResult) -> T,
+        f: impl FnOnce(&CassFutureResult) -> T,
         timeout_duration: Duration,
     ) -> Result<T, FutureError> {
-        self.with_waited_state_timed(|s| f(s.value.as_mut().unwrap()), timeout_duration)
+        self.with_waited_state_timed(|_| f(self.result.get().unwrap()), timeout_duration)
     }
 
     /// Tries to await the future with a given timeout.
@@ -243,7 +248,7 @@ impl CassFuture {
                 let (guard_result, timeout_result) = self
                     .wait_for_value
                     .wait_timeout_while(guard, remaining_timeout, |state| {
-                        state.value.is_none() && state.join_handle.is_none()
+                        self.result.get().is_none() && state.join_handle.is_none()
                     })
                     // unwrap: Error appears only when mutex is poisoned.
                     .unwrap();
@@ -276,7 +281,7 @@ impl CassFuture {
             return CassError::CASS_ERROR_LIB_CALLBACK_ALREADY_SET;
         }
         let bound_cb = BoundCallback { cb, data };
-        if lock.value.is_some() {
+        if self.result.get().is_some() {
             // The value is already available, we need to call the callback ourselves
             mem::drop(lock);
             bound_cb.invoke(self_ptr);
@@ -335,8 +340,12 @@ pub unsafe extern "C" fn cass_future_wait_timed(
 pub unsafe extern "C" fn cass_future_ready(
     future_raw: CassBorrowedSharedPtr<CassFuture, CMut>,
 ) -> cass_bool_t {
-    let state_guard = ArcFFI::as_ref(future_raw).unwrap().state.lock().unwrap();
-    match state_guard.value {
+    let Some(future) = ArcFFI::as_ref(future_raw) else {
+        tracing::error!("Provided null future to cass_future_ready!");
+        return cass_false;
+    };
+
+    match future.result.get() {
         None => cass_false,
         Some(_) => cass_true,
     }
@@ -348,7 +357,7 @@ pub unsafe extern "C" fn cass_future_error_code(
 ) -> CassError {
     ArcFFI::as_ref(future_raw)
         .unwrap()
-        .with_waited_result(|r: &mut CassFutureResult| match r {
+        .with_waited_result(|r: &CassFutureResult| match r {
             Ok(CassResultValue::QueryError(err)) => err.to_cass_error(),
             Err((err, _)) => *err,
             _ => CassError::CASS_OK,
@@ -361,19 +370,26 @@ pub unsafe extern "C" fn cass_future_error_message(
     message: *mut *const ::std::os::raw::c_char,
     message_length: *mut size_t,
 ) {
-    ArcFFI::as_ref(future)
-        .unwrap()
-        .with_waited_state(|state: &mut CassFutureState| {
-            let value = &state.value;
-            let msg = state
-                .err_string
-                .get_or_insert_with(|| match value.as_ref().unwrap() {
-                    Ok(CassResultValue::QueryError(err)) => err.msg(),
-                    Err((_, s)) => s.msg(),
-                    _ => "".to_string(),
-                });
-            unsafe { write_str_to_c(msg.as_str(), message, message_length) };
-        });
+    let Some(future) = ArcFFI::as_ref(future) else {
+        tracing::error!("Provided null future to cass_future_error_message!");
+        unsafe {
+            *message = std::ptr::null();
+            *message_length = 0;
+        }
+        return;
+    };
+
+    future.with_waited_state(|state: &mut CassFutureState| {
+        let value = future.result.get();
+        let msg = state
+            .err_string
+            .get_or_insert_with(|| match value.as_ref().unwrap() {
+                Ok(CassResultValue::QueryError(err)) => err.msg(),
+                Err((_, s)) => s.msg(),
+                _ => "".to_string(),
+            });
+        unsafe { write_str_to_c(msg.as_str(), message, message_length) };
+    });
 }
 
 #[unsafe(no_mangle)]
@@ -387,7 +403,7 @@ pub unsafe extern "C" fn cass_future_get_result(
 ) -> CassOwnedSharedPtr<CassResult, CConst> {
     ArcFFI::as_ref(future_raw)
         .unwrap()
-        .with_waited_result(|r: &mut CassFutureResult| -> Option<Arc<CassResult>> {
+        .with_waited_result(|r: &CassFutureResult| -> Option<Arc<CassResult>> {
             match r.as_ref().ok()? {
                 CassResultValue::QueryResult(qr) => Some(Arc::clone(qr)),
                 _ => None,
@@ -402,7 +418,7 @@ pub unsafe extern "C" fn cass_future_get_error_result(
 ) -> CassOwnedSharedPtr<CassErrorResult, CConst> {
     ArcFFI::as_ref(future_raw)
         .unwrap()
-        .with_waited_result(|r: &mut CassFutureResult| -> Option<Arc<CassErrorResult>> {
+        .with_waited_result(|r: &CassFutureResult| -> Option<Arc<CassErrorResult>> {
             match r.as_ref().ok()? {
                 CassResultValue::QueryError(qr) => Some(Arc::clone(qr)),
                 _ => None,
@@ -417,7 +433,7 @@ pub unsafe extern "C" fn cass_future_get_prepared(
 ) -> CassOwnedSharedPtr<CassPrepared, CConst> {
     ArcFFI::as_ref(future_raw)
         .unwrap()
-        .with_waited_result(|r: &mut CassFutureResult| -> Option<Arc<CassPrepared>> {
+        .with_waited_result(|r: &CassFutureResult| -> Option<Arc<CassPrepared>> {
             match r.as_ref().ok()? {
                 CassResultValue::Prepared(p) => Some(Arc::clone(p)),
                 _ => None,
@@ -433,7 +449,7 @@ pub unsafe extern "C" fn cass_future_tracing_id(
 ) -> CassError {
     ArcFFI::as_ref(future)
         .unwrap()
-        .with_waited_result(|r: &mut CassFutureResult| match r {
+        .with_waited_result(|r: &CassFutureResult| match r {
             Ok(CassResultValue::QueryResult(result)) => match result.tracing_id {
                 Some(id) => {
                     unsafe { *tracing_id = CassUuid::from(id) };
@@ -457,21 +473,7 @@ pub unsafe extern "C" fn cass_future_coordinator(
     future.with_waited_result(|r| match r {
         Ok(CassResultValue::QueryResult(result)) => {
             // unwrap: Coordinator is `None` only for tests.
-            let coordinator_ptr = result.coordinator.as_ref().unwrap() as *const Coordinator;
-
-            // We need to 'extend' the lifetime of returned Coordinator so safe FFI api does not complain.
-            // The lifetime of "result" reference provided to this closure is the lifetime of a mutex guard.
-            // We are guaranteed, that once the future is resolved (i.e. this closure is called), the result will not
-            // be modified in any way. Thus, we can guarantee that returned coordinator lives as long as underlying
-            // CassResult lives (i.e. longer than the lifetime of acquired mutex guard).
-            //
-            // SAFETY: Coordinator's lifetime is tied to the lifetime of underlying CassResult, thus:
-            // 1. Coordinator lives as long as the underlying CassResult lives
-            // 2. Coordinator will not be moved as long as underlying CassResult is not freed
-            // 3. Coordinator is immutable once future is resolved (because CassResult is set once)
-            let coordinator_ref = unsafe { &*coordinator_ptr };
-
-            RefFFI::as_ptr(coordinator_ref)
+            RefFFI::as_ptr(result.coordinator.as_ref().unwrap())
         }
         _ => RefFFI::null(),
     })
diff --git a/scylla-rust-wrapper/src/query_result.rs b/scylla-rust-wrapper/src/query_result.rs
@@ -29,17 +29,20 @@ use std::sync::Arc;
 use thiserror::Error;
 use uuid::Uuid;
 
+#[derive(Debug)]
 pub enum CassResultKind {
     NonRows,
     Rows(CassRowsResult),
 }
 
+#[derive(Debug)]
 pub struct CassRowsResult {
     // Arc: shared with first_row (yoke).
     pub(crate) shared_data: Arc<CassRowsResultSharedData>,
     pub(crate) first_row: Option<RowWithSelfBorrowedResultData>,
 }
 
+#[derive(Debug)]
 pub(crate) struct CassRowsResultSharedData {
     pub(crate) raw_rows: DeserializedMetadataAndRawRows,
     // Arc: shared with CassPrepared
@@ -53,6 +56,7 @@ impl FFI for CassNode {
     type Origin = FromRef;
 }
 
+#[derive(Debug)]
 pub struct CassResult {
     pub tracing_id: Option<Uuid>,
     pub paging_state_response: PagingStateResponse,
@@ -159,6 +163,7 @@ impl<'frame, 'metadata> DeserializeRow<'frame, 'metadata> for CassRawRow<'frame,
 
 /// The lifetime of CassRow is bound to CassResult.
 /// It will be freed, when CassResult is freed.(see #[cass_result_free])
+#[derive(Debug)]
 pub struct CassRow<'result> {
     pub columns: Vec<CassValue<'result>>,
     pub result_metadata: &'result CassResultMetadata,
@@ -230,7 +235,7 @@ mod row_with_self_borrowed_result_data {
 
     /// A simple wrapper over CassRow.
     /// Needed, so we can implement Yokeable for it, instead of implementing it for CassRow.
-    #[derive(Yokeable)]
+    #[derive(Yokeable, Debug)]
     struct CassRowWrapper<'result>(CassRow<'result>);
 
     /// A wrapper over struct which self-borrows the metadata allocated using Arc.
@@ -243,6 +248,7 @@ mod row_with_self_borrowed_result_data {
     ///
     /// This struct is a shared owner of the row bytes and metadata, and self-borrows this data
     /// to the `CassRow` it contains.
+    #[derive(Debug)]
     pub struct RowWithSelfBorrowedResultData(
         Yoke<CassRowWrapper<'static>, Arc<CassRowsResultSharedData>>,
     );
@@ -307,6 +313,7 @@ pub(crate) mod cass_raw_value {
     use scylla::errors::{DeserializationError, TypeCheckError};
     use thiserror::Error;
 
+    #[derive(Debug)]
     pub(crate) struct CassRawValue<'frame, 'metadata> {
         typ: &'metadata ColumnType<'metadata>,
         slice: Option<FrameSlice<'frame>>,
@@ -428,6 +435,7 @@ pub(crate) mod cass_raw_value {
     }
 }
 
+#[derive(Debug)]
 pub struct CassValue<'result> {
     pub(crate) value: CassRawValue<'result, 'result>,
     pub(crate) value_type: &'result Arc<CassDataType>,
