diff --git a/mssql-tds/src/connection/tds_client.rs b/mssql-tds/src/connection/tds_client.rs index 67d4d16f..e11620eb 100644 --- a/mssql-tds/src/connection/tds_client.rs +++ b/mssql-tds/src/connection/tds_client.rs @@ -4761,6 +4761,7 @@ mod tests { }; use crate::test_client_support::byte_stream::tds_client_over_raw_bytes as client_over_bytes; use crate::test_client_support::byte_stream::tds_client_over_raw_bytes_with_column_encryption as client_over_bytes_with_ae; + use crate::test_packet_support::{TestPacketBuilder, create_network_transport_with_gated_data}; use crate::token::tokens::{ ColMetadataToken, CurrentCommand, DoneStatus, DoneToken, InfoToken, Tokens, }; @@ -5122,6 +5123,85 @@ mod tests { } } + #[test] + fn sync_first_cursor_continues_across_packet_boundary() { + use crate::datatypes::sqldatatypes::{TdsDataType, TypeInfo}; + use crate::message::messages::PacketType; + use crate::query::metadata::ColumnMetadata; + use crate::runtime::block_on_sync_first; + use crate::token::tokens::TokenType; + + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("test runtime"); + let value = 0x1234_5678_i32; + let bytes = value.to_le_bytes(); + let first = TestPacketBuilder::new(PacketType::TabularResult) + .end_of_message(false) + .append_byte(TokenType::Row as u8) + .append_bytes(&bytes[..2]) + .build(); + let second = TestPacketBuilder::new(PacketType::TabularResult) + .append_bytes(&bytes[2..]) + .build(); + let (transport, release_second_packet) = { + let _guard = runtime.enter(); + create_network_transport_with_gated_data(&first, &second) + }; + let negotiated_settings = + crate::handler::handler_factory::create_test_negotiated_settings_internal(); + let execution_context = crate::connection::execution_context::ExecutionContext::new(); + let client_context = ClientContext::with_data_source("tcp:localhost,1433"); + let mut client = TdsClient::new( + Box::new(transport), + negotiated_settings, + execution_context, + client_context, + ); + client.current_metadata = Some(Arc::new(ColMetadataToken { + column_count: 1, + columns: vec![ColumnMetadata { + user_type: 0, + flags: 0, + type_info: TypeInfo::fixed_len(TdsDataType::Int4).expect("Int4 is fixed length"), + data_type: TdsDataType::Int4, + column_name: "value".to_string(), + multi_part_name: None, + crypto_metadata: None, + }], + cek_table: Vec::new(), + })); + + let positioned = block_on_sync_first(&runtime, client.next_row_cursor()) + .expect("outside runtime") + .expect("position row"); + assert!(positioned); + let mut column_future = std::pin::pin!(client.read_row_column(0)); + let mut release_second_packet = Some(release_second_packet); + let mut polls = 0; + let counted_column = std::future::poll_fn(|context| { + polls += 1; + let result = column_future.as_mut().poll(context); + if result.is_pending() + && let Some(release) = release_second_packet.take() + { + release.send(()).expect("release second packet"); + } + result + }); + let column = block_on_sync_first(&runtime, counted_column) + .expect("outside runtime") + .expect("decode split column"); + + assert!( + polls >= 2, + "the gated second packet must force the same column future to resume" + ); + assert_eq!(column, CursorColumn::Value(ColumnValues::Int(value))); + } + #[test] fn prepare_reset_connection_routes_mode_to_transport() { let mut client = create_test_client(); diff --git a/mssql-tds/src/lib.rs b/mssql-tds/src/lib.rs index 0d62e386..a57140e2 100644 --- a/mssql-tds/src/lib.rs +++ b/mssql-tds/src/lib.rs @@ -73,6 +73,7 @@ //! - [`error`] — Error definitions. //! - [`message`] — TDS message types (prelogin, login7, etc.). //! - [`query`] — Query metadata and column descriptors. +//! - [`runtime`] — Runtime interoperability for synchronous callers. //! - [`token`] — TDS token stream parsing (COLMETADATA, ROW, DONE, etc.). pub mod connection; @@ -88,6 +89,7 @@ pub(crate) mod handler; pub(crate) mod io; pub mod message; pub mod query; +pub mod runtime; pub mod security; pub(crate) mod sql_identifier; pub(crate) mod ssrp; diff --git a/mssql-tds/src/runtime.rs b/mssql-tds/src/runtime.rs new file mode 100644 index 00000000..31b8d56d --- /dev/null +++ b/mssql-tds/src/runtime.rs @@ -0,0 +1,254 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Runtime interoperability for synchronous consumers. + +use std::future::Future; +use std::task::{Context, Poll, Waker}; + +use tokio::runtime::{Handle, Runtime}; + +/// Failure to start synchronous execution. +#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] +pub enum SyncFirstError { + /// The calling thread already has an entered Tokio runtime context. + #[error("sync-first execution requires a thread without an entered Tokio runtime context")] + RuntimeContext, +} + +/// Runs a future synchronously, entering Tokio only when the first poll suspends. +/// +/// The future is pinned once and polled under `runtime` with the standard +/// library's no-op waker. A ready future returns directly. A pending future is +/// passed, still pinned and without reconstruction, to [`Runtime::block_on`], +/// whose first poll replaces the temporary waker with the runtime's task waker. +/// +/// This is useful at synchronous FFI boundaries where protocol work commonly +/// completes from bytes already buffered by the transport. It does not change +/// the future's cancellation, timeout, or I/O behavior after suspension. +/// +/// # Errors +/// +/// Returns [`SyncFirstError::RuntimeContext`] before polling when +/// [`Handle::try_current`] finds an entered Tokio runtime context. This is +/// deliberately stricter than [`Runtime::block_on`]: Tokio blocking-pool threads +/// also carry an entered handle, even though they may call `block_on`. Callers +/// already using Tokio should stay async instead; synchronous FFI callers should +/// invoke this function from their native, unentered thread. +/// +/// Rejecting every entered context up front guarantees a stateful protocol +/// future is never partially advanced before the fallback execution policy is +/// known to be valid. +#[inline] +pub fn block_on_sync_first(runtime: &Runtime, future: F) -> Result +where + F: Future, +{ + if Handle::try_current().is_ok() { + return Err(SyncFirstError::RuntimeContext); + } + + let mut future = std::pin::pin!(future); + let first_poll = { + let _guard = runtime.enter(); + let mut context = Context::from_waker(Waker::noop()); + future.as_mut().poll(&mut context) + }; + + match first_poll { + Poll::Ready(output) => Ok(output), + Poll::Pending => Ok(runtime.block_on(future)), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{CancelHandle, TdsResult}; + use crate::error::Error; + use std::pin::Pin; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::Waker; + use std::time::Duration; + + struct PendingThenReady { + address: Option, + first_waker: Option, + polls: Arc, + } + + impl Future for PendingThenReady { + type Output = usize; + + fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll { + let address = (&*self as *const Self) as usize; + let poll = self.polls.fetch_add(1, Ordering::Relaxed) + 1; + match self.address { + None => { + self.address = Some(address); + self.first_waker = Some(context.waker().clone()); + Poll::Pending + } + Some(first_address) => { + assert_eq!( + address, first_address, + "the pending future must not be reconstructed or moved" + ); + assert!( + !self + .first_waker + .as_ref() + .expect("first poll stores its waker") + .will_wake(context.waker()), + "Runtime::block_on must replace the temporary no-op waker" + ); + Poll::Ready(poll) + } + } + } + } + + struct CountPolls<'a>(&'a AtomicUsize); + + impl Future for CountPolls<'_> { + type Output = (); + + fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll { + self.0.fetch_add(1, Ordering::Relaxed); + Poll::Ready(()) + } + } + + fn runtime() -> Runtime { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("test runtime") + } + + #[test] + fn ready_future_completes_on_first_poll() { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .expect("runtime without I/O or time drivers"); + + let output = + block_on_sync_first(&runtime, std::future::ready(42)).expect("outside runtime"); + + assert_eq!(output, 42); + } + + #[test] + fn pending_future_resumes_with_runtime_waker_without_reconstruction() { + let runtime = runtime(); + let polls = Arc::new(AtomicUsize::new(0)); + let future = PendingThenReady { + address: None, + first_waker: None, + polls: Arc::clone(&polls), + }; + + let output = block_on_sync_first(&runtime, future).expect("outside runtime"); + + assert_eq!(output, 2); + assert_eq!(polls.load(Ordering::Relaxed), 2); + } + + #[test] + fn pending_future_resumes_on_current_thread_runtime() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("current-thread runtime"); + let polls = Arc::new(AtomicUsize::new(0)); + let future = PendingThenReady { + address: None, + first_waker: None, + polls: Arc::clone(&polls), + }; + + let output = block_on_sync_first(&runtime, future).expect("outside runtime"); + + assert_eq!(output, 2); + assert_eq!(polls.load(Ordering::Relaxed), 2); + } + + #[test] + fn cancellation_survives_the_waker_handoff() { + let runtime = runtime(); + let cancel = CancelHandle::new(); + let child = cancel.child_handle(); + let canceller = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(10)); + cancel.cancel(); + }); + let pending = std::future::pending::>(); + + let result = block_on_sync_first( + &runtime, + CancelHandle::run_until_cancelled(Some(&child), pending), + ) + .expect("outside runtime"); + canceller.join().expect("canceller thread"); + + assert!(matches!(result, Err(Error::OperationCancelledError(_)))); + } + + #[test] + fn exhausted_timeout_survives_the_waker_handoff() { + let runtime = runtime(); + let future = + async { tokio::time::timeout(Duration::ZERO, std::future::pending::<()>()).await }; + + let result = block_on_sync_first(&runtime, future).expect("outside runtime"); + + assert!( + result.is_err(), + "a pending future must not evade a zero timeout" + ); + } + + #[test] + fn nested_runtime_is_rejected_before_polling() { + let runtime = runtime(); + let polls = AtomicUsize::new(0); + + runtime.block_on(async { + let error = block_on_sync_first(&runtime, CountPolls(&polls)) + .expect_err("nested execution must be rejected"); + assert_eq!(error, SyncFirstError::RuntimeContext); + }); + + assert_eq!( + polls.load(Ordering::Relaxed), + 0, + "nested rejection must happen before a stateful future advances" + ); + } + + #[test] + fn blocking_pool_runtime_context_is_rejected_before_polling() { + let runtime = Arc::new(runtime()); + let worker_runtime = Arc::clone(&runtime); + let polls = Arc::new(AtomicUsize::new(0)); + let worker_polls = Arc::clone(&polls); + + let error = runtime.block_on(async move { + tokio::task::spawn_blocking(move || { + block_on_sync_first(&worker_runtime, CountPolls(&worker_polls)) + .expect_err("blocking-pool threads carry an entered runtime handle") + }) + .await + .expect("blocking task") + }); + + assert_eq!(error, SyncFirstError::RuntimeContext); + assert_eq!( + polls.load(Ordering::Relaxed), + 0, + "context rejection must happen before a stateful future advances" + ); + } +} diff --git a/mssql-tds/src/test_packet_support.rs b/mssql-tds/src/test_packet_support.rs index 1d2d9456..a1c832c7 100644 --- a/mssql-tds/src/test_packet_support.rs +++ b/mssql-tds/src/test_packet_support.rs @@ -12,6 +12,7 @@ use byteorder::{BigEndian, ByteOrder, LittleEndian}; use tokio::io::{AsyncWriteExt, DuplexStream, duplex}; +use tokio::sync::oneshot; use crate::connection::client_context::ClientContext; use crate::connection::transport::network_transport::NetworkTransport; @@ -55,6 +56,11 @@ impl TestPacketBuilder { self } + pub(crate) fn end_of_message(&mut self, enabled: bool) -> &mut TestPacketBuilder { + self.data[1] = u8::from(enabled); + self + } + append_method!(append_u16, u16, 2, write_u16); append_method!(append_i16, i16, 2, write_i16); append_method!(append_f32, f32, 4, write_f32); @@ -133,6 +139,29 @@ pub(crate) fn create_network_transport_with_chunked_data( build_duplex_transport(client_side) } +/// Builds a `NetworkTransport` with `first` immediately available and `second` +/// held until the returned gate is opened. +pub(crate) fn create_network_transport_with_gated_data( + first: &[u8], + second: &[u8], +) -> (NetworkTransport, oneshot::Sender<()>) { + let (client_side, mut server_side) = duplex(first.len().max(1)); + let first = first.to_vec(); + let second = second.to_vec(); + let (release, wait) = oneshot::channel(); + tokio::spawn(async move { + if server_side.write_all(&first).await.is_err() { + return; + } + if wait.await.is_err() { + return; + } + let _ = server_side.write_all(&second).await; + }); + + (build_duplex_transport(client_side), release) +} + #[cfg(test)] mod tests { use super::*;