From 2f0f77b72edb43193ad9455fbb7b25b39b121c13 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 12 Aug 2026 13:22:59 -0400 Subject: [PATCH 1/2] feat(sf_core): atomically take result streams --- .../src/apis/database_driver_v1/result_set.rs | 123 +++++++++++++++--- .../src/apis/database_driver_v1/statement.rs | 3 +- sf_core/src/handle_manager.rs | 32 +++-- 3 files changed, 121 insertions(+), 37 deletions(-) diff --git a/sf_core/src/apis/database_driver_v1/result_set.rs b/sf_core/src/apis/database_driver_v1/result_set.rs index e81b3ea412..9bad10dbd5 100644 --- a/sf_core/src/apis/database_driver_v1/result_set.rs +++ b/sf_core/src/apis/database_driver_v1/result_set.rs @@ -390,26 +390,11 @@ async fn snapshot_reader_inputs( // --- DatabaseDriverV1 impl --- impl DatabaseDriverV1 { - /// Builds a fresh Arrow [`RecordBatchReader`] for this result set, lazily - /// from the stored `RowsetData`, so it can be requested multiple times. The - /// protobuf layer wraps it in an `FFI_ArrowArrayStream` at the C boundary. - /// - /// Awaiting this only builds the reader and never blocks. Iterating it, - /// however, must happen in a synchronous context: chunked result sets pull - /// chunks via a blocking channel receiver, so draining from within an async - /// runtime would call `blocking_recv` and panic. Drain after returning from - /// `block_on` (keeping the runtime alive), on a dedicated `std::thread`, or - /// via `tokio::task::spawn_blocking`. - pub async fn result_set_get_stream( + // Build a reader from a retained or atomically consumed result set. + async fn build_result_set_stream( &self, - result_handle: Handle, + rs_ptr: Arc>, ) -> Result, ApiError> { - let rs_ptr = self - .results - .get_obj(result_handle) - .with_context(|| InvalidArgumentSnafu { - argument: "result_handle: ResultSet handle not found".to_string(), - })?; let (data, http_client, prefetch_config, columns) = snapshot_reader_inputs(&rs_ptr).await; let nullable_flags: Vec = columns.iter().map(|c| c.nullable).collect(); @@ -418,7 +403,7 @@ impl DatabaseDriverV1 { } else { Some(nullable_flags.as_slice()) }; - let reader = build_reader_from_rowset_data( + build_reader_from_rowset_data( &data, http_client, &prefetch_config, @@ -426,8 +411,45 @@ impl DatabaseDriverV1 { flags, ) .await - .context(QueryResponseProcessSnafu)?; - Ok(reader) + .context(QueryResponseProcessSnafu) + } + + /// Builds a fresh Arrow [`RecordBatchReader`] while leaving the result-set + /// handle registered, so callers may request the stream again. + /// + /// Awaiting this only builds the reader. Drain chunked readers outside an + /// async runtime, on a dedicated thread, or with `spawn_blocking`. + pub async fn result_set_get_stream( + &self, + result_handle: Handle, + ) -> Result, ApiError> { + let rs_ptr = self + .results + .get_obj(result_handle) + .with_context(|| InvalidArgumentSnafu { + argument: "result_handle: ResultSet handle not found".to_string(), + })?; + self.build_result_set_stream(rs_ptr).await + } + + /// Atomically consumes a result-set handle and builds its Arrow stream. + /// + /// The handle is deregistered before stream construction, and remains + /// released whether construction succeeds or fails. Use this for one-shot + /// consumers to avoid pairing [`Self::result_set_get_stream`] with + /// [`Self::result_set_release`] and duplicating cleanup logic. As with + /// `result_set_get_stream`, drain chunked readers outside an async runtime. + pub async fn result_set_take_stream( + &self, + result_handle: Handle, + ) -> Result, ApiError> { + let rs_ptr = + self.results + .take_obj(result_handle) + .with_context(|| InvalidArgumentSnafu { + argument: "result_handle: ResultSet handle not found".to_string(), + })?; + self.build_result_set_stream(rs_ptr).await } /// Returns chunk metadata (inline data + remote chunk URLs) for this result set. @@ -674,6 +696,65 @@ mod tests { drop(runtime); } + #[test] + fn result_set_take_stream_returns_reader_and_consumes_handle() { + let driver = DatabaseDriverV1::new(); + let data: Data = serde_json::from_str(JSON_ROWSET) + .expect("fixture must deserialize into query_response::Data"); + let descriptor = response_to_descriptor(&data, &WrapperPresets::default()); + let reader_ctx = ReaderContext { + http_client: reqwest::Client::new(), + prefetch_config: PrefetchConfig::default(), + }; + let handle = driver.create_result_set(descriptor, data.into_rowset_data(), reader_ctx); + + let runtime = tokio::runtime::Runtime::new().expect("failed to build tokio runtime"); + let reader = runtime + .block_on(driver.result_set_take_stream(handle)) + .expect("taking an inline JSON result stream should succeed"); + + assert!( + runtime + .block_on(driver.result_set_take_stream(handle)) + .is_err(), + "a consumed handle must not be reusable" + ); + assert!( + driver.result_set_release(handle).is_err(), + "a consumed handle must already be released" + ); + assert_id_name_reader(reader); + drop(runtime); + } + + #[test] + fn result_set_take_stream_consumes_handle_when_stream_build_fails() { + let driver = DatabaseDriverV1::new(); + let data: Data = serde_json::from_str(ARROW_DATA) + .expect("fixture must deserialize into query_response::Data"); + let descriptor = response_to_descriptor(&data, &WrapperPresets::default()); + let reader_ctx = ReaderContext { + http_client: reqwest::Client::new(), + prefetch_config: PrefetchConfig::default(), + }; + let malformed = RowsetData::ArrowSingleChunk { + chunk_base64: "not valid base64".to_string(), + }; + let handle = driver.create_result_set(descriptor, malformed, reader_ctx); + + let runtime = tokio::runtime::Runtime::new().expect("failed to build tokio runtime"); + assert!( + runtime + .block_on(driver.result_set_take_stream(handle)) + .is_err(), + "malformed Arrow data must fail stream construction" + ); + assert!( + driver.result_set_release(handle).is_err(), + "the handle must remain released after construction fails" + ); + } + #[test] fn result_set_get_stream_reads_arrow_rowset() { let driver = DatabaseDriverV1::new(); diff --git a/sf_core/src/apis/database_driver_v1/statement.rs b/sf_core/src/apis/database_driver_v1/statement.rs index 1cdca43f08..ad94981274 100644 --- a/sf_core/src/apis/database_driver_v1/statement.rs +++ b/sf_core/src/apis/database_driver_v1/statement.rs @@ -227,8 +227,7 @@ impl DatabaseDriverV1 { } .fail(); }; - let stream = self.result_set_get_stream(rs_info.handle).await?; - self.result_set_release(rs_info.handle)?; + let stream = self.result_set_take_stream(rs_info.handle).await?; let stmt_ptr = self.statements diff --git a/sf_core/src/handle_manager.rs b/sf_core/src/handle_manager.rs index cac0eb2ee3..769bc06ff3 100644 --- a/sf_core/src/handle_manager.rs +++ b/sf_core/src/handle_manager.rs @@ -80,37 +80,41 @@ impl HandleManager { } } - pub fn delete_handle(&self, handle: Handle) -> bool { - let span = span!(target: "handle_manager", Level::INFO, "Deleting handle", handle_id = handle.id, handle_magic = handle.magic); + /// Atomically deregisters and returns the value for `handle`. Only one caller + /// can take a live handle; later lookups and takes fail. + pub(crate) fn take_obj(&self, handle: Handle) -> Option> { + let span = span!(target: "handle_manager", Level::INFO, "Taking handle", handle_id = handle.id, handle_magic = handle.magic); let _enter = span.enter(); let index = handle.id as usize; let mut handles = self.handles.write_recover(); if index >= handles.len() { - tracing::error!("Handle index out of bounds, cannot delete handle"); - return false; + tracing::error!("Handle index out of bounds, cannot take object"); + return None; } let handle_value = &mut handles[index]; - let magic = handle_value.magic; - - if magic != handle.magic { - tracing::error!("Handle magic mismatch, cannot delete handle"); - return false; + if handle_value.magic != handle.magic { + tracing::error!("Handle magic mismatch, cannot take object"); + return None; } match handle_value.value.take() { - Some(_) => { - tracing::trace!(target: "handle_manager", "Handle deleted successfully"); - true + Some(value) => { + tracing::trace!(target: "handle_manager", "Handle taken successfully"); + Some(value) } None => { - tracing::error!("Handle not found, cannot delete handle"); - false + tracing::error!("Handle not found, cannot take object"); + None } } } + pub fn delete_handle(&self, handle: Handle) -> bool { + self.take_obj(handle).is_some() + } + /// Deregisters and returns every currently-live value matching `pred`, /// leaving non-matching handles untouched. Used by session reaping (e.g. /// `stream_transfer::reap_connection_streams`) to sweep up a connection's From e4f4597c4ad0f05e62a5d24a348e1f030fda4a43 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Wed, 12 Aug 2026 13:50:28 -0400 Subject: [PATCH 2/2] Use consistent span name for handle take --- sf_core/src/handle_manager.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sf_core/src/handle_manager.rs b/sf_core/src/handle_manager.rs index 769bc06ff3..d25387c860 100644 --- a/sf_core/src/handle_manager.rs +++ b/sf_core/src/handle_manager.rs @@ -83,7 +83,7 @@ impl HandleManager { /// Atomically deregisters and returns the value for `handle`. Only one caller /// can take a live handle; later lookups and takes fail. pub(crate) fn take_obj(&self, handle: Handle) -> Option> { - let span = span!(target: "handle_manager", Level::INFO, "Taking handle", handle_id = handle.id, handle_magic = handle.magic); + let span = span!(target: "handle_manager", Level::INFO, "HandleManager::take_obj", handle_id = handle.id, handle_magic = handle.magic); let _enter = span.enter(); let index = handle.id as usize; let mut handles = self.handles.write_recover();