This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new a43f0fecd fix(rust/driver_manager): allow cancel to be truly
concurrent (#4858)
a43f0fecd is described below
commit a43f0fecda799a7d6aa3efbf454aa5a97630d865
Author: David Li <[email protected]>
AuthorDate: Thu Oct 8 15:39:13 2026 -0700
fix(rust/driver_manager): allow cancel to be truly concurrent (#4858)
N.B. the driver exporter is not covered by this PR.
Closes #4817.
Assisted-by: GPT-6 Astra
---
rust/driver_manager/src/lib.rs | 296 ++++++++++--------
rust/driver_manager/tests/cancellation.rs | 495 ++++++++++++++++++++++++++++++
rust/ffi/src/options.rs | 38 +++
3 files changed, 704 insertions(+), 125 deletions(-)
diff --git a/rust/driver_manager/src/lib.rs b/rust/driver_manager/src/lib.rs
index 811886828..f42bcdc58 100644
--- a/rust/driver_manager/src/lib.rs
+++ b/rust/driver_manager/src/lib.rs
@@ -93,17 +93,21 @@
// According to the ADBC specification, objects allow serialized access from
// multiple threads: one thread may make a call, and once finished, another
// thread may make a call. They do not allow concurrent access from multiple
-// threads.
+// threads, except that cancellation must be able to overlap ordinary calls.
//
-// In order to implement these semantics, all mutable FFI objects are wrapped
-// in `Mutex`es. `FFI_Driver` is not wrapped in a `Mutex` because it is
-// an immutable struct of function pointers. Wrapping the driver in a `Mutex`
-// would prevent any parallelism between driver calls, which is not desirable.
+// Databases use a `Mutex`. Connections and statements use `NativeHandle` to
+// serialize ordinary calls while allowing cancellation through raw pointers.
+// Their shared inner objects release the native handles only after all calls
+// (including cancellation) have relinquished their strong references.
+// `FFI_Driver` is not wrapped in a `Mutex` because it is an immutable struct
+// of function pointers. Wrapping the driver in a `Mutex` would prevent any
+// parallelism between driver calls, which is not desirable.
pub mod error;
pub mod profile;
pub mod search;
+use std::cell::UnsafeCell;
use std::collections::HashSet;
use std::ffi::{CString, OsStr};
use std::ops::DerefMut;
@@ -111,11 +115,11 @@ use std::os::raw::c_char;
use std::path::PathBuf;
use std::pin::Pin;
use std::ptr::{null, null_mut};
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, LockResult, Mutex, MutexGuard};
use adbc_ffi::options::{
- check_status, get_option_bytes, get_option_string, set_option_connection,
set_option_database,
- set_option_statement,
+ check_status, get_option_bytes, get_option_string, set_option_connection,
+ set_option_connection_raw, set_option_database, set_option_statement_raw,
};
use arrow_array::ffi::{FFI_ArrowSchema, to_ffi};
use arrow_array::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
@@ -758,7 +762,7 @@ impl Database for ManagedDatabase {
// Initialize the connection.
let connection = self.connection_init(connection)?;
let inner = ManagedConnectionInner {
- connection: Mutex::new(connection),
+ connection: NativeHandle::new(connection),
database: self.inner.clone(),
};
Ok(Self::ConnectionType {
@@ -782,7 +786,7 @@ impl Database for ManagedDatabase {
// Initialize the connection.
let connection = self.connection_init(connection)?;
let inner = ManagedConnectionInner {
- connection: Mutex::new(connection),
+ connection: NativeHandle::new(connection),
database: self.inner.clone(),
};
Ok(Self::ConnectionType {
@@ -791,19 +795,52 @@ impl Database for ManagedDatabase {
}
}
+// Wrap a native handle with a mutex and an UnsafeCell. We need shared access
+// to the native handle for cancellation, so just a mutex isn't sufficient.
+struct NativeHandle<T> {
+ value: UnsafeCell<T>,
+ mu: Mutex<()>,
+}
+
+// SAFETY: Used only for connection and statement FFI handles. Ordinary access
is
+// serialized by operation, and the only overlapping calls are cancellation,
+// which the driver must make thread-safe under the ADBC contract. Raw pointers
+// avoid creating aliased Rust references. The enclosing Arc keeps the handle,
+// its parents, and the driver library alive until all calls finish. Release
runs
+// only in the enclosing inner object's Drop, with no concurrent access.
+unsafe impl<T: Send> Sync for NativeHandle<T> {}
+
+impl<T> NativeHandle<T> {
+ fn new(value: T) -> Self {
+ Self {
+ value: UnsafeCell::new(value),
+ mu: Mutex::new(()),
+ }
+ }
+
+ fn lock(&self) -> LockResult<MutexGuard<'_, ()>> {
+ self.mu.lock()
+ }
+
+ fn get(&self) -> *mut T {
+ self.value.get()
+ }
+}
+
struct ManagedConnectionInner {
- connection: Mutex<adbc_ffi::FFI_AdbcConnection>,
+ connection: NativeHandle<adbc_ffi::FFI_AdbcConnection>,
database: Arc<ManagedDatabaseInner>,
}
impl Drop for ManagedConnectionInner {
fn drop(&mut self) {
let driver = &self.database.driver.driver;
- if let Ok(mut connection) = self.connection.lock() {
+ if !self.connection.mu.is_poisoned() {
+ let connection = self.connection.get();
let method = driver_method!(driver, ConnectionRelease);
// TODO(alexandreyc): how should we handle `ConnectionRelease`
failing?
// See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
- unsafe { method(connection.deref_mut(), null_mut()) };
+ unsafe { method(connection, null_mut()) };
}
// We could still drop here but if the lock is poisoned, we have no
// clue what the status is. Since a panic comes from Rust code,
@@ -832,15 +869,17 @@ impl adbc_core::CancelHandle for ConnectionCancelHandle {
));
}
let driver = &inner.database.driver.driver;
- let mut connection = inner.connection.lock().map_err(|e| {
- Error::with_message_and_status(
- format!("[Driver Manager] connection is poisoned: {e:?}"),
- Status::Internal,
- )
- })?;
+ let Some(method) = driver.ConnectionCancel else {
+ return Err(Error::with_message_and_status(
+ "[Driver Manager] ConnectionCancel not supported by
driver",
+ Status::NotImplemented,
+ ));
+ };
+ let connection = inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
- let method = driver_method!(driver, ConnectionCancel);
- let status = unsafe { method(connection.deref_mut(), &mut error) };
+ // SAFETY: The upgraded Arc prevents release; ADBC requires
cancellation
+ // to be thread-safe, including when an ordinary operation is
running.
+ let status = unsafe { method(connection, &mut error) };
check_status(status, error)
} else {
Ok(())
@@ -863,18 +902,19 @@ impl Optionable for ManagedConnection {
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let method = driver_method!(driver, ConnectionGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
- method(connection.deref_mut(), key, value, length, error)
+ method(connection, key, value, length, error)
};
get_option_bytes(key, populate, driver)
}
@@ -883,16 +923,16 @@ impl Optionable for ManagedConnection {
let key = CString::new(key.as_ref())?;
let mut value: f64 = f64::default();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetOptionDouble);
- let status =
- unsafe { method(connection.deref_mut(), key.as_ptr(), &mut value,
&mut error) };
+ let status = unsafe { method(connection, key.as_ptr(), &mut value,
&mut error) };
check_status(status, error)?;
Ok(value)
}
@@ -901,53 +941,51 @@ impl Optionable for ManagedConnection {
let key = CString::new(key.as_ref())?;
let mut value: i64 = 0;
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetOptionInt);
- let status =
- unsafe { method(connection.deref_mut(), key.as_ptr(), &mut value,
&mut error) };
+ let status = unsafe { method(connection, key.as_ptr(), &mut value,
&mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let method = driver_method!(driver, ConnectionGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
- method(connection.deref_mut(), key, value, length, error)
+ method(connection, key, value, length, error)
};
get_option_string(key, populate, driver)
}
fn set_option(&mut self, key: Self::Option, value: OptionValue) ->
Result<()> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
- set_option_connection(
- driver,
- connection.deref_mut(),
- self.driver_version(),
- key,
- value,
- )
+ let connection = self.inner.connection.get();
+ // SAFETY: The operation guard serializes ordinary calls; the owning
Arc
+ // keeps the handle alive. Only driver-supported cancellation can
overlap.
+ unsafe { set_option_connection_raw(driver, connection,
self.driver_version(), key, value) }
}
}
@@ -956,20 +994,21 @@ impl Connection for ManagedConnection {
fn new_statement(&mut self) -> Result<Self::StatementType> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut statement = adbc_ffi::FFI_AdbcStatement::default();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementNew);
- let status = unsafe { method(connection.deref_mut(), &mut statement,
&mut error) };
+ let status = unsafe { method(connection, &mut statement, &mut error) };
check_status(status, error)?;
let inner = Arc::new(ManagedStatementInner {
- statement: Mutex::new(statement),
+ statement: NativeHandle::new(statement),
connection: self.inner.clone(),
});
@@ -984,29 +1023,31 @@ impl Connection for ManagedConnection {
fn commit(&mut self) -> Result<()> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionCommit);
- let status = unsafe { method(connection.deref_mut(), &mut error) };
+ let status = unsafe { method(connection, &mut error) };
check_status(status, error)
}
fn rollback(&mut self) -> Result<()> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionRollback);
- let status = unsafe { method(connection.deref_mut(), &mut error) };
+ let status = unsafe { method(connection, &mut error) };
check_status(status, error)
}
@@ -1022,23 +1063,16 @@ impl Connection for ManagedConnection {
.map(|c| (c.as_ptr(), c.len()))
.unwrap_or((null(), 0));
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetInfo);
- let status = unsafe {
- method(
- connection.deref_mut(),
- codes_ptr,
- codes_len,
- &mut stream,
- &mut error,
- )
- };
+ let status = unsafe { method(connection, codes_ptr, codes_len, &mut
stream, &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
@@ -1083,19 +1117,20 @@ impl Connection for ManagedConnection {
};
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetObjects);
let mut stream = FFI_ArrowArrayStream::empty();
let status = unsafe {
method(
- connection.deref_mut(),
+ connection,
depth.into(),
catalog_ptr,
db_schema_ptr,
@@ -1136,17 +1171,18 @@ impl Connection for ManagedConnection {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetStatistics);
let status = unsafe {
method(
- connection.deref_mut(),
+ connection,
catalog_ptr,
db_schema_ptr,
table_name_ptr,
@@ -1169,15 +1205,16 @@ impl Connection for ManagedConnection {
}
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetStatisticNames);
- let status = unsafe { method(connection.deref_mut(), &mut stream, &mut
error) };
+ let status = unsafe { method(connection, &mut stream, &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
@@ -1199,17 +1236,18 @@ impl Connection for ManagedConnection {
let mut schema = FFI_ArrowSchema::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetTableSchema);
let status = unsafe {
method(
- connection.deref_mut(),
+ connection,
catalog_ptr,
db_schema_ptr,
table_name_ptr,
@@ -1224,15 +1262,16 @@ impl Connection for ManagedConnection {
fn get_table_types(&self) -> Result<Box<dyn RecordBatchReader + Send +
'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetTableTypes);
- let status = unsafe { method(connection.deref_mut(), &mut stream, &mut
error) };
+ let status = unsafe { method(connection, &mut stream, &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
@@ -1244,18 +1283,19 @@ impl Connection for ManagedConnection {
) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().map_err(|e| {
+ let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionReadPartition);
let partition = partition.as_ref();
let status = unsafe {
method(
- connection.deref_mut(),
+ connection,
partition.as_ptr(),
partition.len(),
&mut stream,
@@ -1269,7 +1309,7 @@ impl Connection for ManagedConnection {
}
struct ManagedStatementInner {
- statement: Mutex<adbc_ffi::FFI_AdbcStatement>,
+ statement: NativeHandle<adbc_ffi::FFI_AdbcStatement>,
connection: Arc<ManagedConnectionInner>,
}
/// Implementation of [Statement].
@@ -1302,15 +1342,17 @@ impl adbc_core::CancelHandle for StatementCancelHandle {
));
}
let driver = &inner.connection.database.driver.driver;
- let mut statement = inner.statement.lock().map_err(|e| {
- Error::with_message_and_status(
- format!("[Driver Manager] statement is poisoned: {e:?}"),
- Status::Internal,
- )
- })?;
+ let Some(method) = driver.StatementCancel else {
+ return Err(Error::with_message_and_status(
+ "[Driver Manager] StatementCancel not supported by driver",
+ Status::NotImplemented,
+ ));
+ };
+ let statement = inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
- let method = driver_method!(driver, StatementCancel);
- let status = unsafe { method(statement.deref_mut(), &mut error) };
+ // SAFETY: The upgraded Arc prevents release; ADBC requires
cancellation
+ // to be thread-safe, including when an ordinary operation is
running.
+ let status = unsafe { method(statement, &mut error) };
check_status(status, error)
} else {
Ok(())
@@ -1321,33 +1363,35 @@ impl adbc_core::CancelHandle for StatementCancelHandle {
impl Statement for ManagedStatement {
fn bind(&mut self, batch: RecordBatch) -> Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementBind);
let batch: StructArray = batch.into();
let (mut array, mut schema) = to_ffi(&batch.to_data())?;
- let status = unsafe { method(statement.deref_mut(), &mut array, &mut
schema, &mut error) };
+ let status = unsafe { method(statement, &mut array, &mut schema, &mut
error) };
check_status(status, error)?;
Ok(())
}
fn bind_stream(&mut self, reader: Box<dyn RecordBatchReader + Send>) ->
Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementBindStream);
let mut stream = FFI_ArrowArrayStream::new(reader);
- let status = unsafe { method(statement.deref_mut(), &mut stream, &mut
error) };
+ let status = unsafe { method(statement, &mut stream, &mut error) };
check_status(status, error)?;
Ok(())
}
@@ -1360,16 +1404,17 @@ impl Statement for ManagedStatement {
fn execute(&mut self) -> Result<Box<dyn RecordBatchReader + Send +
'static>> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteQuery);
let mut stream = FFI_ArrowArrayStream::empty();
- let status = unsafe { method(statement.deref_mut(), &mut stream,
null_mut(), &mut error) };
+ let status = unsafe { method(statement, &mut stream, null_mut(), &mut
error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
@@ -1377,51 +1422,47 @@ impl Statement for ManagedStatement {
fn execute_schema(&mut self) -> Result<arrow_schema::Schema> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteSchema);
let mut schema = FFI_ArrowSchema::empty();
- let status = unsafe { method(statement.deref_mut(), &mut schema, &mut
error) };
+ let status = unsafe { method(statement, &mut schema, &mut error) };
check_status(status, error)?;
Ok((&schema).try_into()?)
}
fn execute_update(&mut self) -> Result<Option<i64>> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteQuery);
let mut rows_affected: i64 = -1;
- let status = unsafe {
- method(
- statement.deref_mut(),
- null_mut(),
- &mut rows_affected,
- &mut error,
- )
- };
+ let status = unsafe { method(statement, null_mut(), &mut
rows_affected, &mut error) };
check_status(status, error)?;
Ok((rows_affected != -1).then_some(rows_affected))
}
fn execute_partitions(&mut self) -> Result<PartitionedResult> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecutePartitions);
let mut schema = FFI_ArrowSchema::empty();
@@ -1429,7 +1470,7 @@ impl Statement for ManagedStatement {
let mut rows_affected: i64 = -1;
let status = unsafe {
method(
- statement.deref_mut(),
+ statement,
&mut schema,
&mut partitions,
&mut rows_affected,
@@ -1449,31 +1490,33 @@ impl Statement for ManagedStatement {
fn get_parameter_schema(&self) -> Result<arrow_schema::Schema> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetParameterSchema);
let mut schema = FFI_ArrowSchema::empty();
- let status = unsafe { method(statement.deref_mut(), &mut schema, &mut
error) };
+ let status = unsafe { method(statement, &mut schema, &mut error) };
check_status(status, error)?;
Ok((&schema).try_into()?)
}
fn prepare(&mut self) -> Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementPrepare);
- let status = unsafe { method(statement.deref_mut(), &mut error) };
+ let status = unsafe { method(statement, &mut error) };
check_status(status, error)?;
Ok(())
}
@@ -1481,32 +1524,33 @@ impl Statement for ManagedStatement {
fn set_sql_query(&mut self, query: impl AsRef<str>) -> Result<()> {
let query = CString::new(query.as_ref())?;
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementSetSqlQuery);
- let status = unsafe { method(statement.deref_mut(), query.as_ptr(),
&mut error) };
+ let status = unsafe { method(statement, query.as_ptr(), &mut error) };
check_status(status, error)?;
Ok(())
}
fn set_substrait_plan(&mut self, plan: impl AsRef<[u8]>) -> Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementSetSubstraitPlan);
let plan = plan.as_ref();
- let status =
- unsafe { method(statement.deref_mut(), plan.as_ptr(), plan.len(),
&mut error) };
+ let status = unsafe { method(statement, plan.as_ptr(), plan.len(),
&mut error) };
check_status(status, error)?;
Ok(())
}
@@ -1517,18 +1561,19 @@ impl Optionable for ManagedStatement {
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let method = driver_method!(driver, StatementGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
- method(statement.deref_mut(), key, value, length, error)
+ method(statement, key, value, length, error)
};
get_option_bytes(key, populate, driver)
}
@@ -1537,15 +1582,16 @@ impl Optionable for ManagedStatement {
let key = CString::new(key.as_ref())?;
let mut value: f64 = f64::default();
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetOptionDouble);
- let status = unsafe { method(statement.deref_mut(), key.as_ptr(), &mut
value, &mut error) };
+ let status = unsafe { method(statement, key.as_ptr(), &mut value, &mut
error) };
check_status(status, error)?;
Ok(value)
}
@@ -1554,63 +1600,63 @@ impl Optionable for ManagedStatement {
let key = CString::new(key.as_ref())?;
let mut value: i64 = 0;
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetOptionInt);
- let status = unsafe { method(statement.deref_mut(), key.as_ptr(), &mut
value, &mut error) };
+ let status = unsafe { method(statement, key.as_ptr(), &mut value, &mut
error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
+ let statement = self.inner.statement.get();
let method = driver_method!(driver, StatementGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
- method(statement.deref_mut(), key, value, length, error)
+ method(statement, key, value, length, error)
};
get_option_string(key, populate, driver)
}
fn set_option(&mut self, key: Self::Option, value: OptionValue) ->
Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().map_err(|e| {
+ let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
- set_option_statement(
- driver,
- statement.deref_mut(),
- self.driver_version(),
- key,
- value,
- )
+ let statement = self.inner.statement.get();
+ // SAFETY: The operation guard serializes ordinary calls; the owning
Arc
+ // keeps the handle alive. Only driver-supported cancellation can
overlap.
+ unsafe { set_option_statement_raw(driver, statement,
self.driver_version(), key, value) }
}
}
-impl Drop for ManagedStatement {
+impl Drop for ManagedStatementInner {
fn drop(&mut self) {
- let driver = self.ffi_driver();
- if let Ok(mut statement) = self.inner.statement.lock() {
+ let driver = &self.connection.database.driver.driver;
+ if !self.statement.mu.is_poisoned() {
+ let statement = self.statement.get();
let method = driver_method!(driver, StatementRelease);
// TODO(alexandreyc): how should we handle `StatementRelease`
failing?
// See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
- unsafe { method(statement.deref_mut(), null_mut()) };
+ unsafe { method(statement, null_mut()) };
}
// We could still drop here but if the lock is poisoned, we have no
// clue what the status is. Since a panic comes from Rust code,
diff --git a/rust/driver_manager/tests/cancellation.rs
b/rust/driver_manager/tests/cancellation.rs
new file mode 100644
index 000000000..f9d260bf4
--- /dev/null
+++ b/rust/driver_manager/tests/cancellation.rs
@@ -0,0 +1,495 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Tests of cancellation through the driver manager, focusing on making sure
+// cancellation is truly concurrent.
+
+use std::cell::RefCell;
+use std::ffi::{c_int, c_void};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering::SeqCst};
+use std::thread;
+use std::time::{Duration, Instant};
+
+use adbc_core::error::{AdbcStatusCode, Result, Status};
+use adbc_core::options::AdbcVersion;
+use adbc_core::{CancelHandle, Connection, Database, Driver, Statement};
+use adbc_driver_manager::{ManagedConnection, ManagedDriver, ManagedStatement};
+use adbc_ffi::{
+ FFI_AdbcConnection, FFI_AdbcDatabase, FFI_AdbcDriver,
FFI_AdbcDriverInitFunc, FFI_AdbcError,
+ FFI_AdbcStatement,
+};
+use arrow_array::ffi_stream::FFI_ArrowArrayStream;
+
+const TIMEOUT: Duration = Duration::from_secs(2);
+const CALLBACK_TIMEOUT: Duration = Duration::from_secs(10);
+
+// ------------------------------------------------------------
+// Simulated driver implementation
+// ------------------------------------------------------------
+
+struct AtomicFlag {
+ mu: std::sync::Mutex<bool>,
+ cv: std::sync::Condvar,
+}
+
+impl AtomicFlag {
+ fn new(value: bool) -> Self {
+ Self {
+ mu: std::sync::Mutex::new(value),
+ cv: std::sync::Condvar::new(),
+ }
+ }
+
+ fn clear(&self) {
+ let mut guard = self.mu.lock().unwrap();
+ *guard = false;
+ self.cv.notify_all();
+ }
+
+ fn set(&self) {
+ let mut guard = self.mu.lock().unwrap();
+ *guard = true;
+ self.cv.notify_all();
+ }
+
+ fn wait(&self, timeout: Duration) -> bool {
+ let (guard, _) = self
+ .cv
+ .wait_timeout_while(self.mu.lock().unwrap(), timeout, |flag|
!*flag)
+ .unwrap();
+ *guard
+ }
+}
+
+// Track what's been done with the driver so we can verify invariants later.
+struct DriverObjectState {
+ // count of non-cancel operations
+ operations: AtomicUsize,
+ // count of cancel operations
+ cancellations: AtomicUsize,
+ // allow the operation to proceed
+ release_operation: AtomicFlag,
+ // allow the cancel to proceed
+ release_cancel: AtomicFlag,
+ // status to return from the cancel callback
+ cancel_status: AtomicU8,
+ // whether the operation or cancel timed out waiting for the other to
release
+ timed_out: AtomicBool,
+ statement_releases: AtomicUsize,
+ connection_releases: AtomicUsize,
+ database_releases: AtomicUsize,
+ driver_releases: AtomicUsize,
+}
+
+impl DriverObjectState {
+ fn new() -> Self {
+ Self {
+ operations: AtomicUsize::new(0),
+ cancellations: AtomicUsize::new(0),
+ release_operation: AtomicFlag::new(false),
+ release_cancel: AtomicFlag::new(true),
+ cancel_status: AtomicU8::new(Status::Ok.into()),
+ timed_out: AtomicBool::new(false),
+ statement_releases: AtomicUsize::new(0),
+ connection_releases: AtomicUsize::new(0),
+ database_releases: AtomicUsize::new(0),
+ driver_releases: AtomicUsize::new(0),
+ }
+ }
+
+ // Do "something" with the driver, waiting on release_operation
+ fn operate(&self) -> AdbcStatusCode {
+ self.operations.fetch_add(1, SeqCst);
+ if !self.release_operation.wait(CALLBACK_TIMEOUT) {
+ self.timed_out.store(true, SeqCst);
+ return Status::Internal.into();
+ }
+ Status::Cancelled.into()
+ }
+
+ // Cancel operate() by flagging release_operation, waiting on
release_cancel first
+ fn cancel(&self) -> AdbcStatusCode {
+ self.cancellations.fetch_add(1, SeqCst);
+ if !self.release_cancel.wait(CALLBACK_TIMEOUT) {
+ self.timed_out.store(true, SeqCst);
+ return Status::Internal.into();
+ }
+ self.release_operation.set();
+ self.cancel_status.load(SeqCst)
+ }
+
+ fn releases(&self) -> [usize; 4] {
+ [
+ self.statement_releases.load(SeqCst),
+ self.connection_releases.load(SeqCst),
+ self.database_releases.load(SeqCst),
+ self.driver_releases.load(SeqCst),
+ ]
+ }
+
+ // SAFETY: Each private_data is a Box<Arc<State>> initialized below.
Callbacks
+ // clone the Arc before blocking and never borrow the mutable FFI handle
itself.
+ unsafe fn from_raw(data: *mut c_void) -> Arc<Self> {
+ unsafe { Arc::clone(&*data.cast::<Arc<Self>>()) }
+ }
+}
+
+fn wait_for(condition: impl Fn() -> bool, timeout: Duration) -> bool {
+ let deadline = Instant::now() + timeout;
+ while !condition() && Instant::now() < deadline {
+ thread::sleep(Duration::from_millis(50));
+ }
+ condition()
+}
+
+thread_local! {
+ // We need a way to pass state to the driver impl, so use a thread local.
+ static INITIAL_STATE: RefCell<Option<Arc<DriverObjectState>>> = const {
RefCell::new(None) };
+}
+
+// ------------------------------------------------------------
+// Simulated driver implementation (driver callbacks)
+// ------------------------------------------------------------
+
+macro_rules! release_callback {
+ ($name:ident, $ty:ty, $counter:ident) => {
+ unsafe extern "C" fn $name(handle: *mut $ty, _: *mut FFI_AdbcError) ->
AdbcStatusCode {
+ unsafe {
+ if (*handle).private_data.is_null() {
+ return Status::InvalidState.into();
+ }
+ let state =
Box::from_raw((*handle).private_data.cast::<Arc<DriverObjectState>>());
+ state.$counter.fetch_add(1, SeqCst);
+ (*handle).private_data = std::ptr::null_mut();
+ }
+ Status::Ok.into()
+ }
+ };
+}
+
+release_callback!(driver_release, FFI_AdbcDriver, driver_releases);
+release_callback!(database_release, FFI_AdbcDatabase, database_releases);
+release_callback!(connection_release, FFI_AdbcConnection, connection_releases);
+release_callback!(statement_release, FFI_AdbcStatement, statement_releases);
+
+unsafe extern "C" fn database_new(
+ database: *mut FFI_AdbcDatabase,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe {
+ let state =
DriverObjectState::from_raw((*(*database).private_driver).private_data);
+ (*database).private_data = Box::into_raw(Box::new(state)).cast();
+ }
+ Status::Ok.into()
+}
+
+unsafe extern "C" fn database_init(
+ _: *mut FFI_AdbcDatabase,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ Status::Ok.into()
+}
+
+unsafe extern "C" fn connection_new(
+ _: *mut FFI_AdbcConnection,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ Status::Ok.into()
+}
+
+unsafe extern "C" fn connection_init(
+ connection: *mut FFI_AdbcConnection,
+ database: *mut FFI_AdbcDatabase,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe {
+ (*connection).private_data =
Box::into_raw(Box::new(DriverObjectState::from_raw(
+ (*database).private_data,
+ )))
+ .cast();
+ }
+ Status::Ok.into()
+}
+
+unsafe extern "C" fn statement_new(
+ connection: *mut FFI_AdbcConnection,
+ statement: *mut FFI_AdbcStatement,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe {
+ (*statement).private_data =
Box::into_raw(Box::new(DriverObjectState::from_raw(
+ (*connection).private_data,
+ )))
+ .cast();
+ }
+ Status::Ok.into()
+}
+
+unsafe extern "C" fn commit(
+ connection: *mut FFI_AdbcConnection,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe { DriverObjectState::from_raw((*connection).private_data).operate()
}
+}
+
+unsafe extern "C" fn execute(
+ statement: *mut FFI_AdbcStatement,
+ _: *mut FFI_ArrowArrayStream,
+ _: *mut i64,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe { DriverObjectState::from_raw((*statement).private_data).operate() }
+}
+
+unsafe extern "C" fn connection_cancel(
+ handle: *mut FFI_AdbcConnection,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe { DriverObjectState::from_raw((*handle).private_data).cancel() }
+}
+
+unsafe extern "C" fn statement_cancel(
+ handle: *mut FFI_AdbcStatement,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ unsafe { DriverObjectState::from_raw((*handle).private_data).cancel() }
+}
+
+unsafe extern "C" fn init<const CANCEL: bool>(
+ _: c_int,
+ driver: *mut c_void,
+ _: *mut FFI_AdbcError,
+) -> AdbcStatusCode {
+ let Some(state) = INITIAL_STATE.with(|state| state.borrow_mut().take())
else {
+ return Status::Internal.into();
+ };
+ unsafe {
+ *driver.cast::<FFI_AdbcDriver>() = FFI_AdbcDriver {
+ private_data: Box::into_raw(Box::new(state)).cast(),
+ release: Some(driver_release),
+ DatabaseNew: Some(database_new),
+ DatabaseInit: Some(database_init),
+ DatabaseRelease: Some(database_release),
+ ConnectionNew: Some(connection_new),
+ ConnectionInit: Some(connection_init),
+ ConnectionRelease: Some(connection_release),
+ ConnectionCommit: Some(commit),
+ ConnectionCancel: if CANCEL {
+ Some(connection_cancel)
+ } else {
+ None
+ },
+ StatementNew: Some(statement_new),
+ StatementRelease: Some(statement_release),
+ StatementExecuteQuery: Some(execute),
+ StatementCancel: if CANCEL { Some(statement_cancel) } else { None
},
+ ..FFI_AdbcDriver::default()
+ };
+ }
+ Status::Ok.into()
+}
+
+// ------------------------------------------------------------
+// Test cases
+// ------------------------------------------------------------
+
+#[derive(Clone, Copy, Debug)]
+enum ObjectKind {
+ Connection,
+ Statement,
+}
+
+const OBJECT_KINDS: [ObjectKind; 2] = [ObjectKind::Connection,
ObjectKind::Statement];
+
+#[derive(Clone)]
+enum Object {
+ Connection(ManagedConnection),
+ Statement(ManagedStatement),
+}
+
+impl Object {
+ fn cancel_handle(&self) -> Box<dyn CancelHandle> {
+ match self {
+ Self::Connection(connection) => connection.get_cancel_handle(),
+ Self::Statement(statement) => statement.get_cancel_handle(),
+ }
+ }
+
+ fn operate(&mut self) -> Result<()> {
+ match self {
+ Self::Connection(connection) => connection.commit(),
+ Self::Statement(statement) => statement.execute().map(|_| ()),
+ }
+ }
+}
+
+fn setup<const CANCEL: bool>(
+ kind: ObjectKind,
+ version: AdbcVersion,
+) -> Result<(Arc<DriverObjectState>, Object)> {
+ let state = Arc::new(DriverObjectState::new());
+ INITIAL_STATE.with(|slot| *slot.borrow_mut() = Some(state.clone()));
+ let mut driver =
+ ManagedDriver::load_static(&(init::<CANCEL> as
FFI_AdbcDriverInitFunc), version)?;
+ let database = driver.new_database()?;
+ let mut connection = database.new_connection()?;
+ let object = match kind {
+ ObjectKind::Connection => Object::Connection(connection),
+ ObjectKind::Statement =>
Object::Statement(connection.new_statement()?),
+ };
+ Ok((state, object))
+}
+
+type TestResult = std::result::Result<(), Box<dyn std::error::Error>>;
+
+#[test]
+fn cancel_is_concurrent() -> TestResult {
+ for kind in OBJECT_KINDS {
+ let (state, mut object) = setup::<true>(kind, AdbcVersion::V110)?;
+ let cancel = object.cancel_handle();
+ let execution = thread::spawn(move || (object.operate(), object));
+ let entered = wait_for(|| state.operations.load(SeqCst) == 1, TIMEOUT);
+ let cancellation = thread::spawn(move || cancel.try_cancel());
+ let overlapped = wait_for(|| state.cancellations.load(SeqCst) == 1,
TIMEOUT);
+ state.release_operation.set(); // in case cancel didn't work, unblock
+ let (result, object) = execution.join().map_err(|_| "worker thread
panicked")?;
+ cancellation
+ .join()
+ .map_err(|_| "worker thread panicked")??;
+ assert!(entered, "{kind:?} did not enter the driver");
+ assert!(overlapped, "{kind:?} blocked cancellation");
+ assert_eq!(
+ result.err().map(|error| error.status),
+ Some(Status::Cancelled)
+ );
+ assert!(!state.timed_out.load(SeqCst));
+ drop(object);
+ }
+ Ok(())
+}
+
+#[test]
+fn cancellation_keeps_native_handles_alive() -> TestResult {
+ for kind in OBJECT_KINDS {
+ let (state, object) = setup::<true>(kind, AdbcVersion::V110)?;
+ let clone = object.clone();
+ drop(clone);
+ assert_eq!(
+ state.releases(),
+ [0; 4],
+ "dropping a clone released {kind:?}"
+ );
+ state.release_cancel.clear();
+ let cancel = object.cancel_handle();
+ let expired = object.cancel_handle();
+ let cancellation = thread::spawn(move || cancel.try_cancel());
+ let entered = wait_for(|| state.cancellations.load(SeqCst) == 1,
TIMEOUT);
+ drop(object);
+ let releases_during_cancel = state.releases();
+ state.release_cancel.set();
+ cancellation
+ .join()
+ .map_err(|_| "worker thread panicked")??;
+ assert!(entered);
+ assert_eq!(releases_during_cancel, [0; 4]);
+ let expected = match kind {
+ ObjectKind::Connection => [0, 1, 1, 1],
+ ObjectKind::Statement => [1, 1, 1, 1],
+ };
+ assert_eq!(state.releases(), expected);
+ assert!(!state.timed_out.load(SeqCst));
+ expired.try_cancel()?;
+ assert_eq!(state.cancellations.load(SeqCst), 1);
+ }
+ Ok(())
+}
+
+#[test]
+fn concurrent_cancellation_and_callback_errors() -> TestResult {
+ for kind in OBJECT_KINDS {
+ let (state, object) = setup::<true>(kind, AdbcVersion::V110)?;
+ state.release_cancel.clear();
+ state
+ .cancel_status
+ .store(Status::InvalidState.into(), SeqCst);
+ let cancel: Arc<dyn CancelHandle> = object.cancel_handle().into();
+ let second_cancel = cancel.clone();
+ let first = thread::spawn(move || cancel.try_cancel());
+ let second = thread::spawn(move || second_cancel.try_cancel());
+ let overlapped = wait_for(|| state.cancellations.load(SeqCst) == 2,
TIMEOUT);
+ state.release_cancel.set();
+ let results = [
+ first.join().map_err(|_| "worker thread panicked")?,
+ second.join().map_err(|_| "worker thread panicked")?,
+ ];
+ assert!(overlapped, "{kind:?} cancellation calls did not overlap");
+ for result in results {
+ assert_eq!(
+ result.err().map(|error| error.status),
+ Some(Status::InvalidState)
+ );
+ }
+ assert!(!state.timed_out.load(SeqCst));
+ }
+ Ok(())
+}
+
+fn unsupported<const CANCEL: bool>(kind: ObjectKind, version: AdbcVersion) ->
TestResult {
+ let (state, mut object) = setup::<CANCEL>(kind, version)?;
+ let cancel = object.cancel_handle();
+ let execution = thread::spawn(move || (object.operate(), object));
+ let entered = wait_for(|| state.operations.load(SeqCst) == 1, TIMEOUT);
+ let finished = Arc::new(AtomicBool::new(false));
+ let cancel_finished = finished.clone();
+ let cancellation = thread::spawn(move || {
+ let result = cancel.try_cancel();
+ cancel_finished.store(true, SeqCst);
+ result
+ });
+ let returned = wait_for(|| finished.load(SeqCst), TIMEOUT);
+ state.release_operation.set();
+ let (result, object) = execution.join().map_err(|_| "worker thread
panicked")?;
+ let cancel_result = cancellation.join().map_err(|_| "worker thread
panicked")?;
+ assert!(
+ entered && returned,
+ "unsupported {kind:?} cancellation blocked"
+ );
+ assert_eq!(
+ cancel_result.err().map(|error| error.status),
+ Some(Status::NotImplemented)
+ );
+ assert_eq!(
+ result.err().map(|error| error.status),
+ Some(Status::Cancelled)
+ );
+ assert_eq!(state.cancellations.load(SeqCst), 0);
+ assert!(!state.timed_out.load(SeqCst));
+ let expired = object.cancel_handle();
+ drop(object);
+ expired.try_cancel()?;
+ Ok(())
+}
+
+#[test]
+fn unsupported_cancellation_does_not_wait() -> TestResult {
+ for kind in OBJECT_KINDS {
+ unsupported::<false>(kind, AdbcVersion::V110)?;
+ unsupported::<true>(kind, AdbcVersion::V100)?;
+ }
+ Ok(())
+}
diff --git a/rust/ffi/src/options.rs b/rust/ffi/src/options.rs
index 8266af252..41adcfe83 100644
--- a/rust/ffi/src/options.rs
+++ b/rust/ffi/src/options.rs
@@ -167,6 +167,25 @@ pub fn set_option_connection(
version: AdbcVersion,
key: impl AsRef<str>,
value: OptionValue,
+) -> Result<()> {
+ // SAFETY: The exclusive borrow keeps the handle alive and prevents
concurrent access.
+ unsafe { set_option_connection_raw(driver, connection, version, key,
value) }
+}
+
+/// Set an option through a raw connection handle.
+///
+/// # Safety
+///
+/// The handle must be valid for the driver and remain alive throughout this
call.
+/// Ordinary operations on the handle must be serialized. Concurrent
cancellation
+/// is permitted only as supported by the driver's ADBC cancellation contract.
+/// No references to the handle may conflict with accesses made by the driver.
+pub unsafe fn set_option_connection_raw(
+ driver: &FFI_AdbcDriver,
+ connection: *mut FFI_AdbcConnection,
+ version: AdbcVersion,
+ key: impl AsRef<str>,
+ value: OptionValue,
) -> Result<()> {
let key = CString::new(key.as_ref())?;
let mut error = FFI_AdbcError::with_driver(driver);
@@ -213,6 +232,25 @@ pub fn set_option_statement(
version: AdbcVersion,
key: impl AsRef<str>,
value: OptionValue,
+) -> Result<()> {
+ // SAFETY: The exclusive borrow keeps the handle alive and prevents
concurrent access.
+ unsafe { set_option_statement_raw(driver, statement, version, key, value) }
+}
+
+/// Set an option through a raw statement handle.
+///
+/// # Safety
+///
+/// The handle must be valid for the driver and remain alive throughout this
call.
+/// Ordinary operations on the handle must be serialized. Concurrent
cancellation
+/// is permitted only as supported by the driver's ADBC cancellation contract.
+/// No references to the handle may conflict with accesses made by the driver.
+pub unsafe fn set_option_statement_raw(
+ driver: &FFI_AdbcDriver,
+ statement: *mut FFI_AdbcStatement,
+ version: AdbcVersion,
+ key: impl AsRef<str>,
+ value: OptionValue,
) -> Result<()> {
let key = CString::new(key.as_ref())?;
let mut error = FFI_AdbcError::with_driver(driver);
