diff --git a/mssql-odbc/.gitignore b/mssql-odbc/.gitignore index 8c1e5fa6..b0037cf4 100644 --- a/mssql-odbc/.gitignore +++ b/mssql-odbc/.gitignore @@ -12,6 +12,15 @@ tests/e2e/CMakeFiles/ tests/e2e/cmake_install.cmake tests/e2e/_deps/ +# CMake perf benchmark build output and results +tests/perf/build/ +tests/perf/build_*/ +tests/perf/results/ +tests/perf/CMakeCache.txt +tests/perf/CMakeFiles/ +tests/perf/cmake_install.cmake +tests/perf/_deps/ + # Compiled objects and libraries *.o *.obj diff --git a/mssql-odbc/src/api/close_cursor.rs b/mssql-odbc/src/api/close_cursor.rs index 5ae54825..2b809427 100644 --- a/mssql-odbc/src/api/close_cursor.rs +++ b/mssql-odbc/src/api/close_cursor.rs @@ -132,6 +132,8 @@ fn sql_free_stmt_close_safe(statement_handle: SqlHandle, stmt: &StmtHandle) -> S pub(super) fn reset_cursor_state(stmt_state: &mut crate::handles::stmt::StmtState) { stmt_state.clear_state(STMT_STATE_CURSOR_OPEN | STMT_STATE_EXEC_CONTEXT); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); stmt_state.column_metadata.clear(); stmt_state.pending_row_counts.clear(); } diff --git a/mssql-odbc/src/api/driver_connect.rs b/mssql-odbc/src/api/driver_connect.rs index a4d193d0..2cdfff1b 100644 --- a/mssql-odbc/src/api/driver_connect.rs +++ b/mssql-odbc/src/api/driver_connect.rs @@ -338,7 +338,7 @@ fn do_connect( let has_server_info = post_tds_info_messages(state, &info_messages); - state.client = Some(client); + state.client = Some(Box::new(client)); state.connection_state = ConnectionState::Connected; debug!("SQLDriverConnectW: connected successfully"); diff --git a/mssql-odbc/src/api/exec_common.rs b/mssql-odbc/src/api/exec_common.rs index d5aa7486..c081bc95 100644 --- a/mssql-odbc/src/api/exec_common.rs +++ b/mssql-odbc/src/api/exec_common.rs @@ -40,7 +40,7 @@ pub(super) fn claim_connection( stmt: &StmtHandle, statement_handle: SqlHandle, op: &str, -) -> Result { +) -> Result, SqlReturn> { let Ok(mut dbc_state) = dbc.inner.lock() else { error!("{op}: dbc mutex poisoned"); clear_exec_started(stmt); @@ -88,7 +88,11 @@ pub(super) fn claim_connection( /// Returns `client` to the DBC and releases the busy claim. Used on the /// DDL/DML success path and on error recovery. -pub(super) fn return_client_idle(dbc: &DbcHandle, statement_handle: SqlHandle, client: TdsClient) { +pub(super) fn return_client_idle( + dbc: &DbcHandle, + statement_handle: SqlHandle, + client: Box, +) { if let Ok(mut dbc_state) = dbc.inner.lock() { dbc_state.client = Some(client); if dbc_state.active_stmt == Some(statement_handle) { @@ -109,7 +113,7 @@ pub(super) fn return_client_idle(dbc: &DbcHandle, statement_handle: SqlHandle, c pub(super) fn try_claim_idle_client( dbc: &DbcHandle, statement_handle: SqlHandle, -) -> Option { +) -> Option> { let Ok(mut dbc_state) = dbc.inner.lock() else { return None; }; @@ -123,7 +127,7 @@ pub(super) fn try_claim_idle_client( /// Returns `client` to the DBC but **keeps** the busy claim — used when a /// cursor is left open for `SQLFetch`. -pub(super) fn return_client_busy(dbc: &DbcHandle, client: TdsClient) { +pub(super) fn return_client_busy(dbc: &DbcHandle, client: Box) { if let Ok(mut dbc_state) = dbc.inner.lock() { dbc_state.client = Some(client); } @@ -136,7 +140,7 @@ pub(super) fn fail_with_tds( dbc: &DbcHandle, stmt: &StmtHandle, statement_handle: SqlHandle, - mut client: TdsClient, + mut client: Box, err: &TdsError, ) -> SqlReturn { let info_messages = client.take_info_messages(); @@ -256,7 +260,7 @@ pub(super) fn finish_execute( dbc: &DbcHandle, stmt: &StmtHandle, statement_handle: SqlHandle, - mut client: TdsClient, + mut client: Box, op: &str, ) -> SqlReturn { let metadata = client.get_metadata().clone(); diff --git a/mssql-odbc/src/api/exec_direct.rs b/mssql-odbc/src/api/exec_direct.rs index 19103ef4..8a41964b 100644 --- a/mssql-odbc/src/api/exec_direct.rs +++ b/mssql-odbc/src/api/exec_direct.rs @@ -107,6 +107,8 @@ fn sql_exec_direct_w_safe( stmt_state.clear_state(STMT_STATE_EXEC_CONTEXT); stmt_state.column_metadata.clear(); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); stmt_state.row_count = -1; stmt_state.pending_row_counts.clear(); stmt_state.prepared_sql = None; @@ -285,7 +287,7 @@ mod tests { ]); { let mut ds = dbc.inner.lock().unwrap(); - ds.client = Some(client); + ds.client = Some(Box::new(client)); // active_stmt stays None => connection idle and claimable. } @@ -325,7 +327,7 @@ mod tests { ]); { let mut ds = dbc.inner.lock().unwrap(); - ds.client = Some(client); + ds.client = Some(Box::new(client)); } let stmt = unsafe { handle_from_raw::(h.stmt) }; diff --git a/mssql-odbc/src/api/execute.rs b/mssql-odbc/src/api/execute.rs index d2a6d61a..f820d0b4 100644 --- a/mssql-odbc/src/api/execute.rs +++ b/mssql-odbc/src/api/execute.rs @@ -159,6 +159,8 @@ fn stage_execution(stmt: &StmtHandle) -> Result { stmt_state.clear_state(STMT_STATE_EXEC_CONTEXT); stmt_state.column_metadata.clear(); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); stmt_state.row_count = -1; stmt_state.pending_row_counts.clear(); stmt_state.set_state(STMT_STATE_EXEC_STARTED); diff --git a/mssql-odbc/src/api/fetch.rs b/mssql-odbc/src/api/fetch.rs index 7e1e88f9..1c2212e8 100644 --- a/mssql-odbc/src/api/fetch.rs +++ b/mssql-odbc/src/api/fetch.rs @@ -13,7 +13,6 @@ use crate::api::odbc_types::{ use crate::error::free_errors; use crate::handles::stmt::STMT_STATE_CURSOR_OPEN; use crate::handles::{HandleType, StmtHandle, handle_from_raw}; -use mssql_tds::connection::tds_client::ResultSet; /// Implements SQLFetch for the current forward-only result set. /// @@ -54,8 +53,35 @@ fn sql_fetch_safe(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn { fetch_rows_next(statement_handle, stmt) } +/// Rows decoded per wire round through the async decode path. The per-row cost +/// of a cursor is dominated by building and polling that state machine, so +/// draining a batch amortizes it; the cap bounds the memory a single fetch can +/// buffer for wide rows. +const FETCH_BATCH_ROWS: usize = 64; + /// Row materialization step for one forward fetch operation. fn fetch_rows_next(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn { + // Fast path: a previous fetch already decoded this row. The batch is only + // ever non-empty while this statement holds the connection, so the busy and + // cursor checks below are already satisfied and the DBC lock is not needed. + // + // Server INFO is not drained here. It cannot be attributed to a specific + // row anyway, and it stays on the client's buffer until the next call that + // can report it (see the SQL_NO_DATA arm below). + { + let Ok(mut stmt_state) = stmt.inner.lock() else { + error!("SQLFetch: stmt mutex poisoned serving prefetched row"); + return SQL_ERROR; + }; + if let Some(row) = stmt_state.row_batch.pop_front() { + if let Some(previous) = stmt_state.current_row.replace(row) { + stmt_state.row_batch_spare.push(previous); + } + stmt_state.reset_get_data_cursor(); + return SQL_SUCCESS; + } + } + let dbc = stmt.parent_dbc(); let mut client = { @@ -130,10 +156,23 @@ fn fetch_rows_next(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn } } - let fetch_result = dbc.runtime.block_on(client.next_row()); + let spare = match stmt.inner.lock() { + Ok(mut ss) => { + let mut spare = std::mem::take(&mut ss.row_batch_spare); + if let Some(previous) = ss.current_row.take() { + spare.push(previous); + } + spare + } + Err(_) => Vec::new(), + }; + let mut rows = Vec::new(); + let fetch_result = + dbc.runtime + .block_on(client.fetch_rows_batch(&mut rows, spare, FETCH_BATCH_ROWS)); match fetch_result { - Ok(Some(row)) => { + Ok(count) if count > 0 => { let Ok(mut stmt_state) = stmt.inner.lock() else { error!("SQLFetch: stmt mutex poisoned storing row"); if let Ok(mut ds) = dbc.inner.lock() { @@ -144,7 +183,10 @@ fn fetch_rows_next(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn } return SQL_ERROR; }; - stmt_state.current_row = Some(row); + let mut rows = rows.into_iter(); + stmt_state.current_row = rows.next(); + stmt_state.row_batch.extend(rows); + stmt_state.reset_get_data_cursor(); // Drain INFO only after the lock is held so a poisoned mutex cannot // silently drop the messages. let info_messages = client.take_info_messages(); @@ -163,7 +205,7 @@ fn fetch_rows_next(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn SQL_SUCCESS } } - Ok(None) => { + Ok(_) => { // End of current rowset. SQLFetch must return SQL_NO_DATA here per the // cursor contract, and SQL_NO_DATA cannot be upgraded to // SQL_SUCCESS_WITH_INFO — so this call has no way to signal "there are @@ -201,6 +243,7 @@ fn fetch_rows_next(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn return SQL_ERROR; }; stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); // Don't clear CURSOR_OPEN here: the cursor stays open until // SQLMoreResults / SQLCloseCursor / SQLFreeStmt(SQL_CLOSE). drop(stmt_state); @@ -215,6 +258,7 @@ fn fetch_rows_next(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlReturn error!(%e, "SQLFetch: row fetch failed"); if let Ok(mut stmt_state) = stmt.inner.lock() { stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); stmt_state.clear_state(STMT_STATE_CURSOR_OPEN); post_tds_error(&mut stmt_state, &e, SQLSTATE_HY000); let info_messages = client.take_info_messages(); @@ -333,7 +377,7 @@ mod tests { let dbc_handle = unsafe { handle_from_raw::(h.dbc) }; { let mut dbc_state = dbc_handle.inner.lock().unwrap(); - dbc_state.client = Some(client); + dbc_state.client = Some(Box::new(client)); dbc_state.active_stmt = Some(h.stmt); } diff --git a/mssql-odbc/src/api/get_data.rs b/mssql-odbc/src/api/get_data.rs index 1714c82a..8b31f927 100644 --- a/mssql-odbc/src/api/get_data.rs +++ b/mssql-odbc/src/api/get_data.rs @@ -6,25 +6,26 @@ use tracing::{debug, error}; use super::odbc_types::{ - SQL_C_CHAR, SQL_C_WCHAR, SQL_ERROR, SQL_INVALID_HANDLE, SQL_NULL_DATA, SQL_SUCCESS, - SQL_SUCCESS_WITH_INFO, SqlHandle, SqlLen, SqlPointer, SqlReturn, SqlSmallInt, SqlUSmallInt, + SQL_C_CHAR, SQL_C_WCHAR, SQL_ERROR, SQL_INVALID_HANDLE, SQL_NO_DATA, SQL_NULL_DATA, + SQL_SUCCESS, SQL_SUCCESS_WITH_INFO, SqlHandle, SqlLen, SqlPointer, SqlReturn, SqlSmallInt, + SqlUSmallInt, }; use super::sqlstate::*; use crate::api::odbc_types::SqlWChar; use crate::api::util::{copy_with_nul, write_if_some}; use crate::error::{free_errors, post_sql_error}; -use crate::handles::stmt::STMT_STATE_CURSOR_OPEN; +use crate::handles::stmt::{GetDataCursor, GetDataPayload, STMT_STATE_CURSOR_OPEN}; use crate::handles::{HandleType, StmtHandle, handle_from_raw}; -use mssql_tds::datatypes::column_values::ColumnValues; +use mssql_tds::datatypes::column_values::{ColumnValues, SqlTime}; +use mssql_tds::datatypes::sql_string::EncodingType; /// Implements SQLGetData for current-row retrieval. /// -/// Phase 1 scope: +/// Scope: /// - Requires an open cursor and a current fetched row. -/// - Supports only `SQL_C_CHAR` output. -/// - Supports basic scalar conversion to UTF-8 text. -/// - Repeated calls on the same column do not advance an offset; each call -/// returns the same prefix for the current value (no chunked streaming yet). +/// - Supports `SQL_C_CHAR` and `SQL_C_WCHAR` output. +/// - Supports chunked reads: repeated calls on the same column advance through +/// the value and report `SQL_NO_DATA` once it has been fully delivered. pub(crate) unsafe fn sql_get_data( statement_handle: SqlHandle, column_number: SqlUSmallInt, @@ -110,13 +111,13 @@ fn sql_get_data_safe( return SQL_ERROR; } - let Some(row) = stmt_state.current_row.as_ref() else { + let Some(row_len) = stmt_state.current_row.as_ref().map(Vec::len) else { post_sql_error(&mut stmt_state, SQLSTATE_24000, 0, "No current row"); return SQL_ERROR; }; let col_index = usize::from(column_number); - if col_index == 0 || col_index > row.len() { + if col_index == 0 || col_index > row_len { post_diag(&mut stmt_state, ERR_INVALID_DESCRIPTOR_INDEX); return SQL_ERROR; } @@ -133,85 +134,308 @@ fn sql_get_data_safe( // Output buffer capacity in element units (u8 for SQL_C_CHAR, SqlWChar for // SQL_C_WCHAR). buffer_length is always in bytes per the ODBC spec. - let buf_elements = if target_type == SQL_C_WCHAR { + let wide = target_type == SQL_C_WCHAR; + let buf_elements = if wide { (buffer_length as usize) / std::mem::size_of::() } else { buffer_length as usize }; - let value = &row[col_index - 1]; - if matches!(value, ColumnValues::Null) { - unsafe { write_if_some(strlen_or_ind_ptr, SQL_NULL_DATA) }; - // Write a NUL terminator into the caller buffer when there's room. The - // helper handles null `dst` and zero-length uniformly. - if target_type == SQL_C_WCHAR { - unsafe { - copy_with_nul(target_value_ptr as *mut SqlWChar, buf_elements, &[]); - } - } else { - unsafe { - copy_with_nul(target_value_ptr as *mut u8, buf_elements, &[]); + // Continue an in-progress chunked read of this column, if any. A different + // column or a different target type restarts the read. + match stmt_state.get_data_cursor.take() { + Some(cursor) if cursor.column == column_number && cursor.wide == wide => { + let Some(payload) = cursor.payload else { + // The value was already delivered in full. + return SQL_NO_DATA; + }; + return continue_chunked_read( + &mut stmt_state, + column_number, + wide, + payload, + cursor.offset, + target_value_ptr, + buf_elements, + strlen_or_ind_ptr, + ); + } + _ => {} + } + + // Zero-copy fast path: an all-ASCII `varchar`/`char` value that fits the + // caller's buffer is copied straight out of the row. This covers the bulk + // of `SQL_C_CHAR` reads and skips both the code-page transcode and the + // intermediate payload allocation that the general path needs for chunking. + if !wide && !target_value_ptr.is_null() && buf_elements > 0 { + let direct_len = stmt_state.current_row.as_ref().and_then(|row| { + let ColumnValues::String(s) = &row[col_index - 1] else { + return None; + }; + let single_byte = matches!( + s.encoding_type(), + EncodingType::Utf8 | EncodingType::LcidBased(_) + ); + if !single_byte || s.bytes.len() >= buf_elements || !s.bytes.is_ascii() { + return None; } + unsafe { copy_with_nul(target_value_ptr as *mut u8, buf_elements, &s.bytes) }; + Some(s.bytes.len()) + }); + if let Some(len) = direct_len { + unsafe { write_if_some(strlen_or_ind_ptr, len as SqlLen) }; + stmt_state.get_data_cursor = Some(GetDataCursor::exhausted(column_number, wide)); + return SQL_SUCCESS; } - return SQL_SUCCESS; } - let Some(as_text) = column_value_to_text(value) else { - post_sql_error( - &mut stmt_state, - SQLSTATE_HYC00, - 0, - "Column type conversion not yet implemented", - ); - return SQL_ERROR; + // Convert the cell before touching `stmt_state` mutably, so the borrow of + // the current row ends here. + let converted = { + let Some(row) = stmt_state.current_row.as_ref() else { + post_sql_error(&mut stmt_state, SQLSTATE_24000, 0, "No current row"); + return SQL_ERROR; + }; + let value = &row[col_index - 1]; + if matches!(value, ColumnValues::Null) { + Converted::Null + } else if wide { + match column_value_to_utf16(value) { + Some(v) => Converted::Wide(v), + None => Converted::Unsupported, + } + } else { + match column_value_to_utf8_bytes(value) { + Some(v) => Converted::Narrow(v), + None => Converted::Unsupported, + } + } }; - if target_type == SQL_C_WCHAR { - let utf16: Vec = as_text.encode_utf16().collect(); - write_string_result( + match converted { + Converted::Unsupported => { + post_unsupported_conversion(&mut stmt_state); + SQL_ERROR + } + Converted::Null => { + unsafe { write_if_some(strlen_or_ind_ptr, SQL_NULL_DATA) }; + // Write a NUL terminator into the caller buffer when there's room. + // The helper handles null `dst` and zero-length uniformly. + if wide { + unsafe { + copy_with_nul(target_value_ptr as *mut SqlWChar, buf_elements, &[]); + } + } else { + unsafe { + copy_with_nul(target_value_ptr as *mut u8, buf_elements, &[]); + } + } + stmt_state.get_data_cursor = Some(GetDataCursor::exhausted(column_number, wide)); + SQL_SUCCESS + } + Converted::Wide(utf16) => deliver_first_chunk( &mut stmt_state, - &utf16, + column_number, + wide, + utf16, target_value_ptr as *mut SqlWChar, buf_elements, strlen_or_ind_ptr, - ) - } else { - write_string_result( + GetDataPayload::Wide, + ), + Converted::Narrow(bytes) => deliver_first_chunk( &mut stmt_state, - as_text.as_bytes(), + column_number, + wide, + bytes, target_value_ptr as *mut u8, buf_elements, strlen_or_ind_ptr, - ) + GetDataPayload::Narrow, + ), + } +} + +/// A cell converted to the element width requested by the caller. +enum Converted { + Null, + Narrow(Vec), + Wide(Vec), + /// The column type has no text conversion yet. + Unsupported, +} + +fn post_unsupported_conversion(stmt_state: &mut crate::handles::stmt::StmtState) { + post_sql_error( + stmt_state, + SQLSTATE_HYC00, + 0, + "Column type conversion not yet implemented", + ); +} + +/// Writes the first chunk of a freshly converted value and records how much of +/// it remains, so a truncated value can be resumed by the next call. +/// +/// The full payload is buffered only when the value did not fit; a value that +/// fits leaves behind a cursor with no buffer. +#[allow(clippy::too_many_arguments)] +fn deliver_first_chunk( + stmt_state: &mut crate::handles::stmt::StmtState, + column_number: SqlUSmallInt, + wide: bool, + payload: Vec, + target_value_ptr: *mut T, + buf_elements: usize, + strlen_or_ind_ptr: *mut SqlLen, + wrap: fn(Vec) -> GetDataPayload, +) -> SqlReturn { + let written = write_chunk( + stmt_state, + &payload, + target_value_ptr, + buf_elements, + strlen_or_ind_ptr, + ); + match written { + ChunkOutcome::Complete => { + stmt_state.get_data_cursor = Some(GetDataCursor::exhausted(column_number, wide)); + SQL_SUCCESS + } + ChunkOutcome::Truncated(delivered) => { + stmt_state.get_data_cursor = Some(GetDataCursor { + column: column_number, + wide, + payload: Some(wrap(payload)), + offset: delivered, + }); + SQL_SUCCESS_WITH_INFO + } + } +} + +/// Resumes a chunked read from `offset` using the already-converted payload, so +/// a long value is converted once rather than once per chunk. +#[allow(clippy::too_many_arguments)] +fn continue_chunked_read( + stmt_state: &mut crate::handles::stmt::StmtState, + column_number: SqlUSmallInt, + wide: bool, + payload: GetDataPayload, + offset: usize, + target_value_ptr: SqlPointer, + buf_elements: usize, + strlen_or_ind_ptr: *mut SqlLen, +) -> SqlReturn { + macro_rules! resume { + ($buf:expr, $ptr_ty:ty, $wrap:expr) => {{ + let outcome = write_chunk( + stmt_state, + &$buf[offset..], + target_value_ptr as *mut $ptr_ty, + buf_elements, + strlen_or_ind_ptr, + ); + match outcome { + ChunkOutcome::Complete => { + stmt_state.get_data_cursor = + Some(GetDataCursor::exhausted(column_number, wide)); + SQL_SUCCESS + } + ChunkOutcome::Truncated(delivered) => { + stmt_state.get_data_cursor = Some(GetDataCursor { + column: column_number, + wide, + payload: Some($wrap($buf)), + offset: offset + delivered, + }); + SQL_SUCCESS_WITH_INFO + } + } + }}; } + + match payload { + GetDataPayload::Narrow(buf) => resume!(buf, u8, GetDataPayload::Narrow), + GetDataPayload::Wide(buf) => resume!(buf, SqlWChar, GetDataPayload::Wide), + } +} + +/// Result of copying one chunk into the application buffer. +enum ChunkOutcome { + /// The remaining value fit entirely. + Complete, + /// The buffer filled up; the payload element count that was delivered. + Truncated(usize), } /// Writes `src` to the caller's output buffer with ODBC string semantics: -/// the indicator (when present) reports the untruncated byte length, the -/// payload is NUL-terminated within the buffer, and truncation is reported via -/// SQLSTATE 01004 + `SQL_SUCCESS_WITH_INFO`. +/// the indicator (when present) reports the untruncated byte length of what +/// remains, the payload is NUL-terminated within the buffer, and truncation is +/// reported via SQLSTATE 01004 + `SQL_SUCCESS_WITH_INFO`. /// /// `buf_elements` is the buffer capacity in units of `T` (not bytes). /// /// The caller-provided pointers are written through small `unsafe` blocks /// inside this function; both pointer arguments are obligations of the FFI /// caller (validated against the buffer length passed by the DM). -fn write_string_result( +fn write_chunk( stmt_state: &mut crate::handles::stmt::StmtState, src: &[T], target_value_ptr: *mut T, buf_elements: usize, strlen_or_ind_ptr: *mut SqlLen, -) -> SqlReturn { +) -> ChunkOutcome { let byte_len = std::mem::size_of_val(src) as SqlLen; unsafe { write_if_some(strlen_or_ind_ptr, byte_len) }; let truncated = unsafe { copy_with_nul(target_value_ptr, buf_elements, src) }; if truncated { post_diag(stmt_state, ERR_STRING_RIGHT_TRUNCATION); - SQL_SUCCESS_WITH_INFO + // `copy_with_nul` reserves the final element for the terminator. + ChunkOutcome::Truncated(buf_elements.saturating_sub(1)) } else { - SQL_SUCCESS + ChunkOutcome::Complete + } +} + +/// Converts a column value to UTF-8 bytes for `SQL_C_CHAR` output. +fn column_value_to_utf8_bytes(v: &ColumnValues) -> Option> { + match v { + // Already UTF-8 on the wire: hand back the bytes without transcoding. + ColumnValues::String(s) if matches!(s.encoding_type(), EncodingType::Utf8) => { + Some(s.bytes.clone()) + } + // Every single-byte SQL Server code page agrees with US-ASCII below + // 0x80, so all-ASCII payloads are already their own UTF-8 encoding. + // The check is vectorized and skips a full code-page transcode, which + // is the common case for `varchar` columns. + ColumnValues::String(s) + if matches!(s.encoding_type(), EncodingType::LcidBased(_)) && s.bytes.is_ascii() => + { + Some(s.bytes.clone()) + } + _ => column_value_to_text(v).map(String::into_bytes), + } +} + +/// Converts a column value to UTF-16 code units for `SQL_C_WCHAR` output. +/// +/// Character data arrives from SQL Server as UTF-16LE, which is exactly what +/// `SQL_C_WCHAR` wants. Reinterpreting those bytes avoids a UTF-16 -> UTF-8 -> +/// UTF-16 round trip through an intermediate `String`. +fn column_value_to_utf16(v: &ColumnValues) -> Option> { + if let ColumnValues::String(s) = v + && let Some(bytes) = s.as_utf16_bytes() + && bytes.len() % 2 == 0 + { + return Some( + bytes + .chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect(), + ); } + column_value_to_text(v).map(|t| t.encode_utf16().collect()) } fn column_value_to_text(v: &ColumnValues) -> Option { @@ -226,10 +450,116 @@ fn column_value_to_text(v: &ColumnValues) -> Option { ColumnValues::String(s) => Some(s.to_utf8_string()), ColumnValues::Uuid(u) => Some(u.to_string()), ColumnValues::Null => Some(String::new()), + ColumnValues::Decimal(d) | ColumnValues::Numeric(d) => Some(d.to_string()), + ColumnValues::Date(d) => Some(format_date(d.get_days())), + ColumnValues::Time(t) => Some(format_time(t)), + ColumnValues::DateTime2(dt) => Some(format!( + "{} {}", + format_date(dt.days), + format_time(&dt.time) + )), + ColumnValues::DateTimeOffset(dto) => { + let (sign, mins) = if dto.offset < 0 { + ('-', (-(dto.offset as i32)) as u32) + } else { + ('+', dto.offset as u32) + }; + Some(format!( + "{} {} {sign}{:02}:{:02}", + format_date(dto.datetime2.days), + format_time(&dto.datetime2.time), + mins / 60, + mins % 60 + )) + } + // `datetime` counts days from 1900-01-01 in 1/300-second ticks. + ColumnValues::DateTime(dt) => { + let days = DAYS_0001_TO_1900 as i64 + dt.days as i64; + let nanos = (dt.time as u64) * 10_000_000 / 3; + Some(format!( + "{} {}", + format_date(days.clamp(0, u32::MAX as i64) as u32), + format_nanos(nanos, 3) + )) + } + ColumnValues::SmallDateTime(dt) => { + let days = DAYS_0001_TO_1900 + u32::from(dt.days); + let nanos = u64::from(dt.time) * 60 * 1_000_000_000; + Some(format!("{} {}", format_date(days), format_nanos(nanos, 0))) + } + ColumnValues::SmallMoney(m) => Some(format_money(i64::from(m.int_val))), + ColumnValues::Money(m) => { + let scaled = (i64::from(m.msb_part) << 32) | (i64::from(m.lsb_part) & 0xFFFF_FFFF); + Some(format_money(scaled)) + } + ColumnValues::Bytes(b) => { + let mut s = String::with_capacity(b.len() * 2); + for byte in b { + use std::fmt::Write; + let _ = write!(s, "{byte:02X}"); + } + Some(s) + } + ColumnValues::Xml(x) => Some(x.as_string()), _ => None, } } +/// Days from 0001-01-01 to 1900-01-01, the epoch used by `datetime` and +/// `smalldatetime`. +const DAYS_0001_TO_1900: u32 = 693_595; + +/// Formats a day count from 0001-01-01 as `YYYY-MM-DD`. +/// +/// Uses the civil-from-days algorithm, shifting the year to start in March so +/// the leap day falls at the end of the cycle and month lengths follow a +/// regular pattern. +fn format_date(days_since_year_one: u32) -> String { + // Re-base onto the 1970-01-01 era the algorithm is defined against. + const DAYS_0001_TO_1970: i64 = 719_162; + let z = days_since_year_one as i64 - DAYS_0001_TO_1970 + 719_468; + let era = z.div_euclid(146_097); + let doe = z.rem_euclid(146_097); + let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; + let y = yoe + era * 400; + let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); + let mp = (5 * doy + 2) / 153; + let d = doy - (153 * mp + 2) / 5 + 1; + let m = if mp < 10 { mp + 3 } else { mp - 9 }; + let y = if m <= 2 { y + 1 } else { y }; + format!("{y:04}-{m:02}-{d:02}") +} + +/// Formats a [`SqlTime`] as `HH:MM:SS[.fffffff]`, honouring its scale. +fn format_time(t: &SqlTime) -> String { + format_nanos(t.time_nanoseconds, t.scale) +} + +/// Formats nanoseconds since midnight as `HH:MM:SS[.fffffff]`, emitting +/// `scale` fractional digits (none when `scale` is 0). +fn format_nanos(nanos: u64, scale: u8) -> String { + let secs = nanos / 1_000_000_000; + let frac = nanos % 1_000_000_000; + let (h, m, s) = (secs / 3600, (secs / 60) % 60, secs % 60); + if scale == 0 { + return format!("{h:02}:{m:02}:{s:02}"); + } + let scale = scale.min(9) as u32; + let divisor = 10u64.pow(9 - scale); + format!( + "{h:02}:{m:02}:{s:02}.{:0width$}", + frac / divisor, + width = scale as usize + ) +} + +/// Formats a money value stored as a ×10⁴ scaled integer. +fn format_money(scaled: i64) -> String { + let sign = if scaled < 0 { "-" } else { "" }; + let abs = scaled.unsigned_abs(); + format!("{sign}{}.{:04}", abs / 10_000, abs % 10_000) +} + #[cfg(test)] mod tests { use super::*; @@ -237,6 +567,52 @@ mod tests { use crate::test_support::TestHandles; use mssql_tds::datatypes::sql_string::SqlString; + #[test] + fn format_date_known_epochs() { + // 1900-01-01 is the SQL Server datetime epoch. + assert_eq!(format_date(DAYS_0001_TO_1900), "1900-01-01"); + assert_eq!(format_date(0), "0001-01-01"); + assert_eq!(format_date(719_162), "1970-01-01"); + // 2000 is a leap year; day 60 of that year is 2000-02-29. + assert_eq!(format_date(DAYS_0001_TO_1900 + 36_583), "2000-02-29"); + } + + #[test] + fn format_money_scales_and_signs() { + assert_eq!(format_money(0), "0.0000"); + assert_eq!(format_money(1), "0.0001"); + assert_eq!(format_money(123_456), "12.3456"); + assert_eq!(format_money(-123_456), "-12.3456"); + assert_eq!( + format_money(9_223_372_036_854_775_807), + "922337203685477.5807" + ); + } + + #[test] + fn money_mixed_endian_reassembly() { + // TDS transmits `money` as MSB i32 then LSB i32. A negative value has an + // all-ones MSB word that must not sign-extend the LSB word. + let reassemble = + |msb: i32, lsb: i32| (i64::from(msb) << 32) | (i64::from(lsb) & 0xFFFF_FFFF); + let expected: i64 = -123_456; + let msb = (expected >> 32) as i32; + let lsb = expected as i32; + assert_eq!(reassemble(msb, lsb), expected); + assert_eq!(format_money(reassemble(msb, lsb)), "-12.3456"); + + let big: i64 = 9_223_372_036_854_775_807; + assert_eq!(reassemble((big >> 32) as i32, big as i32), big); + } + + #[test] + fn format_nanos_honours_scale() { + assert_eq!(format_nanos(0, 0), "00:00:00"); + assert_eq!(format_nanos(3_723_000_000_000, 0), "01:02:03"); + assert_eq!(format_nanos(3_723_123_456_700, 7), "01:02:03.1234567"); + assert_eq!(format_nanos(3_723_123_456_700, 3), "01:02:03.123"); + } + #[test] fn get_data_null_handle() { let ret = unsafe { diff --git a/mssql-odbc/src/api/get_type_info.rs b/mssql-odbc/src/api/get_type_info.rs index c76decde..fd8692a3 100644 --- a/mssql-odbc/src/api/get_type_info.rs +++ b/mssql-odbc/src/api/get_type_info.rs @@ -149,6 +149,8 @@ fn sql_get_type_info_w_safe( stmt_state.clear_state(STMT_STATE_EXEC_CONTEXT); stmt_state.column_metadata.clear(); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); stmt_state.prepared_sql = None; // A cached prepared plan is superseded; release its server handle // (deferred) once we hold the client below. diff --git a/mssql-odbc/src/api/more_results.rs b/mssql-odbc/src/api/more_results.rs index 6ef0a2b7..68417d1a 100644 --- a/mssql-odbc/src/api/more_results.rs +++ b/mssql-odbc/src/api/more_results.rs @@ -115,6 +115,8 @@ fn sql_more_results_safe(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlR // Refresh the count for the newly-positioned result set (-1 for a SELECT). stmt_state.row_count = client.last_rows_affected(); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); // Drain INFO only after the lock is held. let info_messages = client.take_info_messages(); let has_server_info = post_tds_info_messages(&mut stmt_state, &info_messages); @@ -148,6 +150,8 @@ fn sql_more_results_safe(statement_handle: SqlHandle, stmt: &StmtHandle) -> SqlR // SQLRowCount now that we are positioned on it. stmt_state.row_count = client.last_rows_affected(); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); let info_messages = client.take_info_messages(); let has_server_info = post_tds_info_messages(&mut stmt_state, &info_messages); drop(stmt_state); @@ -246,7 +250,7 @@ mod tests { } { let mut ds = dbc.inner.lock().unwrap(); - ds.client = Some(client); + ds.client = Some(Box::new(client)); ds.active_stmt = Some(h.stmt); } first diff --git a/mssql-odbc/src/api/prepare.rs b/mssql-odbc/src/api/prepare.rs index 74873673..8e8f70d0 100644 --- a/mssql-odbc/src/api/prepare.rs +++ b/mssql-odbc/src/api/prepare.rs @@ -115,6 +115,8 @@ fn sql_prepare_w_safe(stmt: &StmtHandle, sql: String) -> SqlReturn { stmt_state.orphan_prepared_handle(); stmt_state.column_metadata.clear(); stmt_state.current_row = None; + stmt_state.reset_get_data_cursor(); + stmt_state.discard_row_batch(); stmt_state.clear_state(STMT_STATE_EXEC_CONTEXT); stmt_state.set_state(STMT_STATE_PREPARED); diff --git a/mssql-odbc/src/handles/dbc.rs b/mssql-odbc/src/handles/dbc.rs index 54d14ead..0c180831 100644 --- a/mssql-odbc/src/handles/dbc.rs +++ b/mssql-odbc/src/handles/dbc.rs @@ -60,7 +60,7 @@ pub(crate) struct DbcState { /// one statement may hold an open cursor per connection at a time. pub(crate) active_stmt: Option<*mut c_void>, /// Active TDS connection, present only when `connection_state == Connected`. - pub(crate) client: Option, + pub(crate) client: Option>, /// Pre-connect access token set via `SQL_COPT_SS_ACCESS_TOKEN`. /// Consumed by `SQLDriverConnect` to select `AccessToken` authentication. pub(crate) access_token: Option, diff --git a/mssql-odbc/src/handles/stmt.rs b/mssql-odbc/src/handles/stmt.rs index 7d9a713f..741164cf 100644 --- a/mssql-odbc/src/handles/stmt.rs +++ b/mssql-odbc/src/handles/stmt.rs @@ -71,6 +71,16 @@ pub(crate) struct StmtState { pub(crate) pending_unprepare: Option, /// Current fetched row, populated by SQLFetch for later SQLGetData support. pub(crate) current_row: Option>, + /// Rows decoded ahead of the cursor by the last batched fetch, oldest first. + /// `SQLFetch` serves from here before going back to the wire, which keeps + /// the async decode state machine off the per-row path. + pub(crate) row_batch: VecDeque>, + /// Row buffers already consumed by the cursor, recycled by the next batch. + pub(crate) row_batch_spare: Vec>, + /// Chunked-read progress for `SQLGetData`, so successive calls on the same + /// column advance through a long value instead of re-returning the same + /// prefix. Reset by `SQLFetch` and whenever a different column is read. + pub(crate) get_data_cursor: Option, /// Rows affected by the last execution, reported by `SQLRowCount`. `-1` /// means "not available" (no statement executed yet, a result-returning /// SELECT, DDL, or `SET NOCOUNT ON`) — matching msodbcsql's @@ -99,11 +109,69 @@ pub(crate) struct StmtState { pub(crate) state_flags: u32, } +/// Chunked-read progress for `SQLGetData`. +/// +/// ODBC lets an application pull a long value in pieces by calling +/// `SQLGetData` repeatedly on the same column until it returns `SQL_NO_DATA`. +/// That requires the driver to remember how much of the value it has already +/// handed back; without it the application receives the same prefix forever and +/// the loop never terminates. +#[derive(Debug)] +pub(crate) struct GetDataCursor { + /// 1-based column ordinal this cursor tracks. + pub(crate) column: SqlUSmallInt, + /// Whether the buffered payload was produced for `SQL_C_WCHAR`. A change of + /// target type restarts the read rather than reusing an incompatible buffer. + pub(crate) wide: bool, + /// The converted value, buffered only when a chunk had to be truncated. + /// `None` means the value was delivered in full and the next call on this + /// column reports `SQL_NO_DATA`. + pub(crate) payload: Option, + /// Number of elements already delivered from `payload`. + pub(crate) offset: usize, +} + +/// Buffered `SQLGetData` value, in the element width of the requested C type. +#[derive(Debug)] +pub(crate) enum GetDataPayload { + /// `SQL_C_CHAR` payload. + Narrow(Vec), + /// `SQL_C_WCHAR` payload. + Wide(Vec), +} + +impl GetDataCursor { + /// Cursor for a value that was delivered in full, so the next read on the + /// same column reports `SQL_NO_DATA`. Holds no buffer. + pub(crate) fn exhausted(column: SqlUSmallInt, wide: bool) -> Self { + Self { + column, + wide, + payload: None, + offset: 0, + } + } +} + impl StmtState { pub(crate) fn has_state(&self, mask: u32) -> bool { (self.state_flags & mask) != 0 } + /// Discards any chunked `SQLGetData` progress. Called whenever the current + /// row changes so each column's read restarts from the beginning. + pub(crate) fn reset_get_data_cursor(&mut self) { + self.get_data_cursor = None; + } + + /// Drops rows read ahead of the cursor. Called whenever the cursor is + /// closed or repositioned onto a different result set, so a stale batch is + /// never served after the underlying rowset has changed. + pub(crate) fn discard_row_batch(&mut self) { + self.row_batch.clear(); + self.row_batch_spare.clear(); + } + pub(crate) fn set_state(&mut self, mask: u32) { self.state_flags |= mask; } @@ -160,6 +228,9 @@ impl StmtHandle { prepared_handle: None, pending_unprepare: None, current_row: None, + row_batch: VecDeque::new(), + row_batch_spare: Vec::new(), + get_data_cursor: None, row_count: -1, pending_row_counts: VecDeque::new(), row_array_size: 1, diff --git a/mssql-odbc/tests/perf/CMakeLists.txt b/mssql-odbc/tests/perf/CMakeLists.txt new file mode 100644 index 00000000..a2f46075 --- /dev/null +++ b/mssql-odbc/tests/perf/CMakeLists.txt @@ -0,0 +1,92 @@ +cmake_minimum_required(VERSION 3.15) +project(odbc_perf LANGUAGES C CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +# Perf numbers from a debug build are meaningless — default to Release when the +# caller did not pick a configuration for a single-config generator. +if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE Release CACHE STRING "" FORCE) +endif() + +find_package(Threads QUIET) +if(NOT Threads_FOUND) + set(CMAKE_THREAD_LIBS_INIT "-lpthread") + set(THREADS_PREFER_PTHREAD_FLAG ON) +endif() + +# --------------------------------------------------------------------------- +# Google Benchmark via FetchContent (mirrors how tests/e2e fetches GoogleTest) +# --------------------------------------------------------------------------- +include(FetchContent) +set(BENCHMARK_ENABLE_TESTING OFF CACHE BOOL "" FORCE) +set(BENCHMARK_ENABLE_GTEST_TESTS OFF CACHE BOOL "" FORCE) +set(BENCHMARK_ENABLE_INSTALL OFF CACHE BOOL "" FORCE) +set(BENCHMARK_USE_BUNDLED_GTEST OFF CACHE BOOL "" FORCE) +FetchContent_Declare( + benchmark + GIT_REPOSITORY https://github.com/google/benchmark.git + GIT_TAG v1.9.1 +) +FetchContent_MakeAvailable(benchmark) + +# --------------------------------------------------------------------------- +# Platform-specific ODBC libraries +# --------------------------------------------------------------------------- +if(WIN32) + set(ODBC_LIBRARIES odbc32 odbccp32) +else() + find_library(ODBC_LIB NAMES odbc odbc2 iodbc) + if(NOT ODBC_LIB) + message(FATAL_ERROR "Could not find an ODBC library (unixODBC or iODBC). " + "Install unixODBC-dev / libiodbc2-dev or equivalent.") + endif() + find_library(ODBCINST_LIB NAMES odbcinst iodbcinst) + set(ODBC_LIBRARIES ${ODBC_LIB}) + if(ODBCINST_LIB) + list(APPEND ODBC_LIBRARIES ${ODBCINST_LIB}) + endif() + + find_path(ODBC_INCLUDE_DIR sql.h + PATHS /opt/homebrew/include /usr/local/include /usr/include + ) + if(NOT ODBC_INCLUDE_DIR) + message(FATAL_ERROR "Could not find sql.h. Ensure unixODBC headers are installed.") + endif() +endif() + +# --------------------------------------------------------------------------- +# Unicode build mode (after FetchContent so the vendored benchmark build is not +# compiled with UNICODE/_UNICODE). Windows defaults to ON because mssql-odbc +# exports only wide (W) entry points. +# --------------------------------------------------------------------------- +if(WIN32) + set(_perf_unicode_default ON) +else() + set(_perf_unicode_default OFF) +endif() +option(ODBC_PERF_FORCE_UNICODE "Build perf benchmarks with UNICODE/_UNICODE" ${_perf_unicode_default}) +if(ODBC_PERF_FORCE_UNICODE) + add_compile_definitions(UNICODE _UNICODE) +endif() + +# --------------------------------------------------------------------------- +# perf_lib – shared config + connection helpers +# --------------------------------------------------------------------------- +add_library(perf_lib STATIC lib/perf_fixture.cpp) +target_include_directories(perf_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) +if(ODBC_INCLUDE_DIR) + target_include_directories(perf_lib PUBLIC ${ODBC_INCLUDE_DIR}) +endif() +target_link_libraries(perf_lib PUBLIC benchmark::benchmark ${ODBC_LIBRARIES}) + +function(add_odbc_bench BENCH_NAME) + add_executable(${BENCH_NAME} ${ARGN}) + target_link_libraries(${BENCH_NAME} PRIVATE perf_lib) +endfunction() + +add_odbc_bench(connect_bench benches/connect_bench.cpp) +add_odbc_bench(exec_bench benches/exec_bench.cpp) +add_odbc_bench(fetch_bench benches/fetch_bench.cpp) +add_odbc_bench(datatype_bench benches/datatype_bench.cpp) diff --git a/mssql-odbc/tests/perf/README.md b/mssql-odbc/tests/perf/README.md new file mode 100644 index 00000000..843bc1a6 --- /dev/null +++ b/mssql-odbc/tests/perf/README.md @@ -0,0 +1,183 @@ +# mssql-odbc performance benchmarks + +A [Google Benchmark](https://github.com/google/benchmark) suite that measures the +`mssql-odbc` driver through the ODBC Driver Manager and compares it against +Microsoft's `msodbcsql18`. + +The benchmark binaries are driver-agnostic: they pick a driver purely from the +`Driver={...}` keyword in the connection string. The same binary therefore +produces both sides of the comparison, so the numbers are apples-to-apples — +same client code, same driver manager, same server, same session. + +This implements "Layer 2 — mssql-odbc via ODBC DM" from +[`docs/perf-testing-plan.md`](../../../docs/perf-testing-plan.md). + +## What is measured + +| Binary | Cases | +| --- | --- | +| `connect_bench` | `SQLDriverConnect` + `SQLDisconnect` round trip, statement handle alloc/free | +| `exec_bench` | `SQLExecDirect`, prepare-once/execute-many, prepare+execute each time, parameterized execute | +| `fetch_bench` | Narrow and wide row fetch at 100 / 1 000 / 10 000 rows | +| `datatype_bench` | Per-type retrieval cost: int, bigint, decimal, float, varchar, nvarchar, datetime2, uniqueidentifier, varchar(max) | + +Row-producing cases build a session-scoped `#perf_rows` temp table once per +fixture, so the measured loop is the client-side fetch path rather than server +query planning. + +## Prerequisites + +- SQL Server reachable from this machine (a local instance is fine). +- `msodbcsql18` installed (it ships the driver manager on Windows and provides + the reference numbers). +- CMake 3.20+, Ninja, and a C++17 toolchain. + - **Windows:** Visual Studio 2022 with the "Desktop development with C++" + workload. `run_perf.ps1` locates `vcvars64.bat` via `vswhere` and selects a + toolset that actually ships the CRT import libraries — some installs have a + newer toolset directory without them, which otherwise fails with + `LNK1104: cannot open file 'MSVCRTD.lib'`. + - **Linux/macOS:** `unixodbc-dev` (or `unixodbc` via Homebrew) for `sql.h`. +- The Rust driver registered with the driver manager (see below). + +### One-time driver registration + +The comparison selects drivers by name, so `mssql-odbc` needs its own driver +manager entry pointing at a path the developer can overwrite. This is the only +step that needs elevated rights, and it is done once. + +**Windows** (elevated PowerShell): + +```powershell +New-Item -Path 'C:\odbc-dev' -ItemType Directory -Force +$key = 'HKLM:\Software\ODBC\ODBCINST.INI\mssql-odbc dev' +New-Item -Path $key -Force | Out-Null +Set-ItemProperty -Path $key -Name 'Driver' -Value 'C:\odbc-dev\mssql-odbc-dev.dll' +Set-ItemProperty -Path $key -Name 'Setup' -Value 'C:\odbc-dev\mssql-odbc-dev.dll' +Set-ItemProperty -Path 'HKLM:\Software\ODBC\ODBCINST.INI\ODBC Drivers' ` + -Name 'mssql-odbc dev' -Value 'Installed' +``` + +`run_perf.ps1` then copies each fresh `target\release\msodbcsql18.dll` over +`C:\odbc-dev\mssql-odbc-dev.dll`, so no further registry writes are needed and +the benchmark run itself does not require Administrator. + +**Linux/macOS**, add to `~/.odbcinst.ini`: + +```ini +[mssql-odbc dev] +Description = mssql-odbc development build +Driver = /home//.odbc-dev/libmssql_odbc.so +``` + +### Test database and login + +Any login with permission to create session temp tables works. To create a +dedicated one: + +```sql +CREATE DATABASE odbcperf; +GO +CREATE LOGIN odbcperf WITH PASSWORD = '', CHECK_POLICY = OFF; +ALTER SERVER ROLE sysadmin ADD MEMBER odbcperf; +GO +``` + +Windows integrated authentication also works — omit `-Uid`/`-Pwd` and the +connection string uses `Trusted_Connection=Yes`. + +## Running + +```powershell +# Full comparison: builds the driver, builds the benches, runs both drivers. +.\run_perf.ps1 -Uid odbcperf -Pwd '' + +# One binary, longer measurement window. +.\run_perf.ps1 -Bench fetch_bench -MinTime 2.0 -Repetitions 5 + +# Re-run without rebuilding, restricted to a subset of cases. +.\run_perf.ps1 -SkipBuild -Filter 'BM_Fetch.*' + +# Measure the Rust driver only. +.\run_perf.ps1 -RefDriver '' +``` + +```bash +./run_perf.sh --uid odbcperf --pwd '' +./run_perf.sh --bench fetch_bench --min-time 2.0 +``` + +Raw Google Benchmark JSON lands in `results/`, one file per binary per driver. +The runner prints a comparison of median real times: + +``` +Benchmark mssql-odbc msodbcsql18 Ratio Verdict +--------- ---------- ----------- ----- ------- +BM_Connect_Disconnect 2.57 ms 11.01 ms 0.23 faster +BM_Fetch_NarrowRows/10000 17.34 ms 4.87 ms 3.56 slower +BM_Type_Decimal ERROR 7.49 ms - unsupported +``` + +`Ratio` is `mssql-odbc / msodbcsql18`, so below 1.0 means the Rust driver is +faster. A case the driver cannot service at all is reported as `unsupported` +rather than being dropped, so capability gaps stay visible next to the timings. + +## Environment variables + +The suite reuses the [`tests/e2e`](../e2e/README.md) contract, so a single +environment drives both suites. `run_perf.ps1` sets these for you. + +| Variable | Meaning | +| --- | --- | +| `ODBC_TEST_DRIVER` | Driver name; this is what selects mssql-odbc vs msodbcsql18 | +| `ODBC_TEST_SERVER` | Server address | +| `ODBC_TEST_DATABASE` | Database name | +| `ODBC_TEST_UID` / `ODBC_TEST_PWD` | SQL login; omit both for integrated auth | +| `ODBC_TEST_DSN` | Connect by DSN instead of by driver name | +| `ODBC_TEST_CONNSTR` | Full connection string, overriding everything above | +| `ODBC_TEST_TRUST_CERT` | `TrustServerCertificate` value (default `Yes`) | +| `ODBC_TEST_ENCRYPT` | `Encrypt` value | + +Running a binary directly is sometimes handy while iterating: + +```powershell +$env:ODBC_TEST_DRIVER = 'mssql-odbc dev' +$env:ODBC_TEST_SERVER = 'localhost' +.\build\fetch_bench.exe --benchmark_filter=BM_Fetch_NarrowRows +``` + +## Building manually + +```powershell +cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DODBC_PERF_FORCE_UNICODE=ON +cmake --build build +``` + +`ODBC_PERF_FORCE_UNICODE` (default `ON` on Windows) compiles the benchmark +sources with `UNICODE`/`_UNICODE` so `SQLTCHAR` is `wchar_t`, matching how real +Windows ODBC applications are built. It is applied only to the benchmark +targets, not to the vendored Google Benchmark library. + +## Interpreting results + +- Benchmarks talk to a real server, so a share of every measurement is network + and server time that both drivers pay equally. Ratios are more meaningful than + absolute numbers, and `BM_AllocFree_Stmt` (no server round trip) is the + cleanest signal of pure driver overhead. +- Run with `-Repetitions 3` or more and check the `_cv` rows in the raw output. + A coefficient of variation above roughly 5% means the machine is too noisy to + trust small differences. +- Only compare runs collected on the same machine against the same server. + +## Driver capability notes + +The benchmarks are written against what `mssql-odbc` currently exports, which +constrains the harness: + +- `SQLBindCol` is not exported, so all retrieval goes through `SQLGetData`. +- `SQLGetData` supports `SQL_C_CHAR` and `SQL_C_WCHAR` only; the fixture reads + every column as `SQL_C_CHAR`. +- `SQLBindParameter` supports `SQL_C_CHAR` bound to the `varchar` family only, + so `BM_ParameterizedExecute` passes its parameter as text. + +When these gaps close, the corresponding cases can be widened to also measure +bound-column fetch and native C-type conversion. diff --git a/mssql-odbc/tests/perf/benches/connect_bench.cpp b/mssql-odbc/tests/perf/benches/connect_bench.cpp new file mode 100644 index 00000000..3095a3db --- /dev/null +++ b/mssql-odbc/tests/perf/benches/connect_bench.cpp @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// Connection-establishment cost: handle allocation, login handshake, TLS. + +#include "perf_fixture.h" + +namespace { + +/// Full connect/disconnect cycle — the dominant cost for short-lived processes +/// and the scenario where TLS negotiation and login round-trips show up. +void BM_Connect_Disconnect(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + + perf::Env env; + PERF_REQUIRE(env.ok(), state, "SQLAllocHandle(ENV) failed"); + + { + perf::Conn probe(env); + PERF_REQUIRE(probe.ok(), state, probe.error().c_str()); + } + + for (auto _ : state) { + perf::Conn conn(env); + if (!conn.ok()) { + state.SkipWithError(conn.error().c_str()); + return; + } + benchmark::DoNotOptimize(conn.dbc()); + } + state.SetItemsProcessed(state.iterations()); +} + +/// Handle allocation only — isolates driver-side bookkeeping from network cost. +void BM_AllocFree_Stmt(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + + perf::Env env; + PERF_REQUIRE(env.ok(), state, "SQLAllocHandle(ENV) failed"); + perf::Conn conn(env); + PERF_REQUIRE(conn.ok(), state, conn.error().c_str()); + + for (auto _ : state) { + SQLHSTMT stmt = SQL_NULL_HSTMT; + if (!SQL_SUCCEEDED(SQLAllocHandle(SQL_HANDLE_STMT, conn.dbc(), &stmt))) { + state.SkipWithError("SQLAllocHandle(STMT) failed"); + return; + } + SQLFreeHandle(SQL_HANDLE_STMT, stmt); + } + state.SetItemsProcessed(state.iterations()); +} + +BENCHMARK(BM_Connect_Disconnect)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_AllocFree_Stmt)->Unit(benchmark::kMicrosecond); + +} // namespace + +BENCHMARK_MAIN(); diff --git a/mssql-odbc/tests/perf/benches/datatype_bench.cpp b/mssql-odbc/tests/perf/benches/datatype_bench.cpp new file mode 100644 index 00000000..c2811e37 --- /dev/null +++ b/mssql-odbc/tests/perf/benches/datatype_bench.cpp @@ -0,0 +1,103 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// Per-type decode cost. Each case fetches the same number of rows so the delta +// between cases is attributable to the column type's conversion path. + +#include "perf_fixture.h" + +#include + +namespace { + +constexpr int kRows = 5000; + +struct Session { + perf::Env env; + perf::Conn conn{env}; +}; + +/// Fetch |rows| rows of a single expression. +void TypeBench(benchmark::State& state, const std::string& expr, int rows = kRows) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + + SqlTString sql = perf::ToSqlTStr("SELECT TOP (" + std::to_string(rows) + ") " + expr + + " AS v FROM sys.all_objects a CROSS JOIN " + "sys.all_objects b"); + + int64_t fetched = 0; + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLExecDirect(s.conn.stmt(), sql.data(), SQL_NTS))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + std::string err; + int64_t n = perf::DrainRows(s.conn.stmt(), &err); + if (n < 0) { + state.SkipWithError(err.c_str()); + return; + } + fetched += n; + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(fetched); +} + +void BM_Type_Int(benchmark::State& state) { + TypeBench(state, "CAST(a.object_id AS INT)"); +} + +void BM_Type_BigInt(benchmark::State& state) { + TypeBench(state, "CAST(a.object_id AS BIGINT)"); +} + +void BM_Type_Decimal(benchmark::State& state) { + TypeBench(state, "CAST(a.object_id AS DECIMAL(18,4))"); +} + +void BM_Type_Float(benchmark::State& state) { + TypeBench(state, "CAST(a.object_id AS FLOAT)"); +} + +void BM_Type_Varchar(benchmark::State& state) { + TypeBench(state, "CAST(REPLICATE('x', 100) AS VARCHAR(100))"); +} + +void BM_Type_NVarchar(benchmark::State& state) { + TypeBench(state, "CAST(REPLICATE(N'\u00e9', 100) AS NVARCHAR(100))"); +} + +void BM_Type_DateTime2(benchmark::State& state) { + TypeBench(state, "CAST('2026-01-02 03:04:05.1234567' AS DATETIME2)"); +} + +void BM_Type_Guid(benchmark::State& state) { + TypeBench(state, "CAST('6F9619FF-8B86-D011-B42D-00C04FC964FF' AS UNIQUEIDENTIFIER)"); +} + +/// A value larger than the harness read buffer, so this measures the chunked +/// SQLGetData path rather than a single-shot copy. Row count is deliberately +/// small: each row moves ~20 KB, and a driver with a quadratic chunked-read path +/// makes this case orders of magnitude slower than the others. +void BM_Type_VarcharMax(benchmark::State& state) { + TypeBench(state, "CAST(REPLICATE(CAST('y' AS VARCHAR(MAX)), 20000) AS VARCHAR(MAX))", + 100); +} + +BENCHMARK(BM_Type_Int)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_BigInt)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_Decimal)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_Float)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_Varchar)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_NVarchar)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_DateTime2)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_Guid)->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Type_VarcharMax)->Unit(benchmark::kMillisecond); + +} // namespace + +BENCHMARK_MAIN(); diff --git a/mssql-odbc/tests/perf/benches/exec_bench.cpp b/mssql-odbc/tests/perf/benches/exec_bench.cpp new file mode 100644 index 00000000..a4146ced --- /dev/null +++ b/mssql-odbc/tests/perf/benches/exec_bench.cpp @@ -0,0 +1,164 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// Statement-execution cost on an already-established connection: round-trip +// latency, the prepare/execute split, and parameter marshalling. + +#include "perf_fixture.h" + +#include + +namespace { + +/// Shared connection across all iterations so the measurement excludes login. +struct Session { + perf::Env env; + perf::Conn conn{env}; +}; + +/// Minimal round-trip — measures per-statement overhead with almost no server +/// work and no result-set decoding. +void BM_ExecDirect_SelectOne(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + + SqlTString sql = perf::ToSqlTStr("SELECT 1"); + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLExecDirect(s.conn.stmt(), sql.data(), SQL_NTS))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + std::string err; + if (perf::DrainRows(s.conn.stmt(), &err) < 0) { + state.SkipWithError(err.c_str()); + return; + } + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(state.iterations()); +} + +/// SQLPrepare once, SQLExecute repeatedly — the path a well-written application +/// uses for a hot statement. Isolates execute cost from statement compilation. +void BM_PreparedExecute(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + + SqlTString sql = perf::ToSqlTStr("SELECT 1"); + PERF_REQUIRE(SQL_SUCCEEDED(SQLPrepare(s.conn.stmt(), sql.data(), SQL_NTS)), state, + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLExecute(s.conn.stmt()))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + std::string err; + if (perf::DrainRows(s.conn.stmt(), &err) < 0) { + state.SkipWithError(err.c_str()); + return; + } + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(state.iterations()); +} + +/// Prepare + execute every iteration — shows the cost a driver pays when an +/// application does not cache prepared handles. +void BM_PrepareExecute_EachTime(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + + SqlTString sql = perf::ToSqlTStr("SELECT 1"); + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLPrepare(s.conn.stmt(), sql.data(), SQL_NTS)) || + !SQL_SUCCEEDED(SQLExecute(s.conn.stmt()))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + std::string err; + if (perf::DrainRows(s.conn.stmt(), &err) < 0) { + state.SkipWithError(err.c_str()); + return; + } + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(state.iterations()); +} + +/// Parameterised execute — adds SQLBindParameter marshalling to the round trip. +/// Bound as SQL_C_CHAR → SQL_VARCHAR because that is the conversion pair +/// mssql-odbc currently implements; msodbcsql accepts it too, so both drivers +/// are measured on identical work. +void BM_ParameterizedExecute(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + + char value[32] = "perf-parameter-value"; + SQLLEN ind = SQL_NTS; + PERF_REQUIRE(SQL_SUCCEEDED(SQLBindParameter(s.conn.stmt(), 1, SQL_PARAM_INPUT, + SQL_C_CHAR, SQL_VARCHAR, sizeof(value) - 1, + 0, value, sizeof(value), &ind)), + state, perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + + SqlTString sql = perf::ToSqlTStr("SELECT ?"); + PERF_REQUIRE(SQL_SUCCEEDED(SQLPrepare(s.conn.stmt(), sql.data(), SQL_NTS)), state, + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLExecute(s.conn.stmt()))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + std::string err; + if (perf::DrainRows(s.conn.stmt(), &err) < 0) { + state.SkipWithError(err.c_str()); + return; + } + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(state.iterations()); +} + +/// SQLExecDirect without draining the result set — isolates the request/response +/// round trip from result-set consumption. Paired with BM_ExecDirect_SelectOne, +/// the difference is the fetch-and-close cost for a single trivial row. +void BM_ExecDirect_NoDrain(benchmark::State& state) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + + SqlTString sql = perf::ToSqlTStr("SELECT 1"); + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLExecDirect(s.conn.stmt(), sql.data(), SQL_NTS))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(state.iterations()); +} + +BENCHMARK(BM_ExecDirect_SelectOne)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_ExecDirect_NoDrain)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_PreparedExecute)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_PrepareExecute_EachTime)->Unit(benchmark::kMicrosecond); +BENCHMARK(BM_ParameterizedExecute)->Unit(benchmark::kMicrosecond); + +} // namespace + +BENCHMARK_MAIN(); diff --git a/mssql-odbc/tests/perf/benches/fetch_bench.cpp b/mssql-odbc/tests/perf/benches/fetch_bench.cpp new file mode 100644 index 00000000..9c0a1a92 --- /dev/null +++ b/mssql-odbc/tests/perf/benches/fetch_bench.cpp @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// Result-set retrieval throughput: token-stream decoding, row materialisation, +// and SQLGetData conversion cost across row counts and row widths. + +#include "perf_fixture.h" + +#include + +namespace { + +/// Rows are served from a session temp table populated once, so the measurement +/// reflects fetch cost rather than the server's row-generation cost. +constexpr int kSourceRows = 20000; + +struct Session { + perf::Env env; + perf::Conn conn{env}; + + bool Prepare() { + return conn.ok() && + conn.Exec("SELECT TOP (" + std::to_string(kSourceRows) + + ") ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS id, " + "CAST('row payload text for fetch benchmarking' AS VARCHAR(64)) " + "AS payload " + "INTO #perf_rows " + "FROM sys.all_objects a CROSS JOIN sys.all_objects b"); + } +}; + +void FetchRange(benchmark::State& state, const std::string& columns, + bool get_data = true) { + PERF_REQUIRE(perf::Config::Instance().HasConnection(), state, + "no connection configured (set ODBC_TEST_SERVER)"); + Session s; + PERF_REQUIRE(s.conn.ok(), state, s.conn.error().c_str()); + PERF_REQUIRE(s.Prepare(), state, s.conn.error().c_str()); + + const int64_t rows = state.range(0); + SqlTString sql = perf::ToSqlTStr("SELECT TOP (" + std::to_string(rows) + ") " + + columns + " FROM #perf_rows"); + + int64_t fetched = 0; + for (auto _ : state) { + if (!SQL_SUCCEEDED(SQLExecDirect(s.conn.stmt(), sql.data(), SQL_NTS))) { + state.SkipWithError( + perf::DiagText(SQL_HANDLE_STMT, s.conn.stmt()).c_str()); + return; + } + std::string err; + int64_t n = get_data ? perf::DrainRows(s.conn.stmt(), &err) + : perf::DrainRowsNoGetData(s.conn.stmt(), &err); + if (n < 0) { + state.SkipWithError(err.c_str()); + return; + } + fetched += n; + perf::CloseCursor(s.conn.stmt()); + } + state.SetItemsProcessed(fetched); +} + +/// Narrow rows — dominated by per-row protocol and cursor overhead. +void BM_Fetch_NarrowRows(benchmark::State& state) { + FetchRange(state, "id"); +} + +/// Wide rows — adds per-column SQLGetData conversion to the per-row cost. +void BM_Fetch_WideRows(benchmark::State& state) { + FetchRange(state, + "id, payload, payload AS p2, payload AS p3, payload AS p4, " + "payload AS p5, payload AS p6, payload AS p7"); +} + +/// Cursor-only advance — isolates SQLFetch cost from SQLGetData conversion. +void BM_Fetch_RowsOnly(benchmark::State& state) { + FetchRange(state, "id", /*get_data=*/false); +} + +BENCHMARK(BM_Fetch_NarrowRows) + ->Arg(100) + ->Arg(1000) + ->Arg(10000) + ->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Fetch_WideRows) + ->Arg(100) + ->Arg(1000) + ->Arg(10000) + ->Unit(benchmark::kMillisecond); +BENCHMARK(BM_Fetch_RowsOnly)->Arg(10000)->Unit(benchmark::kMillisecond); + +} // namespace + +BENCHMARK_MAIN(); diff --git a/mssql-odbc/tests/perf/include/perf_fixture.h b/mssql-odbc/tests/perf/include/perf_fixture.h new file mode 100644 index 00000000..87a32f40 --- /dev/null +++ b/mssql-odbc/tests/perf/include/perf_fixture.h @@ -0,0 +1,152 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. +// +// perf_fixture.h – Shared plumbing for the ODBC performance benchmarks. +// +// The benchmarks drive the driver through the ODBC Driver Manager, so the exact +// same binary measures mssql-odbc and msodbcsql18. Which driver is exercised is +// decided purely by the `Driver={...}` name in the connection string +// (ODBC_TEST_DRIVER), so a comparison run never has to touch the registry. + +#pragma once + +#ifdef _WIN32 +#include +#endif + +#include +#include + +#include + +#include +#include + +#ifndef SQL_OV_ODBC3_80 +#define SQL_OV_ODBC3_80 380UL +#endif + +/// SQLTCHAR-based string type for ODBC API calls. +using SqlTString = std::basic_string; + +namespace perf { + +SqlTString ToSqlTStr(const std::string& s); +std::string ToNarrow(const SqlTString& s); + +/// Concatenated SQLSTATE + message text for every diagnostic record on |handle|. +std::string DiagText(SQLSMALLINT handle_type, SQLHANDLE handle); + +// --------------------------------------------------------------------------- +// Config – connection info from environment variables +// --------------------------------------------------------------------------- +// Shares the ODBC_TEST_* contract with tests/e2e so one environment drives both +// suites. ODBC_TEST_DRIVER is what selects the driver under measurement. +// +// ODBC_TEST_SERVER – server hostname (required unless CONNSTR/DSN) +// ODBC_TEST_DATABASE – database name (default: tempdb) +// ODBC_TEST_UID/PWD – SQL login (omit for integrated auth) +// ODBC_TEST_DRIVER – driver name (default: ODBC Driver 18 for SQL Server) +// ODBC_TEST_DSN – DSN name (overrides driver/server) +// ODBC_TEST_CONNSTR – full connection string (overrides everything above) +// ODBC_TEST_TRUST_CERT – TrustServerCertificate (default: Yes) +// ODBC_TEST_ENCRYPT – Encrypt value (default: driver default) +// --------------------------------------------------------------------------- +class Config { +public: + static const Config& Instance(); + + const std::string& Driver() const { return driver_; } + bool HasConnection() const { + return !connstr_.empty() || !dsn_.empty() || !server_.empty(); + } + + /// Connection string assembled from the environment. + SqlTString ConnectionString() const; + +private: + Config(); + static std::string GetEnv(const char* name, const char* fallback = ""); + + std::string dsn_; + std::string server_; + std::string database_; + std::string uid_; + std::string pwd_; + std::string driver_; + std::string connstr_; + std::string trust_cert_; + std::string encrypt_; +}; + +// --------------------------------------------------------------------------- +// RAII handle wrappers +// --------------------------------------------------------------------------- + +/// ODBC environment handle set to ODBC 3.80. +class Env { +public: + Env(); + ~Env(); + Env(const Env&) = delete; + Env& operator=(const Env&) = delete; + + SQLHENV get() const { return env_; } + bool ok() const { return env_ != SQL_NULL_HENV; } + +private: + SQLHENV env_ = SQL_NULL_HENV; +}; + +/// A connected HDBC plus one HSTMT, both released on destruction. +class Conn { +public: + explicit Conn(const Env& env); + ~Conn(); + Conn(const Conn&) = delete; + Conn& operator=(const Conn&) = delete; + + SQLHDBC dbc() const { return dbc_; } + SQLHSTMT stmt() const { return stmt_; } + bool ok() const { return connected_ && stmt_ != SQL_NULL_HSTMT; } + + /// Diagnostic text explaining why ok() is false, or why the last Exec failed. + const std::string& error() const { return error_; } + + /// Run a statement, draining every row and result set it produces. + bool Exec(const std::string& sql); + +private: + SQLHDBC dbc_ = SQL_NULL_HDBC; + SQLHSTMT stmt_ = SQL_NULL_HSTMT; + bool connected_ = false; + std::string error_; +}; + +/// Read every row/column of the current result set via SQLGetData and return the +/// number of rows consumed, or -1 on error. +/// +/// SQLGetData is used rather than SQLBindCol because mssql-odbc does not export +/// SQLBindCol yet; keeping both drivers on the same retrieval path is what makes +/// the comparison meaningful. +int64_t DrainRows(SQLHSTMT stmt, std::string* error); + +/// Advances the cursor to exhaustion without retrieving any column data. +/// +/// Isolates the per-row cost of the cursor and token-stream path from the +/// per-column `SQLGetData` conversion cost measured by `DrainRows`. +int64_t DrainRowsNoGetData(SQLHSTMT stmt, std::string* error); + +/// Close the cursor so the statement can be reused for the next iteration. +void CloseCursor(SQLHSTMT stmt); + +/// Bail out of a benchmark with a diagnostic instead of reporting bogus timings. +#define PERF_REQUIRE(cond, state, msg) \ + do { \ + if (!(cond)) { \ + (state).SkipWithError(msg); \ + return; \ + } \ + } while (0) + +} // namespace perf diff --git a/mssql-odbc/tests/perf/lib/perf_fixture.cpp b/mssql-odbc/tests/perf/lib/perf_fixture.cpp new file mode 100644 index 00000000..cd64597c --- /dev/null +++ b/mssql-odbc/tests/perf/lib/perf_fixture.cpp @@ -0,0 +1,284 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#include "perf_fixture.h" + +#include +#include + +namespace perf { + +SqlTString ToSqlTStr(const std::string& s) { + return SqlTString(s.begin(), s.end()); +} + +std::string ToNarrow(const SqlTString& s) { + return std::string(s.begin(), s.end()); +} + +std::string DiagText(SQLSMALLINT handle_type, SQLHANDLE handle) { + SQLTCHAR state[8] = {}; + SQLINTEGER native = 0; + SQLTCHAR msg[1024] = {}; + SQLSMALLINT msg_len = 0; + std::ostringstream oss; + bool found = false; + + for (SQLSMALLINT rec = 1;; rec++) { + SQLRETURN rc = SQLGetDiagRec( + handle_type, handle, rec, state, &native, msg, + static_cast(sizeof(msg) / sizeof(SQLTCHAR)), &msg_len); + if (!SQL_SUCCEEDED(rc)) { + break; + } + if (found) { + oss << " | "; + } + oss << "[" << ToNarrow(SqlTString(state)) << "] " << ToNarrow(SqlTString(msg)); + found = true; + } + return found ? oss.str() : "(no diagnostic)"; +} + +// --------------------------------------------------------------------------- +// Config +// --------------------------------------------------------------------------- + +const Config& Config::Instance() { + static Config cfg; + return cfg; +} + +Config::Config() + : dsn_(GetEnv("ODBC_TEST_DSN")), + server_(GetEnv("ODBC_TEST_SERVER")), + database_(GetEnv("ODBC_TEST_DATABASE", "tempdb")), + uid_(GetEnv("ODBC_TEST_UID")), + pwd_(GetEnv("ODBC_TEST_PWD")), + driver_(GetEnv("ODBC_TEST_DRIVER", "ODBC Driver 18 for SQL Server")), + connstr_(GetEnv("ODBC_TEST_CONNSTR")), + trust_cert_(GetEnv("ODBC_TEST_TRUST_CERT", "Yes")), + encrypt_(GetEnv("ODBC_TEST_ENCRYPT")) {} + +std::string Config::GetEnv(const char* name, const char* fallback) { +#ifdef _WIN32 + char* buf = nullptr; + size_t len = 0; + if (_dupenv_s(&buf, &len, name) == 0 && buf != nullptr) { + std::string val(buf); + free(buf); + if (!val.empty()) { + return val; + } + } + return fallback ? fallback : ""; +#else + const char* val = std::getenv(name); + return (val && val[0]) ? std::string(val) + : (fallback ? std::string(fallback) : std::string()); +#endif +} + +SqlTString Config::ConnectionString() const { + if (!connstr_.empty()) { + return ToSqlTStr(connstr_); + } + + std::ostringstream cs; + if (!dsn_.empty()) { + cs << "DSN=" << dsn_ << ";"; + } else { + cs << "Driver={" << driver_ << "};Server=" << server_ << ";"; + } + cs << "Database=" << database_ << ";TrustServerCertificate=" << trust_cert_ << ";"; + if (!encrypt_.empty()) { + cs << "Encrypt=" << encrypt_ << ";"; + } + if (!uid_.empty()) { + cs << "Uid=" << uid_ << ";Pwd=" << pwd_ << ";"; + } else { + cs << "Trusted_Connection=Yes;"; + } + return ToSqlTStr(cs.str()); +} + +// --------------------------------------------------------------------------- +// Env +// --------------------------------------------------------------------------- + +Env::Env() { + if (!SQL_SUCCEEDED(SQLAllocHandle(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &env_))) { + env_ = SQL_NULL_HENV; + return; + } + SQLRETURN rc = SQLSetEnvAttr(env_, SQL_ATTR_ODBC_VERSION, + reinterpret_cast(SQL_OV_ODBC3_80), 0); + if (!SQL_SUCCEEDED(rc)) { + SQLFreeHandle(SQL_HANDLE_ENV, env_); + env_ = SQL_NULL_HENV; + } +} + +Env::~Env() { + if (env_ != SQL_NULL_HENV) { + SQLFreeHandle(SQL_HANDLE_ENV, env_); + } +} + +// --------------------------------------------------------------------------- +// Conn +// --------------------------------------------------------------------------- + +Conn::Conn(const Env& env) { + if (!env.ok()) { + error_ = "environment handle allocation failed"; + return; + } + if (!SQL_SUCCEEDED(SQLAllocHandle(SQL_HANDLE_DBC, env.get(), &dbc_))) { + error_ = "SQLAllocHandle(DBC) failed: " + DiagText(SQL_HANDLE_ENV, env.get()); + dbc_ = SQL_NULL_HDBC; + return; + } + + SqlTString conn_str = Config::Instance().ConnectionString(); + SQLRETURN rc = SQLDriverConnect(dbc_, nullptr, conn_str.data(), + static_cast(conn_str.size()), nullptr, 0, + nullptr, SQL_DRIVER_NOPROMPT); + if (!SQL_SUCCEEDED(rc)) { + error_ = "SQLDriverConnect failed: " + DiagText(SQL_HANDLE_DBC, dbc_); + return; + } + connected_ = true; + + if (!SQL_SUCCEEDED(SQLAllocHandle(SQL_HANDLE_STMT, dbc_, &stmt_))) { + error_ = "SQLAllocHandle(STMT) failed: " + DiagText(SQL_HANDLE_DBC, dbc_); + stmt_ = SQL_NULL_HSTMT; + } +} + +Conn::~Conn() { + if (stmt_ != SQL_NULL_HSTMT) { + SQLFreeHandle(SQL_HANDLE_STMT, stmt_); + } + if (connected_) { + SQLDisconnect(dbc_); + } + if (dbc_ != SQL_NULL_HDBC) { + SQLFreeHandle(SQL_HANDLE_DBC, dbc_); + } +} + +bool Conn::Exec(const std::string& sql) { + if (stmt_ == SQL_NULL_HSTMT) { + error_ = "no statement handle"; + return false; + } + SqlTString text = ToSqlTStr(sql); + SQLRETURN rc = SQLExecDirect(stmt_, text.data(), SQL_NTS); + if (!SQL_SUCCEEDED(rc)) { + error_ = "SQLExecDirect(" + sql + ") failed: " + DiagText(SQL_HANDLE_STMT, stmt_); + CloseCursor(stmt_); + return false; + } + do { + std::string drain_err; + if (DrainRows(stmt_, &drain_err) < 0) { + error_ = drain_err; + CloseCursor(stmt_); + return false; + } + } while (SQLMoreResults(stmt_) == SQL_SUCCESS); + CloseCursor(stmt_); + return true; +} + +// --------------------------------------------------------------------------- +// Row draining +// --------------------------------------------------------------------------- + +int64_t DrainRows(SQLHSTMT stmt, std::string* error) { + // Upper bound on SQLGetData calls for a single column, so a driver that never + // terminates a truncation sequence fails the run instead of hanging it. + constexpr int kMaxChunksPerColumn = 100000; + + SQLSMALLINT col_count = 0; + if (!SQL_SUCCEEDED(SQLNumResultCols(stmt, &col_count))) { + if (error) { + *error = "SQLNumResultCols failed: " + DiagText(SQL_HANDLE_STMT, stmt); + } + return -1; + } + if (col_count == 0) { + return 0; + } + + // Sized to hold a full inline (non-LOB) column in one SQLGetData call; larger + // values loop below, which is exactly what a real application does. + char buf[8192]; + int64_t rows = 0; + + for (;;) { + SQLRETURN rc = SQLFetch(stmt); + if (rc == SQL_NO_DATA) { + break; + } + if (!SQL_SUCCEEDED(rc)) { + if (error) { + *error = "SQLFetch failed: " + DiagText(SQL_HANDLE_STMT, stmt); + } + return -1; + } + for (SQLSMALLINT col = 1; col <= col_count; col++) { + SQLLEN ind = 0; + SQLRETURN grc = SQL_SUCCESS; + // Standard truncation idiom: a long value comes back as repeated + // SQL_SUCCESS_WITH_INFO chunks terminated by SQL_SUCCESS. The cap + // guards against a driver that never terminates the sequence. + for (int chunk = 0; chunk < kMaxChunksPerColumn; chunk++) { + grc = SQLGetData(stmt, static_cast(col), SQL_C_CHAR, buf, + static_cast(sizeof(buf)), &ind); + if (grc == SQL_NO_DATA) { + break; + } + if (!SQL_SUCCEEDED(grc)) { + if (error) { + *error = "SQLGetData failed: " + DiagText(SQL_HANDLE_STMT, stmt); + } + return -1; + } + benchmark::DoNotOptimize(buf); + if (grc != SQL_SUCCESS_WITH_INFO) { + break; + } + } + } + rows++; + } + return rows; +} + +int64_t DrainRowsNoGetData(SQLHSTMT stmt, std::string* error) { + int64_t rows = 0; + for (;;) { + SQLRETURN rc = SQLFetch(stmt); + if (rc == SQL_NO_DATA) { + break; + } + if (!SQL_SUCCEEDED(rc)) { + if (error) { + *error = "SQLFetch failed: " + DiagText(SQL_HANDLE_STMT, stmt); + } + return -1; + } + rows++; + } + return rows; +} + +void CloseCursor(SQLHSTMT stmt) { + // SQLCloseCursor returns 24000 when no cursor is open (e.g. after a DDL + // statement); that is expected, so the return code is deliberately ignored. + SQLCloseCursor(stmt); +} + +} // namespace perf diff --git a/mssql-odbc/tests/perf/run_perf.ps1 b/mssql-odbc/tests/perf/run_perf.ps1 new file mode 100644 index 00000000..950f137d --- /dev/null +++ b/mssql-odbc/tests/perf/run_perf.ps1 @@ -0,0 +1,277 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +# Build the Rust ODBC driver and the C++ Google Benchmark suite, then run the +# same benchmark binaries against mssql-odbc and msodbcsql18 and print a +# side-by-side comparison. +# +# Unlike run_e2e.ps1, this script does NOT require Administrator. Both drivers +# are selected by the `Driver={...}` name in the connection string, so a +# comparison run never edits the registry — it only refreshes the DLL that the +# already-registered Rust driver entry points at. +# +# Prerequisites (one-time, requires Administrator): +# The Rust driver must be registered under -RustDriver pointing at a path this +# user can overwrite, e.g. +# HKLM\Software\ODBC\ODBCINST.INI\mssql-odbc dev +# Driver = C:\odbc-dev\mssql-odbc-dev.dll +# Setup = C:\odbc-dev\mssql-odbc-dev.dll +# HKLM\Software\ODBC\ODBCINST.INI\ODBC Drivers +# mssql-odbc dev = Installed +# +# Usage: +# .\run_perf.ps1 +# .\run_perf.ps1 -Bench fetch_bench -MinTime 2.0 +# .\run_perf.ps1 -SkipBuild -Filter 'BM_Fetch.*' + +param( + # Driver name (as registered with the Driver Manager) for the Rust driver. + [string]$RustDriver = 'mssql-odbc dev', + # Reference driver to compare against. Empty string runs the Rust driver only. + [string]$RefDriver = 'ODBC Driver 18 for SQL Server', + [string]$Server = 'localhost', + [string]$Database = 'odbcperf', + [string]$Uid = '', + [string]$Pwd = '', + # Restrict to one benchmark binary (connect_bench|exec_bench|fetch_bench|datatype_bench). + [string]$Bench = '', + # Google Benchmark --benchmark_filter regex. + [string]$Filter = '', + # Seconds of measurement per benchmark case. + [double]$MinTime = 1.0, + # Repetitions per case; >1 makes the median/stddev columns meaningful. + [int]$Repetitions = 3, + [switch]$SkipBuild, + [string]$ResultsDir = '' +) + +$ErrorActionPreference = 'Stop' + +$ScriptDir = Split-Path -Parent $MyInvocation.MyCommand.Definition +$OdbcCrateDir = Resolve-Path (Join-Path $ScriptDir '..\..') +$WorkspaceDir = Resolve-Path (Join-Path $OdbcCrateDir '..') +$BuildDir = Join-Path $ScriptDir 'build' + +if (-not $ResultsDir) { $ResultsDir = Join-Path $ScriptDir 'results' } +New-Item -ItemType Directory -Force -Path $ResultsDir | Out-Null + +$AllBenches = @('connect_bench', 'exec_bench', 'fetch_bench', 'datatype_bench') +$Benches = if ($Bench) { @($Bench) } else { $AllBenches } + +function Write-Section([string]$msg) { + Write-Host '' + Write-Host "=== $msg ===" -ForegroundColor Cyan +} + +# --------------------------------------------------------------------------- +# Resolve where the registered Rust driver DLL lives +# --------------------------------------------------------------------------- +function Get-RegisteredDriverPath([string]$name) { + $key = "HKLM:\Software\ODBC\ODBCINST.INI\$name" + if (-not (Test-Path $key)) { return $null } + $v = Get-ItemProperty -Path $key -Name 'Driver' -ErrorAction SilentlyContinue + if ($null -eq $v) { return $null } + return $v.Driver +} + +$RustDriverPath = Get-RegisteredDriverPath $RustDriver +if (-not $RustDriverPath) { + throw "Driver '$RustDriver' is not registered. See the header of this script for the one-time registration steps." +} + +if ($RefDriver) { + $RefDriverPath = Get-RegisteredDriverPath $RefDriver + if (-not $RefDriverPath) { + throw "Reference driver '$RefDriver' is not registered." + } +} + +# --------------------------------------------------------------------------- +# Locate the MSVC toolchain. CMake's Visual Studio generator cannot always +# discover the compiler on a plain shell, so the C++ build runs inside a +# vcvars64 environment with an explicit toolset that has the CRT libraries. +# --------------------------------------------------------------------------- +function Get-VcVarsPath { + $vswhere = "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" + if (-not (Test-Path $vswhere)) { throw 'vswhere.exe not found; install Visual Studio 2022 with the C++ workload.' } + $root = & $vswhere -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath + if (-not $root) { throw 'No Visual Studio installation with the C++ toolset was found.' } + $vcvars = Join-Path $root 'VC\Auxiliary\Build\vcvars64.bat' + if (-not (Test-Path $vcvars)) { throw "vcvars64.bat not found under $root" } + return $vcvars +} + +# vcvars64 defaults to the newest *installed* toolset, which is not necessarily +# one with the CRT import libraries present. Pick the highest toolset that +# actually has msvcrtd.lib so the CMake compiler probe links. +function Get-VcToolsetVersion([string]$vcvars) { + $msvcRoot = Join-Path (Split-Path -Parent (Split-Path -Parent (Split-Path -Parent $vcvars))) 'Tools\MSVC' + if (-not (Test-Path $msvcRoot)) { return '' } + $candidate = Get-ChildItem $msvcRoot | + Where-Object { Test-Path (Join-Path $_.FullName 'lib\x64\msvcrtd.lib') } | + Sort-Object Name -Descending | + Select-Object -First 1 + if (-not $candidate) { return '' } + # vcvars expects a major.minor version (e.g. 14.44), not the full toolset number. + $parts = $candidate.Name.Split('.') + return "$($parts[0]).$($parts[1])" +} + +function Invoke-InVcVars([string]$command) { + $vcvars = Get-VcVarsPath + $ver = Get-VcToolsetVersion $vcvars + $verArg = if ($ver) { " -vcvars_ver=$ver" } else { '' } + & $env:ComSpec /c "call `"$vcvars`"$verArg >nul && $command" + if ($LASTEXITCODE -ne 0) { throw "Command failed under vcvars: $command" } +} + +# --------------------------------------------------------------------------- +# Build +# --------------------------------------------------------------------------- +if (-not $SkipBuild) { + Write-Section 'Building Rust driver (release)' + Push-Location $OdbcCrateDir + try { + cargo build --release + if ($LASTEXITCODE -ne 0) { throw 'cargo build failed' } + } finally { + Pop-Location + } + + $BuiltDll = Join-Path $WorkspaceDir 'target\release\msodbcsql18.dll' + if (-not (Test-Path $BuiltDll)) { throw "Built driver not found at $BuiltDll" } + + Write-Host "[ DRIVER ] Copying $BuiltDll -> $RustDriverPath" + Copy-Item -Path $BuiltDll -Destination $RustDriverPath -Force + + Write-Section 'Building C++ benchmarks (Release)' + Invoke-InVcVars "cd /d `"$ScriptDir`" && cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DODBC_PERF_FORCE_UNICODE=ON" + Invoke-InVcVars "cd /d `"$ScriptDir`" && cmake --build build" +} + +# --------------------------------------------------------------------------- +# Run one benchmark binary against one driver +# --------------------------------------------------------------------------- +function Invoke-Bench([string]$benchName, [string]$driverName, [string]$label) { + $exe = Join-Path $BuildDir "Release\$benchName.exe" + if (-not (Test-Path $exe)) { $exe = Join-Path $BuildDir "$benchName.exe" } + if (-not (Test-Path $exe)) { throw "Benchmark binary not found: $benchName" } + + $out = Join-Path $ResultsDir "$benchName.$label.json" + + $env:ODBC_TEST_DRIVER = $driverName + $env:ODBC_TEST_SERVER = $Server + $env:ODBC_TEST_DATABASE = $Database + $env:ODBC_TEST_UID = $Uid + $env:ODBC_TEST_PWD = $Pwd + + $benchArgs = @( + "--benchmark_out=$out", + '--benchmark_out_format=json', + "--benchmark_min_time=$($MinTime)s", + "--benchmark_repetitions=$Repetitions", + '--benchmark_report_aggregates_only=true' + ) + if ($Filter) { $benchArgs += "--benchmark_filter=$Filter" } + + Write-Host "[ RUN ] $benchName driver='$driverName'" + & $exe @benchArgs | Out-Host + if ($LASTEXITCODE -ne 0) { + Write-Host "[ FAIL ] $benchName exited $LASTEXITCODE for '$driverName'" -ForegroundColor Red + } + return $out +} + +# --------------------------------------------------------------------------- +# Compare two result files +# --------------------------------------------------------------------------- +function Compare-Results([string]$rustJson, [string]$refJson) { + if (-not (Test-Path $rustJson) -or -not (Test-Path $refJson)) { return @() } + + # Keep only median aggregates so a single row per case is compared. A case + # that failed is recorded with a null time so it shows up as an error row + # rather than silently disappearing from the table. + function Read-Medians([string]$path) { + $map = @{} + $data = Get-Content $path -Raw | ConvertFrom-Json + foreach ($b in $data.benchmarks) { + $name = $b.name -replace '_(median|mean|stddev|cv)$', '' + if ($b.PSObject.Properties.Name -contains 'error_occurred' -and $b.error_occurred) { + $map[$name] = $null + continue + } + if ($b.run_type -eq 'aggregate' -and $b.aggregate_name -ne 'median') { continue } + if ($map.ContainsKey($name) -and $null -eq $map[$name]) { continue } + $map[$name] = $b.real_time * (Get-TimeScale $b.time_unit) + } + return $map + } + + $rust = Read-Medians $rustJson + $ref = Read-Medians $refJson + + $rows = @() + foreach ($name in $rust.Keys) { + if (-not $ref.ContainsKey($name)) { continue } + $r = $rust[$name] + $m = $ref[$name] + $rows += [pscustomobject]@{ + Benchmark = $name + 'mssql-odbc' = if ($null -eq $r) { 'ERROR' } else { Format-Duration $r } + 'msodbcsql18' = if ($null -eq $m) { 'ERROR' } else { Format-Duration $m } + Ratio = if ($null -eq $r -or $null -eq $m -or $m -le 0) { '-' } else { [math]::Round($r / $m, 2) } + Verdict = if ($null -eq $r) { 'unsupported' } + elseif ($null -eq $m -or $m -le 0) { 'n/a' } + elseif ($r -lt $m * 0.95) { 'faster' } + elseif ($r -gt $m * 1.05) { 'slower' } + else { 'parity' } + } + } + return $rows +} + +# Normalize Google Benchmark time units to nanoseconds. +function Get-TimeScale([string]$unit) { + switch ($unit) { + 'ns' { 1 } + 'us' { 1000 } + 'ms' { 1000000 } + 's' { 1000000000 } + default { 1 } + } +} + +function Format-Duration([double]$ns) { + if ($ns -ge 1000000) { return "{0:N2} ms" -f ($ns / 1000000) } + if ($ns -ge 1000) { return "{0:N2} us" -f ($ns / 1000) } + return "{0:N2} ns" -f $ns +} + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +Write-Section 'Configuration' +Write-Host " Rust driver : $RustDriver ($RustDriverPath)" +if ($RefDriver) { Write-Host " Reference : $RefDriver ($RefDriverPath)" } +Write-Host " Server : $Server / $Database" +Write-Host " Auth : $(if ($Uid) { "SQL login '$Uid'" } else { 'Windows integrated' })" +Write-Host " Benchmarks : $($Benches -join ', ')" + +$allRows = @() +foreach ($b in $Benches) { + Write-Section "Benchmark: $b" + $rustJson = Invoke-Bench $b $RustDriver 'mssql-odbc' + if ($RefDriver) { + $refJson = Invoke-Bench $b $RefDriver 'msodbcsql18' + $allRows += Compare-Results $rustJson $refJson + } +} + +if ($allRows.Count -gt 0) { + Write-Section 'Comparison (median real time; Ratio < 1.0 = mssql-odbc faster)' + $allRows | Sort-Object { if ($_.Ratio -eq '-') { [double]::MaxValue } else { [double]$_.Ratio } } | + Format-Table -AutoSize | Out-String | ForEach-Object { Write-Output $_ } +} + +Write-Host '' +Write-Host "Raw JSON results in: $ResultsDir" diff --git a/mssql-odbc/tests/perf/run_perf.sh b/mssql-odbc/tests/perf/run_perf.sh new file mode 100755 index 00000000..93d26e3c --- /dev/null +++ b/mssql-odbc/tests/perf/run_perf.sh @@ -0,0 +1,269 @@ +#!/bin/bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +# Build the Rust ODBC driver and the C++ Google Benchmark suite, then run the +# same benchmark binaries against mssql-odbc and msodbcsql18 and print a +# side-by-side comparison. +# +# For Unix-like platforms (Linux, macOS) that use unixODBC. +# For Windows, see run_perf.ps1. +# +# Both drivers are selected by the `Driver={...}` keyword in the connection +# string, so the same binary produces both sides of the comparison. +# +# Prerequisites: +# The Rust driver must be registered under --rust-driver in odbcinst.ini at a +# path this user can overwrite, e.g. in ~/.odbcinst.ini: +# [mssql-odbc dev] +# Driver = /home//.odbc-dev/libmsodbcsql18.so +# +# Usage: +# ./run_perf.sh [--uid=USER] [--pwd=PASSWORD] [--server=HOST] [--database=DB] +# [--bench=NAME] [--filter=REGEX] [--min-time=SECONDS] +# [--repetitions=N] [--skip-build] [--rust-driver=NAME] +# [--ref-driver=NAME] [--results-dir=PATH] +# +# Examples: +# ./run_perf.sh --uid=odbcperf --pwd='' +# ./run_perf.sh --bench=fetch_bench --min-time=2.0 --repetitions=5 +# ./run_perf.sh --skip-build --filter='BM_Fetch.*' +# ./run_perf.sh --ref-driver='' # measure mssql-odbc only + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ODBC_CRATE_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)" +WORKSPACE_DIR="$(cd "$ODBC_CRATE_DIR/.." && pwd)" +BUILD_DIR="$SCRIPT_DIR/build" + +RUST_DRIVER="mssql-odbc dev" +REF_DRIVER="ODBC Driver 18 for SQL Server" +SERVER="localhost" +DATABASE="odbcperf" +UID_ARG="" +PWD_ARG="" +BENCH="" +FILTER="" +MIN_TIME="1.0" +REPETITIONS="3" +SKIP_BUILD=0 +RESULTS_DIR="$SCRIPT_DIR/results" + +ALL_BENCHES=(connect_bench exec_bench fetch_bench datatype_bench) + +for arg in "$@"; do + case "$arg" in + --rust-driver=*) RUST_DRIVER="${arg#*=}" ;; + --ref-driver=*) REF_DRIVER="${arg#*=}" ;; + --server=*) SERVER="${arg#*=}" ;; + --database=*) DATABASE="${arg#*=}" ;; + --uid=*) UID_ARG="${arg#*=}" ;; + --pwd=*) PWD_ARG="${arg#*=}" ;; + --bench=*) BENCH="${arg#*=}" ;; + --filter=*) FILTER="${arg#*=}" ;; + --min-time=*) MIN_TIME="${arg#*=}" ;; + --repetitions=*) REPETITIONS="${arg#*=}" ;; + --results-dir=*) RESULTS_DIR="${arg#*=}" ;; + --skip-build) SKIP_BUILD=1 ;; + -h|--help) sed -n '2,32p' "${BASH_SOURCE[0]}"; exit 0 ;; + *) echo "Unknown argument: $arg" >&2; exit 2 ;; + esac +done + +if [[ -n "$BENCH" ]]; then + BENCHES=("$BENCH") +else + BENCHES=("${ALL_BENCHES[@]}") +fi + +mkdir -p "$RESULTS_DIR" + +section() { printf '\n=== %s ===\n' "$1"; } + +# ---------------------------------------------------------------------------- +# Resolve where the registered Rust driver .so lives so a fresh build can be +# dropped in place without touching odbcinst.ini. +# ---------------------------------------------------------------------------- +registered_driver_path() { + local name="$1" + local ini + for ini in "$HOME/.odbcinst.ini" /etc/odbcinst.ini /usr/local/etc/odbcinst.ini; do + [[ -f "$ini" ]] || continue + awk -v section="[$name]" ' + $0 == section { inside = 1; next } + /^\[/ { inside = 0 } + inside && /^[[:space:]]*[Dd]river[[:space:]]*=/ { + sub(/^[^=]*=[[:space:]]*/, ""); print; exit + } + ' "$ini" | head -n1 + done | head -n1 +} + +RUST_DRIVER_PATH="$(registered_driver_path "$RUST_DRIVER")" +if [[ -z "$RUST_DRIVER_PATH" ]]; then + echo "Driver '$RUST_DRIVER' is not registered in any odbcinst.ini." >&2 + echo "See the header of this script for the one-time registration steps." >&2 + exit 1 +fi + +# ---------------------------------------------------------------------------- +# Build +# ---------------------------------------------------------------------------- +if [[ $SKIP_BUILD -eq 0 ]]; then + section "Building Rust driver (release)" + (cd "$ODBC_CRATE_DIR" && cargo build --release) + + BUILT_LIB="" + for candidate in \ + "$WORKSPACE_DIR/target/release/libmsodbcsql18.so" \ + "$WORKSPACE_DIR/target/release/libmsodbcsql18.dylib"; do + [[ -f "$candidate" ]] && BUILT_LIB="$candidate" && break + done + if [[ -z "$BUILT_LIB" ]]; then + echo "Built driver not found under $WORKSPACE_DIR/target/release" >&2 + exit 1 + fi + + echo "[ DRIVER ] Copying $BUILT_LIB -> $RUST_DRIVER_PATH" + mkdir -p "$(dirname "$RUST_DRIVER_PATH")" + cp -f "$BUILT_LIB" "$RUST_DRIVER_PATH" + + section "Building C++ benchmarks (Release)" + cmake -S "$SCRIPT_DIR" -B "$BUILD_DIR" -DCMAKE_BUILD_TYPE=Release + cmake --build "$BUILD_DIR" -j"$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)" +fi + +# ---------------------------------------------------------------------------- +# Run one benchmark binary against one driver +# ---------------------------------------------------------------------------- +run_bench() { + local bench_name="$1" driver_name="$2" label="$3" + local exe="$BUILD_DIR/$bench_name" + if [[ ! -x "$exe" ]]; then + echo "Benchmark binary not found: $exe" >&2 + exit 1 + fi + + local out="$RESULTS_DIR/$bench_name.$label.json" + local args=( + "--benchmark_out=$out" + --benchmark_out_format=json + "--benchmark_min_time=${MIN_TIME}s" + "--benchmark_repetitions=$REPETITIONS" + --benchmark_report_aggregates_only=true + ) + [[ -n "$FILTER" ]] && args+=("--benchmark_filter=$FILTER") + + echo "[ RUN ] $bench_name driver='$driver_name'" + ODBC_TEST_DRIVER="$driver_name" \ + ODBC_TEST_SERVER="$SERVER" \ + ODBC_TEST_DATABASE="$DATABASE" \ + ODBC_TEST_UID="$UID_ARG" \ + ODBC_TEST_PWD="$PWD_ARG" \ + "$exe" "${args[@]}" || echo "[ FAIL ] $bench_name exited $? for '$driver_name'" +} + +# ---------------------------------------------------------------------------- +# Compare two result files. Emits "namerust_nsref_ns" lines; a failed +# case reports -1 so it stays visible in the table instead of disappearing. +# ---------------------------------------------------------------------------- +compare_results() { + local rust_json="$1" ref_json="$2" + [[ -f "$rust_json" && -f "$ref_json" ]] || return 0 + command -v python3 >/dev/null 2>&1 || return 0 + + python3 - "$rust_json" "$ref_json" <<'PY' +import json, sys + +SCALE = {"ns": 1, "us": 1e3, "ms": 1e6, "s": 1e9} + +def medians(path): + out = {} + with open(path) as fh: + data = json.load(fh) + for b in data.get("benchmarks", []): + name = b["name"] + for suffix in ("_median", "_mean", "_stddev", "_cv"): + if name.endswith(suffix): + name = name[: -len(suffix)] + break + if b.get("error_occurred"): + out[name] = None + continue + if b.get("run_type") == "aggregate" and b.get("aggregate_name") != "median": + continue + if out.get(name, False) is None: + continue + out[name] = b["real_time"] * SCALE.get(b.get("time_unit", "ns"), 1) + return out + +rust, ref = medians(sys.argv[1]), medians(sys.argv[2]) +for name, r in rust.items(): + if name not in ref: + continue + m = ref[name] + print("%s\t%s\t%s" % (name, -1 if r is None else r, -1 if m is None else m)) +PY +} + +format_table() { + awk -F'\t' ' + function fmt(ns) { + if (ns < 0) return "ERROR" + if (ns >= 1e6) return sprintf("%.2f ms", ns / 1e6) + if (ns >= 1e3) return sprintf("%.2f us", ns / 1e3) + return sprintf("%.2f ns", ns) + } + { + r = $2 + 0; m = $3 + 0 + if (r < 0 || m <= 0) { + ratio = "-"; key = 1e18 + verdict = (r < 0) ? "unsupported" : "n/a" + } else { + key = r / m + ratio = sprintf("%.2f", key) + verdict = (r < m * 0.95) ? "faster" : (r > m * 1.05) ? "slower" : "parity" + } + printf "%.9f\t%s|%s|%s|%s|%s\n", key, $1, fmt(r), fmt(m), ratio, verdict + } + ' | sort -k1,1g | cut -f2- | { + printf 'Benchmark|mssql-odbc|msodbcsql18|Ratio|Verdict\n' + cat + } | column -t -s'|' +} + +# ---------------------------------------------------------------------------- +# Main +# ---------------------------------------------------------------------------- +section "Configuration" +echo " Rust driver : $RUST_DRIVER ($RUST_DRIVER_PATH)" +[[ -n "$REF_DRIVER" ]] && echo " Reference : $REF_DRIVER" +echo " Server : $SERVER / $DATABASE" +if [[ -n "$UID_ARG" ]]; then + echo " Auth : SQL login '$UID_ARG'" +else + echo " Auth : integrated" +fi +echo " Benchmarks : ${BENCHES[*]}" + +ROWS_FILE="$(mktemp)" +trap 'rm -f "$ROWS_FILE"' EXIT + +for b in "${BENCHES[@]}"; do + section "Benchmark: $b" + run_bench "$b" "$RUST_DRIVER" "mssql-odbc" + if [[ -n "$REF_DRIVER" ]]; then + run_bench "$b" "$REF_DRIVER" "msodbcsql18" + compare_results "$RESULTS_DIR/$b.mssql-odbc.json" \ + "$RESULTS_DIR/$b.msodbcsql18.json" >> "$ROWS_FILE" + fi +done + +if [[ -s "$ROWS_FILE" ]]; then + section "Comparison (median real time; Ratio < 1.0 = mssql-odbc faster)" + format_table < "$ROWS_FILE" +fi + +echo +echo "Raw JSON results in: $RESULTS_DIR" diff --git a/mssql-tds/src/connection/tds_client.rs b/mssql-tds/src/connection/tds_client.rs index 7fd3155b..9e68c6e3 100644 --- a/mssql-tds/src/connection/tds_client.rs +++ b/mssql-tds/src/connection/tds_client.rs @@ -95,6 +95,9 @@ pub struct TdsClient { /// the metadata it was built for so it is rebuilt when the result set /// changes. `None` until the first encrypted result set is seen. current_decryptor: Option, + /// Row-decode parser context for `current_metadata`, built once per result + /// set instead of once per row. Rebuilt when the metadata `Arc` changes. + current_parser_context: Option, count_map: HashMap, /// Rows affected by the most recent statement; see [`last_rows_affected`](Self::last_rows_affected). last_rows_affected: i64, @@ -193,6 +196,7 @@ impl TdsClient { recovery_context: Box::new(recovery_context), current_metadata: None, current_decryptor: None, + current_parser_context: None, count_map: HashMap::new(), last_rows_affected: -1, dml_result_counts: Vec::new(), @@ -2086,7 +2090,6 @@ impl TdsClient { /// This functions returns to the next row in the result set. /// If there are no more rows, it returns None. - #[instrument(skip(self), level = "info")] pub(crate) async fn get_next_row(&mut self) -> TdsResult>> { let col_count = self .current_metadata @@ -2101,6 +2104,160 @@ impl TdsClient { } } + /// Fetches the next row, reusing `buffer`'s allocation instead of returning + /// a freshly allocated `Vec` per row. + /// + /// Returns `false` when the result set is exhausted, leaving `buffer` + /// cleared. Callers driving a cursor row-by-row (ODBC `SQLFetch`) should + /// prefer this over [`ResultSet::next_row`]: it avoids both the per-row + /// `Vec` allocation and the boxed future that `#[async_trait]` introduces + /// for the trait method. + pub async fn fetch_next_row_into_vec( + &mut self, + buffer: &mut Vec, + ) -> TdsResult { + if !self.maybe_has_unread_rows() { + buffer.clear(); + return Ok(false); + } + let col_count = self + .current_metadata + .as_ref() + .map(|m| m.columns.len()) + .unwrap_or(0); + let mut writer = DefaultRowWriter::with_buffer(std::mem::take(buffer), col_count); + // Boxed deliberately: measurably faster than letting this future be + // materialized inline in the caller's frame on every fetch. + let has_row = Box::pin(self.get_next_row_into(&mut writer)).await?; + *buffer = writer.take_row(); + if !has_row { + buffer.clear(); + } + Ok(has_row) + } + + /// Fetches up to `max_rows` rows in a single decode call, appending them to + /// `out`. + /// + /// The per-row cost of a cursor is dominated by constructing and polling the + /// decode state machine, not by the decoding itself. Draining several rows + /// inside one future amortizes that cost over the whole batch. `spare` + /// supplies row buffers the caller has already consumed so the batch does + /// not allocate a fresh `Vec` per row. + /// + /// Returns the number of rows appended; a short batch means the result set + /// is exhausted. Falls back to a single row when decoding pauses mid-row + /// (incremental PLP / column streaming), preserving the paused-cursor + /// contract of [`Self::fetch_next_row_into_vec`]. + pub async fn fetch_rows_batch( + &mut self, + out: &mut Vec>, + spare: Vec>, + max_rows: usize, + ) -> TdsResult { + if max_rows == 0 || !self.maybe_has_unread_rows() { + return Ok(0); + } + let col_count = self + .current_metadata + .as_ref() + .map(|m| m.columns.len()) + .unwrap_or(0); + let mut writer = DefaultRowWriter::batching(col_count, spare); + // Boxed deliberately: materializing this future inline in the caller's + // frame is measurably slower than the single allocation. + let count = Box::pin(self.get_rows_into(&mut writer, max_rows)).await?; + out.append(&mut writer.take_completed()); + Ok(count) + } + + /// Decodes up to `max_rows` rows through a batch-mode writer inside a single + /// future, hoisting the per-result-set setup (metadata resolution, parser + /// context, decryptor lookup) out of the per-row path. + async fn get_rows_into( + &mut self, + writer: &mut DefaultRowWriter, + max_rows: usize, + ) -> TdsResult { + if self.current_metadata.is_none() { + return Err(UsageError( + "No metadata found while fetching the next row. Have you called the execute method or was the query supposed to return resultset?".to_string(), + )); + } + if !matches!(self.active_row_read_state, ActiveRowReadState::Idle) { + return Err(crate::error::Error::ImplementationError( + "Batched fetch cannot start while a row read is paused".to_string(), + )); + } + + let metadata = Arc::clone(self.current_metadata.as_ref().unwrap()); + let cached = matches!( + &self.current_parser_context, + Some(ParserContext::ColumnMetadata(built_for, _)) if Arc::ptr_eq(built_for, &metadata) + ); + if !cached { + let decryptor = self.resolve_cell_decryptor(&metadata).await?; + self.current_parser_context = Some(ParserContext::ColumnMetadata(metadata, decryptor)); + } + let parser_context = self.current_parser_context.clone().unwrap(); + + let mut count = 0; + while count < max_rows { + // `Instant::now` is a counter read on every call; skip both of them + // when no request timeout is being tracked. + let start = self.remaining_request_timeout.map(|_| Instant::now()); + let batch = self + .transport + .receive_rows_into( + &parser_context, + self.remaining_request_timeout, + self.cancel_handle.as_ref(), + writer, + max_rows - count, + ) + .await?; + if let Some(start) = start { + self.update_remaining_timeout(start); + } + count += batch.rows; + + match batch.stopped_at { + None => break, + Some(RowReadResult::RowWritten) => unreachable!("batch consumes written rows"), + Some(RowReadResult::RowPaused(_) | RowReadResult::PlpPaused(_)) => { + return Err(crate::error::Error::ImplementationError( + "Row decode paused against a non-pausing batch writer".to_string(), + )); + } + Some(RowReadResult::Token(token)) => { + // Boxed: non-row tokens appear once per result set, but the + // handler inlines every token parser's state machine. + if Box::pin(self.handle_row_read_token(token)).await?.is_some() { + break; + } + } + } + } + Ok(count) + } + + /// Diagnostic: byte sizes of the futures on the row-fetch hot path. + /// + /// Exposed so perf work can verify that changes actually shrink the + /// per-row state machines rather than just moving allocations around. + #[doc(hidden)] + pub fn row_future_sizes(&mut self) -> (usize, usize) { + let mut writer = DefaultRowWriter::new(0); + let inner = self.get_next_row_into(&mut writer); + let size_inner = std::mem::size_of_val(&inner); + drop(inner); + let mut buf = Vec::new(); + let outer = self.fetch_next_row_into_vec(&mut buf); + let size_outer = std::mem::size_of_val(&outer); + drop(outer); + (size_inner, size_outer) + } + /// Returns `true` when transparent parameter encryption should be attempted: /// the connection requested Always Encrypted and the server acknowledged the /// feature during login. @@ -2834,7 +2991,6 @@ impl TdsClient { /// /// Uses `receive_row_into` to decode ROW/NBCROW tokens directly through /// `decode_into`, bypassing the intermediate `RowToken { all_values }`. - #[instrument(skip(self, writer), level = "info")] pub(crate) async fn get_next_row_into( &mut self, writer: &mut (dyn RowWriter + Send), @@ -2851,7 +3007,10 @@ impl TdsClient { return self.resume_row_loop(*pause_state, writer).await; } ActiveRowReadState::PlpPaused(mut plp_state) => { - let mut buffer = [0u8; 8192]; + // Heap-allocated deliberately: an inline `[u8; 8192]` would be + // baked into this function's future and paid for on every row, + // not just on the rare paused-PLP path. + let mut buffer = vec![0u8; 8192]; while !plp_state.reached_end() { let start = Instant::now(); let read = self @@ -2878,8 +3037,15 @@ impl TdsClient { } let metadata = Arc::clone(self.current_metadata.as_ref().unwrap()); - let decryptor = self.resolve_cell_decryptor(&metadata).await?; - let parser_context = ParserContext::ColumnMetadata(metadata, decryptor); + let cached = matches!( + &self.current_parser_context, + Some(ParserContext::ColumnMetadata(built_for, _)) if Arc::ptr_eq(built_for, &metadata) + ); + if !cached { + let decryptor = self.resolve_cell_decryptor(&metadata).await?; + self.current_parser_context = Some(ParserContext::ColumnMetadata(metadata, decryptor)); + } + let parser_context = self.current_parser_context.clone().unwrap(); loop { let start = Instant::now(); let result = self @@ -2896,7 +3062,6 @@ impl TdsClient { match result { RowReadResult::RowWritten => { writer.end_row(); - info!("Row Received"); return Ok(true); } RowReadResult::RowPaused(pause_state) => { @@ -2937,7 +3102,6 @@ impl TdsClient { match result { RowReadResult::RowWritten => { writer.end_row(); - info!("Row Received"); Ok(true) } RowReadResult::RowPaused(next_pause) => { @@ -3466,7 +3630,6 @@ impl ResultSet for TdsClient { .unwrap_or(&self.empty_metadata) } - #[instrument(skip(self), level = "info")] async fn next_row(&mut self) -> TdsResult>> { if self.maybe_has_unread_rows() { self.get_next_row().await @@ -3475,7 +3638,6 @@ impl ResultSet for TdsClient { } } - #[instrument(skip(self, writer), level = "info")] async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult { if self.maybe_has_unread_rows() { self.get_next_row_into(writer).await diff --git a/mssql-tds/src/connection/transport/network_transport.rs b/mssql-tds/src/connection/transport/network_transport.rs index 57dbd00e..9a36f4ff 100644 --- a/mssql-tds/src/connection/transport/network_transport.rs +++ b/mssql-tds/src/connection/transport/network_transport.rs @@ -18,9 +18,9 @@ use crate::io::packet_reader::{PacketReader, TdsPacketReader}; use crate::io::packet_writer::PacketWriter; use crate::io::reader_writer::{NetworkReader, NetworkReaderWriter, NetworkWriter}; use crate::io::token_stream::{ - ParserContext, PlpPauseState, RowPauseState, RowReadResult, TdsTokenStreamReader, - read_active_plp_bytes_internal, receive_row_into_internal, receive_token_internal, - resume_row_into_internal, + BatchRowsResult, ParserContext, PlpPauseState, RowPauseState, RowReadResult, + TdsTokenStreamReader, read_active_plp_bytes_internal, receive_row_into_internal, + receive_rows_into_internal, receive_token_internal, resume_row_into_internal, }; use crate::message::attention::AttentionRequest; use crate::message::login_options::TdsVersion; @@ -1143,6 +1143,17 @@ impl TdsPacketReader for NetworkTransport { Ok(result) } + fn try_take_fixed(&mut self, n: usize) -> Option<[u8; 8]> { + debug_assert!(n <= 8); + if n > 8 || !self.tds_read_buffer.do_we_have_enough_data(n) { + return None; + } + let mut out = [0u8; 8]; + out[..n].copy_from_slice(&self.tds_read_buffer.get_slice()[..n]); + self.tds_read_buffer.consume_bytes(n); + Some(out) + } + async fn read_int16_big_endian(&mut self) -> TdsResult { if !self.tds_read_buffer.do_we_have_enough_data(2) { self.read_tds_packet().await?; @@ -1434,6 +1445,38 @@ impl TdsTokenStreamReader for NetworkTransport { result } + async fn receive_rows_into( + &mut self, + context: &ParserContext, + remaining_request_timeout: Option, + cancel_handle: Option<&CancelHandle>, + writer: &mut (dyn RowWriter + Send), + max_rows: usize, + ) -> TdsResult { + let cancellable = CancelHandle::run_until_cancelled( + cancel_handle, + receive_rows_into_internal(self, &*PARSER_REGISTRY, context, writer, max_rows), + ); + let result = match remaining_request_timeout.as_ref() { + Some(t) => match timeout(*t, cancellable).await { + Ok(r) => r, + Err(elapsed) => Err(TimeoutError(TimeoutErrorType::Elapsed(elapsed))), + }, + None => cancellable.await, + }; + + match &result { + Ok(_) => {} + Err(err) => match err { + OperationCancelledError(_) | TimeoutError(_) => { + self.cancel_read_stream_and_wait().await?; + } + _ => {} + }, + } + result + } + async fn resume_row_into( &mut self, pause_state: RowPauseState, diff --git a/mssql-tds/src/datatypes/decoder.rs b/mssql-tds/src/datatypes/decoder.rs index fa6bb22a..ad589581 100644 --- a/mssql-tds/src/datatypes/decoder.rs +++ b/mssql-tds/src/datatypes/decoder.rs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use crate::io::packet_reader::fast; + use async_trait::async_trait; use core::fmt; use std::sync::Arc; @@ -196,7 +198,7 @@ impl PlpChunkStreamReader { pub(crate) async fn begin( reader: &mut (dyn TdsPacketReader + Send + Sync), ) -> TdsResult> { - let raw_len_i64 = reader.read_int64().await?; + let raw_len_i64 = fast::read_int64(reader).await?; let raw_len = raw_len_i64 as u64; let raw_len_usize = raw_len as usize; @@ -241,7 +243,7 @@ impl PlpChunkStreamReader { return Ok(true); } - let chunk_len = reader.read_uint32().await? as usize; + let chunk_len = fast::read_uint32(reader).await? as usize; if chunk_len == 0 { self.reached_end = true; if let PlpChunkReadLength::Known(known_len) = self.length @@ -489,10 +491,10 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let length = reader.read_uint32().await?; - let variant_base_type = reader.read_byte().await?; + let length = fast::read_uint32(reader).await?; + let variant_base_type = fast::read_byte(reader).await?; let tds_type = TdsDataType::try_from(variant_base_type)?; - let variant_prop_bytes = reader.read_byte().await?; + let variant_prop_bytes = fast::read_byte(reader).await?; let bytes_for_type_and_properties_byte = 2; // Use checked arithmetic to prevent integer underflow @@ -584,7 +586,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let scale = reader.read_byte().await?; + let scale = fast::read_byte(reader).await?; Ok(match tds_type { TdsDataType::TimeN => { let time_nanos = self.read_time(reader, data_length as u8, scale).await?; @@ -615,7 +617,7 @@ impl GenericDecoder { T: TdsPacketReader + Send + Sync, { // Decimal/numeric data type has 1 byte length. - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; let TypeInfoVariant::VarLenPrecisionScale(_, _, precision, scale) = metadata.type_info.type_info_variant else { @@ -640,7 +642,7 @@ impl GenericDecoder { if length == 0 { return Ok(None); } - let sign = reader.read_byte().await?; + let sign = fast::read_byte(reader).await?; let is_positive = sign == 1; let number_of_int_parts = (length - 1) >> 2; @@ -661,7 +663,7 @@ impl GenericDecoder { validate_alloc_size(int_parts_len * 4, "read_decimal int_parts")?; let mut int_parts = vec![0i32; int_parts_len]; for part_index in 0..number_of_int_parts { - int_parts[part_index as usize] = reader.read_int32().await?; + int_parts[part_index as usize] = fast::read_int32(reader).await?; } Ok(Some(DecimalParts { @@ -676,8 +678,8 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let days = reader.read_int32().await?; - let ticks = reader.read_uint32().await?; + let days = fast::read_int32(reader).await?; + let ticks = fast::read_uint32(reader).await?; Ok(SqlDateTime { days, time: ticks }) } @@ -686,8 +688,8 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let days = reader.read_uint16().await?; - let minutes = reader.read_uint16().await?; + let days = fast::read_uint16(reader).await?; + let minutes = fast::read_uint16(reader).await?; Ok(SqlSmallDateTime { days, time: minutes, @@ -698,7 +700,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let days = reader.read_uint24().await?; + let days = fast::read_uint24(reader).await?; Ok(SqlDate::unchecked_create(days)) } @@ -707,9 +709,9 @@ impl GenericDecoder { T: TdsPacketReader + Send + Sync, { let scaled_value = match byte_len { - 3 => reader.read_uint24().await? as u64, - 4 => reader.read_uint32().await? as u64, - _ => reader.read_uint40().await?, + 3 => fast::read_uint24(reader).await? as u64, + 4 => fast::read_uint32(reader).await? as u64, + _ => fast::read_uint40(reader).await?, }; // The value from SQL Server is in scaled units based on the scale: @@ -788,7 +790,7 @@ impl GenericDecoder { ))); } }; - let offset = reader.read_int16().await?; + let offset = fast::read_int16(reader).await?; let datetime_offset = SqlDateTimeOffset { datetime2, offset }; Ok(ColumnValues::DateTimeOffset(datetime_offset)) } @@ -798,10 +800,10 @@ impl GenericDecoder { T: TdsPacketReader + Send + Sync, { let value: ColumnValues = match byte_len { - 1 => ColumnValues::TinyInt(reader.read_byte().await?), // Some(reader.read_byte().await? as i64), - 2 => ColumnValues::SmallInt(reader.read_int16().await?), // Some(reader.read_int16().await? as i64), - 4 => ColumnValues::Int(reader.read_int32().await?), - 8 => ColumnValues::BigInt(reader.read_int64().await?), + 1 => ColumnValues::TinyInt(fast::read_byte(reader).await?), // Some(fast::read_byte(reader).await? as i64), + 2 => ColumnValues::SmallInt(fast::read_int16(reader).await?), // Some(fast::read_int16(reader).await? as i64), + 4 => ColumnValues::Int(fast::read_int32(reader).await?), + 8 => ColumnValues::BigInt(fast::read_int64(reader).await?), 0 => ColumnValues::Null, _ => { return Err(crate::error::Error::from(Error::new( @@ -817,7 +819,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let small_money_val = reader.read_int32().await?; + let small_money_val = fast::read_int32(reader).await?; Ok(small_money_val.into()) } @@ -827,8 +829,8 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let msb = reader.read_int32().await?; - let lsb = reader.read_int32().await?; + let msb = fast::read_int32(reader).await?; + let lsb = fast::read_int32(reader).await?; Ok(SqlMoney { lsb_part: lsb, msb_part: msb, @@ -884,7 +886,7 @@ impl GenericDecoder { }; // Read length prefix (USHORTLEN format) - let length_prefix_value = reader.read_uint16().await? as usize; + let length_prefix_value = fast::read_uint16(reader).await? as usize; // Handle NULL (length = 0xFFFF) if length_prefix_value == 0xFFFF { @@ -908,13 +910,13 @@ impl GenericDecoder { } // Read 8-byte header - let layout_format_byte = reader.read_byte().await?; - let layout_version_byte = reader.read_byte().await?; - let dimension_count = reader.read_uint16().await?; - let base_type_byte = reader.read_byte().await?; - let _reserved1 = reader.read_byte().await?; // Reserved - let _reserved2 = reader.read_byte().await?; // Reserved - let _reserved3 = reader.read_byte().await?; // Reserved + let layout_format_byte = fast::read_byte(reader).await?; + let layout_version_byte = fast::read_byte(reader).await?; + let dimension_count = fast::read_uint16(reader).await?; + let base_type_byte = fast::read_byte(reader).await?; + let _reserved1 = fast::read_byte(reader).await?; // Reserved + let _reserved2 = fast::read_byte(reader).await?; // Reserved + let _reserved3 = fast::read_byte(reader).await?; // Reserved // Validate header using enum conversions let _layout_format = VectorLayoutFormat::try_from(layout_format_byte)?; @@ -978,7 +980,7 @@ impl GenericDecoder { where T: TdsPacketReader + Send + Sync, { - let long_len_i64 = reader.read_int64().await?; + let long_len_i64 = fast::read_int64(reader).await?; let long_len = long_len_i64 as u64; // If the length is SQL_PLP_NULL, it means the value is NULL. @@ -1002,7 +1004,7 @@ impl GenericDecoder { 0 }; let mut plp_buffer = vec![0u8; vector_capacity]; - let mut chunk_len = reader.read_uint32().await? as usize; + let mut chunk_len = fast::read_uint32(reader).await? as usize; let mut offset: usize = 0; let mut chunk_count = 0u32; @@ -1063,7 +1065,7 @@ impl GenericDecoder { .read_bytes(&mut plp_buffer[offset..offset + chunk_len]) .await?; offset += chunk_size_read; - chunk_len = reader.read_uint32().await? as usize; + chunk_len = fast::read_uint32(reader).await? as usize; } Ok(Some(plp_buffer)) } @@ -1087,24 +1089,24 @@ impl GenericDecoder { match metadata.data_type { // === Fixed-length integer types === TdsDataType::Int1 => { - writer.write_u8(col, reader.read_byte().await?); + writer.write_u8(col, fast::read_byte(reader).await?); } TdsDataType::Int2 => { - writer.write_i16(col, reader.read_int16().await?); + writer.write_i16(col, fast::read_int16(reader).await?); } TdsDataType::Int4 => { - writer.write_i32(col, reader.read_int32().await?); + writer.write_i32(col, fast::read_int32(reader).await?); } TdsDataType::Int8 => { - writer.write_i64(col, reader.read_int64().await?); + writer.write_i64(col, fast::read_int64(reader).await?); } TdsDataType::IntN => { - let byte_len = reader.read_byte().await?; + let byte_len = fast::read_byte(reader).await?; match byte_len { - 1 => writer.write_u8(col, reader.read_byte().await?), - 2 => writer.write_i16(col, reader.read_int16().await?), - 4 => writer.write_i32(col, reader.read_int32().await?), - 8 => writer.write_i64(col, reader.read_int64().await?), + 1 => writer.write_u8(col, fast::read_byte(reader).await?), + 2 => writer.write_i16(col, fast::read_int16(reader).await?), + 4 => writer.write_i32(col, fast::read_int32(reader).await?), + 8 => writer.write_i64(col, fast::read_int64(reader).await?), 0 => writer.write_null(col), _ => { return Err(crate::error::Error::from(Error::new( @@ -1117,28 +1119,28 @@ impl GenericDecoder { // === Fixed-length float types === TdsDataType::Flt4 => { - writer.write_f32(col, reader.read_float32().await?); + writer.write_f32(col, fast::read_float32(reader).await?); } TdsDataType::Flt8 => { - writer.write_f64(col, reader.read_float64().await?); + writer.write_f64(col, fast::read_float64(reader).await?); } TdsDataType::FltN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; match length { 0 => writer.write_null(col), - 4 => writer.write_f32(col, reader.read_float32().await?), - _ => writer.write_f64(col, reader.read_float64().await?), + 4 => writer.write_f32(col, fast::read_float32(reader).await?), + _ => writer.write_f64(col, fast::read_float64(reader).await?), } } // === Bit types === TdsDataType::Bit => { - writer.write_bool(col, reader.read_byte().await? == 1); + writer.write_bool(col, fast::read_byte(reader).await? == 1); } TdsDataType::BitN => { - let byte_len = reader.read_byte().await?; + let byte_len = fast::read_byte(reader).await?; if byte_len > 0 { - writer.write_bool(col, reader.read_byte().await? == 1); + writer.write_bool(col, fast::read_byte(reader).await? == 1); } else { writer.write_null(col); } @@ -1152,7 +1154,7 @@ impl GenericDecoder { writer.write_money(col, self.read_money8(reader).await?); } TdsDataType::MoneyN => { - let byte_len = reader.read_byte().await?; + let byte_len = fast::read_byte(reader).await?; match byte_len { 4 => writer.write_smallmoney(col, self.read_money4(reader).await?), 8 => writer.write_money(col, self.read_money8(reader).await?), @@ -1166,11 +1168,11 @@ impl GenericDecoder { } // === Decimal / Numeric === - TdsDataType::DecimalN => match self.read_decimal(reader, metadata).await? { + TdsDataType::DecimalN => match Box::pin(self.read_decimal(reader, metadata)).await? { Some(val) => writer.write_decimal(col, val), None => writer.write_null(col), }, - TdsDataType::NumericN => match self.read_decimal(reader, metadata).await? { + TdsDataType::NumericN => match Box::pin(self.read_decimal(reader, metadata)).await? { Some(val) => writer.write_numeric(col, val), None => writer.write_null(col), }, @@ -1191,7 +1193,7 @@ impl GenericDecoder { // === Binary types === TdsDataType::BigBinary => { - let length = reader.read_uint16().await?; + let length = fast::read_uint16(reader).await?; // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { writer.write_null(col); @@ -1208,12 +1210,12 @@ impl GenericDecoder { } TdsDataType::BigVarBinary => { if metadata.is_plp() { - match GenericDecoder::read_plp_bytes(reader).await? { + match Box::pin(GenericDecoder::read_plp_bytes(reader)).await? { Some(bytes) => writer.write_bytes(col, bytes), None => writer.write_null(col), } } else { - let length = reader.read_uint16().await?; + let length = fast::read_uint16(reader).await?; // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { writer.write_null(col); @@ -1235,8 +1237,8 @@ impl GenericDecoder { writer.write_datetime(col, self.read_datetime(reader).await?); } TdsDataType::DateTim4 => { - let daypart = reader.read_uint16().await?; - let timepart = reader.read_uint16().await?; + let daypart = fast::read_uint16(reader).await?; + let timepart = fast::read_uint16(reader).await?; writer.write_smalldatetime( col, SqlSmallDateTime { @@ -1246,7 +1248,7 @@ impl GenericDecoder { ); } TdsDataType::DateTimeN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; match length { 0 => writer.write_null(col), 4 => writer.write_smalldatetime(col, self.read_small_datetime(reader).await?), @@ -1254,7 +1256,7 @@ impl GenericDecoder { } } TdsDataType::DateN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; if length == 0 { writer.write_null(col); } else { @@ -1262,7 +1264,7 @@ impl GenericDecoder { } } TdsDataType::TimeN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; if length == 0 { writer.write_null(col); } else { @@ -1282,42 +1284,40 @@ impl GenericDecoder { } } TdsDataType::DateTime2N => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; if length == 0 { writer.write_null(col); } else { - let cv = self - .read_datetime2( - reader, - length, - metadata.get_scale().ok_or_else(|| { - crate::error::Error::ImplementationError( - "DateTime2N type should have scale".to_string(), - ) - })?, - ) - .await?; + let cv = Box::pin(self.read_datetime2( + reader, + length, + metadata.get_scale().ok_or_else(|| { + crate::error::Error::ImplementationError( + "DateTime2N type should have scale".to_string(), + ) + })?, + )) + .await?; if let ColumnValues::DateTime2(dt2) = cv { writer.write_datetime2(col, dt2); } } } TdsDataType::DateTimeOffsetN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; if length == 0 { writer.write_null(col); } else { - let cv = self - .read_datetime_offset( - reader, - length, - metadata.get_scale().ok_or_else(|| { - crate::error::Error::ImplementationError( - "DateTimeOffsetN type should have scale".to_string(), - ) - })?, - ) - .await?; + let cv = Box::pin(self.read_datetime_offset( + reader, + length, + metadata.get_scale().ok_or_else(|| { + crate::error::Error::ImplementationError( + "DateTimeOffsetN type should have scale".to_string(), + ) + })?, + )) + .await?; if let ColumnValues::DateTimeOffset(dto) = cv { writer.write_datetimeoffset(col, dto); } @@ -1326,7 +1326,7 @@ impl GenericDecoder { // === GUID === TdsDataType::Guid => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; if length == 0 { writer.write_null(col); } else { @@ -1346,7 +1346,7 @@ impl GenericDecoder { // === Fallback: rare types go through decode() → write_column_value() === _ => { - let value = self.decode(reader, metadata).await?; + let value = Box::pin(self.decode(reader, metadata)).await?; write_column_value(writer, col, value); } } @@ -1362,33 +1362,33 @@ impl SqlTypeDecode for GenericDecoder { { let result = match metadata.data_type { TdsDataType::Int1 => { - let value = reader.read_byte().await?; + let value = fast::read_byte(reader).await?; ColumnValues::from(value) } TdsDataType::Int2 => { - let value = reader.read_int16().await?; + let value = fast::read_int16(reader).await?; ColumnValues::SmallInt(value) } TdsDataType::Int4 => { - let value = reader.read_int32().await?; + let value = fast::read_int32(reader).await?; ColumnValues::from(value) } TdsDataType::Int8 => { - let value = reader.read_int64().await?; + let value = fast::read_int64(reader).await?; ColumnValues::BigInt(value) } TdsDataType::Flt4 => { - let value = reader.read_float32().await?; + let value = fast::read_float32(reader).await?; ColumnValues::Real(value) } TdsDataType::Flt8 => { - let value = reader.read_float64().await?; + let value = fast::read_float64(reader).await?; ColumnValues::Float(value) } TdsDataType::Money4 => ColumnValues::SmallMoney(self.read_money4(reader).await?), TdsDataType::Money => ColumnValues::Money(self.read_money8(reader).await?), TdsDataType::MoneyN => { - let byte_len = reader.read_byte().await?; + let byte_len = fast::read_byte(reader).await?; match byte_len { 4 => ColumnValues::SmallMoney(self.read_money4(reader).await?), 8 => ColumnValues::Money(self.read_money8(reader).await?), @@ -1415,7 +1415,7 @@ impl SqlTypeDecode for GenericDecoder { } } TdsDataType::Bit => { - let value = reader.read_byte().await?; + let value = fast::read_byte(reader).await?; ColumnValues::Bit(value == 1) } TdsDataType::NChar @@ -1431,11 +1431,11 @@ impl SqlTypeDecode for GenericDecoder { ColumnValues::DateTime(value) } TdsDataType::IntN => { - let byte_len = reader.read_byte().await?; + let byte_len = fast::read_byte(reader).await?; self.read_intn(reader, byte_len).await? } TdsDataType::BigBinary => { - let length = reader.read_uint16().await?; + let length = fast::read_uint16(reader).await?; // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { ColumnValues::Null @@ -1458,7 +1458,7 @@ impl SqlTypeDecode for GenericDecoder { None => ColumnValues::Null, } } else { - let length = reader.read_uint16().await?; + let length = fast::read_uint16(reader).await?; // 0xFFFF is the USHORTLEN NULL marker (CHARBIN_NULL). if length == 0xFFFF { ColumnValues::Null @@ -1498,36 +1498,36 @@ impl SqlTypeDecode for GenericDecoder { None => ColumnValues::Null, } } - TdsDataType::Vector => self.decode_vector(reader, metadata).await?, + TdsDataType::Vector => Box::pin(self.decode_vector(reader, metadata)).await?, TdsDataType::BitN => { - let byte_len = reader.read_byte().await?; + let byte_len = fast::read_byte(reader).await?; if byte_len > 0 { - let value = reader.read_byte().await?; + let value = fast::read_byte(reader).await?; ColumnValues::Bit(value == 1) } else { ColumnValues::Null } } TdsDataType::Guid => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; Self::read_guid(reader, length).await? } TdsDataType::FltN => { // This is variable length float, hence the length needs to be read first - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; if length == 0 { return Ok(ColumnValues::Null); } if length == 4 { - let value = reader.read_float32().await?; + let value = fast::read_float32(reader).await?; ColumnValues::Real(value) } else { - let value = reader.read_float64().await?; + let value = fast::read_float64(reader).await?; ColumnValues::Float(value) } } TdsDataType::DateTimeN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; // If length is 0, then it is NULL if length == 0 { return Ok(ColumnValues::Null); @@ -1541,11 +1541,11 @@ impl SqlTypeDecode for GenericDecoder { } } TdsDataType::DateN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; return Self::read_daten(reader, length).await; } TdsDataType::TimeN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; match length { 0 => return Ok(ColumnValues::Null), _ => { @@ -1565,7 +1565,7 @@ impl SqlTypeDecode for GenericDecoder { } } TdsDataType::DateTime2N => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; match length { 0 => Ok(ColumnValues::Null), _ => { @@ -1583,7 +1583,7 @@ impl SqlTypeDecode for GenericDecoder { } }?, TdsDataType::DateTimeOffsetN => { - let length = reader.read_byte().await?; + let length = fast::read_byte(reader).await?; match length { 0 => Ok(ColumnValues::Null), _ => { @@ -1601,13 +1601,13 @@ impl SqlTypeDecode for GenericDecoder { } }?, TdsDataType::Image => { - let text_ptr_len = reader.read_byte().await? as usize; + let text_ptr_len = fast::read_byte(reader).await? as usize; let length = if text_ptr_len > 0 { const TIMESTAMP_BYTE_COUNT: usize = 8; reader.skip_bytes(text_ptr_len).await?; reader.skip_bytes(TIMESTAMP_BYTE_COUNT).await?; - reader.read_uint32().await? as usize + fast::read_uint32(reader).await? as usize } else { 0 }; @@ -1637,10 +1637,10 @@ impl SqlTypeDecode for GenericDecoder { None => ColumnValues::Null, } } - TdsDataType::SsVariant => self.read_sql_variant(reader).await?, + TdsDataType::SsVariant => Box::pin(self.read_sql_variant(reader)).await?, TdsDataType::DateTim4 => { - let daypart = reader.read_uint16().await?; - let timepart = reader.read_uint16().await?; + let daypart = fast::read_uint16(reader).await?; + let timepart = fast::read_uint16(reader).await?; ColumnValues::SmallDateTime(SqlSmallDateTime { days: daypart, time: timepart, @@ -1700,12 +1700,12 @@ impl StringDecoder { let encoding_type = get_encoding_type(metadata); if metadata.is_plp() { - match GenericDecoder::read_plp_bytes(reader).await? { + match Box::pin(GenericDecoder::read_plp_bytes(reader)).await? { Some(bytes) => writer.write_string(col, SqlString::new(bytes, encoding_type)), None => writer.write_null(col), } } else if Self::is_long_len_type(metadata.data_type) { - let text_ptr_len = reader.read_byte().await? as usize; + let text_ptr_len = fast::read_byte(reader).await? as usize; if text_ptr_len == 0 { writer.write_null(col); @@ -1715,7 +1715,7 @@ impl StringDecoder { const TIMESTAMP_BYTE_COUNT: usize = 8; reader.skip_bytes(text_ptr_len).await?; reader.skip_bytes(TIMESTAMP_BYTE_COUNT).await?; - let length = reader.read_uint32().await? as usize; + let length = fast::read_uint32(reader).await? as usize; if length > MAX_ALLOC_SIZE { return Err(crate::error::Error::ProtocolError(format!( @@ -1732,11 +1732,12 @@ impl StringDecoder { }; writer.write_string(col, sql_string); } else { - let length = reader.read_uint16().await? as usize; + let length = fast::read_uint16(reader).await? as usize; if length == 0xFFFF { writer.write_null(col); } else { - let mut buffer = vec![0u8; length]; + let mut buffer = writer.take_string_buffer(length); + buffer.resize(length, 0); reader.read_bytes(&mut buffer).await?; writer.write_string(col, SqlString::new(buffer, encoding_type)); } @@ -1786,13 +1787,13 @@ impl SqlTypeDecode for StringDecoder { // Creates SqlString with appropriate encoding type // NULL handling works (textptr_len = 0) // LCID-based decoding implemented (see sql_string.rs) - let text_ptr_len = reader.read_byte().await? as usize; + let text_ptr_len = fast::read_byte(reader).await? as usize; let length = if text_ptr_len > 0 { const TIMESTAMP_BYTE_COUNT: usize = 8; reader.skip_bytes(text_ptr_len).await?; reader.skip_bytes(TIMESTAMP_BYTE_COUNT).await?; - reader.read_uint32().await? as usize + fast::read_uint32(reader).await? as usize } else { // text_ptr_len == 0 means NULL value return Ok(ColumnValues::Null); @@ -1815,7 +1816,7 @@ impl SqlTypeDecode for StringDecoder { }; Ok(ColumnValues::String(sql_string)) } else { - let length = reader.read_uint16().await? as usize; + let length = fast::read_uint16(reader).await? as usize; if length == 0xFFFF { return Ok(ColumnValues::Null); } else { @@ -2096,7 +2097,7 @@ where Ok(match tds_type { // BIGVARBINARYTYPE, BIGBINARYTYPE TdsDataType::BigVarBinary | TdsDataType::BigBinary => { - let _max_length: u16 = reader.read_uint16().await?; + let _max_length: u16 = fast::read_uint16(reader).await?; if data_length as usize > MAX_ALLOC_SIZE { return Err(crate::error::Error::ProtocolError(format!( "SQL Variant binary data length {data_length} exceeds maximum allowed size of {MAX_ALLOC_SIZE} bytes" @@ -2107,8 +2108,8 @@ where ColumnValues::Bytes(buffer) } TdsDataType::NumericN | TdsDataType::DecimalN => { - let precision = reader.read_byte().await?; - let scale = reader.read_byte().await?; + let precision = fast::read_byte(reader).await?; + let scale = fast::read_byte(reader).await?; let decimal_parts = GenericDecoder::read_decimal_data(reader, data_length as u8, precision, scale) .await?; @@ -2151,7 +2152,7 @@ where } let mut collation_bytes = vec![0u8; 5]; reader.read_bytes(&mut collation_bytes).await?; - let _max_length = reader.read_uint16().await? as usize; + let _max_length = fast::read_uint16(reader).await? as usize; let collation: SqlCollation = collation_bytes.as_slice().try_into()?; if data_length as usize > MAX_ALLOC_SIZE { return Err(crate::error::Error::ProtocolError(format!( diff --git a/mssql-tds/src/datatypes/row_writer.rs b/mssql-tds/src/datatypes/row_writer.rs index eb8054e4..fa8f30f9 100644 --- a/mssql-tds/src/datatypes/row_writer.rs +++ b/mssql-tds/src/datatypes/row_writer.rs @@ -19,6 +19,15 @@ use uuid::Uuid; /// enabling consumers (Arrow writers, N-API binary encoders, etc.) to /// receive values without going through the intermediate `ColumnValues` enum. pub trait RowWriter { + /// Returns `false` when this writer never pauses, letting the decoder skip + /// the per-column [`RowWriter::pause_after_column`] check entirely. + /// + /// The default is `true` so overriding only `pause_after_column` stays + /// correct; writers that never pause should override this to `false`. + fn may_pause(&self) -> bool { + true + } + /// Returns `true` to pause row decoding after reading column `col`. /// /// Writers that need incremental column fetch behavior (for example, @@ -103,12 +112,33 @@ pub trait RowWriter { fn write_vector(&mut self, col: usize, val: SqlVector); /// Signals the end of the current row. fn end_row(&mut self); + + /// Supplies a byte buffer of at least `capacity` bytes for a variable-length + /// value the decoder is about to read. + /// + /// Writers that recycle rows can hand back a buffer harvested from a + /// previously consumed row, which removes an allocate/free pair per + /// variable-length column. The buffer is logically empty; the decoder sizes + /// it before filling it. + fn take_string_buffer(&mut self, capacity: usize) -> Vec { + Vec::with_capacity(capacity) + } } /// Default implementation that assembles `Vec`, preserving /// the current decoder behavior. Existing `next_row()` callers see no change. +/// +/// In *batch mode* (see [`DefaultRowWriter::batching`]) each `end_row` pushes +/// the completed row into an internal queue instead of leaving it in place, so +/// one decode call can materialize many rows. Cursor drivers use this to +/// amortize the cost of building and polling the async decode state machine +/// across a whole rowset rather than paying it per row. pub struct DefaultRowWriter { row: Vec, + completed: Vec>, + spare: Vec>, + string_pool: Vec>, + batching: bool, } impl DefaultRowWriter { @@ -116,9 +146,70 @@ impl DefaultRowWriter { pub fn new(col_count: usize) -> Self { Self { row: Vec::with_capacity(col_count), + completed: Vec::new(), + spare: Vec::new(), + string_pool: Vec::new(), + batching: false, + } + } + + /// Creates a writer that reuses `buffer`'s existing allocation. + pub fn with_buffer(mut buffer: Vec, col_count: usize) -> Self { + buffer.clear(); + buffer.reserve(col_count.saturating_sub(buffer.capacity())); + Self { + row: buffer, + completed: Vec::new(), + spare: Vec::new(), + string_pool: Vec::new(), + batching: false, } } + /// Creates a batch-mode writer that queues completed rows internally. + /// + /// `spare` supplies row `Vec`s the caller has finished with; they are + /// recycled instead of allocating a fresh `Vec` per row. + pub fn batching(col_count: usize, spare: Vec>) -> Self { + let mut writer = Self { + row: Vec::new(), + completed: Vec::new(), + spare, + string_pool: Vec::new(), + batching: true, + }; + writer.row = writer.next_row_buffer(col_count); + writer + } + + fn next_row_buffer(&mut self, col_count: usize) -> Vec { + match self.spare.pop() { + Some(mut buf) => { + // Reclaim the byte buffers from the row we're about to drop so + // variable-length columns can refill them instead of asking the + // allocator for a fresh one on every column of every row. + for value in buf.drain(..) { + if let ColumnValues::String(s) = value { + self.string_pool.push(s.into_bytes()); + } + } + buf.reserve(col_count.saturating_sub(buf.capacity())); + buf + } + None => Vec::with_capacity(col_count), + } + } + + /// Number of rows queued by batch mode so far. + pub fn completed_len(&self) -> usize { + self.completed.len() + } + + /// Takes the rows queued by batch mode, leaving the writer ready for reuse. + pub fn take_completed(&mut self) -> Vec> { + std::mem::take(&mut self.completed) + } + /// Takes the completed row, leaving the writer ready for reuse. pub fn take_row(&mut self) -> Vec { std::mem::take(&mut self.row) @@ -126,6 +217,10 @@ impl DefaultRowWriter { } impl RowWriter for DefaultRowWriter { + fn may_pause(&self) -> bool { + false + } + fn write_null(&mut self, _col: usize) { self.row.push(ColumnValues::Null); } @@ -223,7 +318,24 @@ impl RowWriter for DefaultRowWriter { } fn end_row(&mut self) { - // No-op for DefaultRowWriter — row is taken via take_row(). + // Non-batch mode: the row stays in place and is taken via take_row(). + if self.batching { + let col_count = self.row.len(); + let replacement = self.next_row_buffer(col_count); + let finished = std::mem::replace(&mut self.row, replacement); + self.completed.push(finished); + } + } + + fn take_string_buffer(&mut self, capacity: usize) -> Vec { + match self.string_pool.pop() { + Some(mut buf) => { + buf.clear(); + buf.reserve(capacity.saturating_sub(buf.capacity())); + buf + } + None => Vec::with_capacity(capacity), + } } } @@ -264,6 +376,81 @@ mod tests { use super::*; use crate::datatypes::sql_string::EncodingType; + #[test] + fn batching_writer_recycles_string_buffers_from_consumed_rows() { + let mut writer = DefaultRowWriter::batching(1, Vec::new()); + // No pool yet, so the buffer is freshly allocated at the asked capacity. + assert_eq!(writer.take_string_buffer(32).capacity(), 32); + + // Feed back a consumed row carrying a string allocation. + let consumed = vec![ColumnValues::String(SqlString::new( + Vec::with_capacity(128), + EncodingType::Utf16, + ))]; + writer.spare.push(consumed); + + // Recycling the row harvests the string allocation into the pool. + let recycled = writer.next_row_buffer(1); + assert!(recycled.is_empty()); + let buf = writer.take_string_buffer(16); + assert!(buf.is_empty()); + assert_eq!(buf.capacity(), 128); + } + + #[test] + fn default_row_writer_never_pauses() { + assert!(!DefaultRowWriter::new(1).may_pause()); + } + + #[test] + fn batching_writer_queues_rows_and_recycles_spares() { + let mut writer = DefaultRowWriter::batching(2, Vec::new()); + assert_eq!(writer.completed_len(), 0); + + for i in 0..3i32 { + writer.write_i32(0, i); + writer.write_null(1); + writer.end_row(); + } + assert_eq!(writer.completed_len(), 3); + + let rows = writer.take_completed(); + assert_eq!(writer.completed_len(), 0); + assert_eq!(rows.len(), 3); + for (i, row) in rows.iter().enumerate() { + assert_eq!(row.len(), 2); + assert_eq!(row[0], ColumnValues::Int(i as i32)); + assert_eq!(row[1], ColumnValues::Null); + } + } + + #[test] + fn batching_writer_reuses_spare_allocations() { + let spare = vec![Vec::with_capacity(64)]; + let spare_ptr = spare[0].as_ptr(); + + let mut writer = DefaultRowWriter::batching(2, spare); + // The first row buffer should be the recycled one, cleared but keeping + // its capacity. + assert_eq!(writer.row.capacity(), 64); + assert!(writer.row.is_empty()); + assert_eq!(writer.row.as_ptr(), spare_ptr); + + writer.write_i32(0, 7); + writer.end_row(); + // Spare pool exhausted, so the next buffer is freshly allocated. + assert_eq!(writer.take_completed().len(), 1); + } + + #[test] + fn non_batching_writer_keeps_row_in_place_on_end_row() { + let mut writer = DefaultRowWriter::new(1); + writer.write_i32(0, 1); + writer.end_row(); + assert_eq!(writer.completed_len(), 0); + assert_eq!(writer.take_row(), vec![ColumnValues::Int(1)]); + } + #[test] fn default_row_writer_assembles_column_values() { let mut writer = DefaultRowWriter::new(5); diff --git a/mssql-tds/src/datatypes/sql_string.rs b/mssql-tds/src/datatypes/sql_string.rs index 0c00ad07..cdf4f7af 100644 --- a/mssql-tds/src/datatypes/sql_string.rs +++ b/mssql-tds/src/datatypes/sql_string.rs @@ -43,6 +43,12 @@ impl SqlString { } } + /// Consumes the `SqlString`, returning its backing byte buffer so the + /// allocation can be recycled. + pub fn into_bytes(self) -> Vec { + self.bytes + } + /// Creates a UTF-16LE–encoded `SqlString` from a Rust `String`. pub fn from_utf8_string(string: String) -> Self { let utf16_bytes = string diff --git a/mssql-tds/src/datatypes/sqldatatypes.rs b/mssql-tds/src/datatypes/sqldatatypes.rs index cf55081d..b96f9a9b 100644 --- a/mssql-tds/src/datatypes/sqldatatypes.rs +++ b/mssql-tds/src/datatypes/sqldatatypes.rs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use crate::io::packet_reader::fast; + use crate::core::TdsResult; use crate::datatypes::sqltypes::get_time_length_from_scale; use crate::error::Error; @@ -726,7 +728,7 @@ where if let Ok(vdt) = var_len_type { let type_info = match vdt { VariableLengthTypes::TimeN => { - let scale = reader.read_byte().await?; + let scale = fast::read_byte(reader).await?; let length = get_time_length_from_scale(scale)? as usize; trace!( "Parsing TimeN: scale={}, calculated length={}", @@ -739,7 +741,7 @@ where } } VariableLengthTypes::DateTime2N => { - let scale = reader.read_byte().await?; + let scale = fast::read_byte(reader).await?; let time_length = get_time_length_from_scale(scale)? as usize; let length = time_length + 3; // time + 3 bytes for date TypeInfo { @@ -749,7 +751,7 @@ where } } VariableLengthTypes::DateTimeOffsetN => { - let scale = reader.read_byte().await?; + let scale = fast::read_byte(reader).await?; let time_length = get_time_length_from_scale(scale)? as usize; let length = time_length + 3 + 2; // time + 3 bytes for date + 2 bytes for offset TypeInfo { @@ -760,8 +762,8 @@ where } VariableLengthTypes::Vector => { // Vector uses USHORTLEN (u16) for max length and SCALE byte for base type - let length = reader.read_uint16().await? as usize; - let base_type_byte = reader.read_byte().await?; + let length = fast::read_uint16(reader).await? as usize; + let base_type_byte = fast::read_byte(reader).await?; // Validate base type let _base_type = VectorBaseType::try_from(base_type_byte)?; @@ -778,7 +780,7 @@ where | VariableLengthTypes::FltN | VariableLengthTypes::Guid | VariableLengthTypes::BitN => { - let length: usize = reader.read_byte().await? as usize; + let length: usize = fast::read_byte(reader).await? as usize; TypeInfo { tds_type: data_type, length, @@ -796,9 +798,9 @@ where | VariableLengthTypes::Numeric => { let len_byte_count = vdt.get_len_byte_count(); let length = match len_byte_count { - 1 => reader.read_byte().await? as usize, - 2 => reader.read_uint16().await? as usize, - 4 => reader.read_int32().await? as usize, + 1 => fast::read_byte(reader).await? as usize, + 2 => fast::read_uint16(reader).await? as usize, + 4 => fast::read_int32(reader).await? as usize, _ => { unreachable!( "Invalid tds length {:?} for type: {:?}", @@ -806,8 +808,8 @@ where ) } }; - let precision = reader.read_byte().await?; - let scale = reader.read_byte().await?; + let precision = fast::read_byte(reader).await?; + let scale = fast::read_byte(reader).await?; TypeInfo { tds_type: data_type, length, @@ -931,11 +933,11 @@ where if let Ok(pt) = plp_type { let type_info = match pt { PartialLengthType::Udt => { - let len = reader.read_uint16().await? as usize; + let len = fast::read_uint16(reader).await? as usize; let db_name = reader.read_varchar_u8_length().await?; let schema_name = reader.read_varchar_u8_length().await?; let type_name = reader.read_varchar_u8_length().await?; - // let assembly_qualified_name_length = reader.read_uint16().await? as usize; + // let assembly_qualified_name_length = fast::read_uint16(reader).await? as usize; let assembly_qualified_name = reader.read_varchar_u16_length().await?; let assembly_qualified_name: String = match assembly_qualified_name { Some(name) => name, @@ -969,7 +971,7 @@ where type_info_variant: TypeInfoVariant::PartialLen(pt, None, None, None, None), }, PartialLengthType::Xml => { - let schema_present = reader.read_byte().await?; + let schema_present = fast::read_byte(reader).await?; let db_name = if schema_present == 0x01 { Some(reader.read_varchar_u8_length().await?) } else { @@ -1030,10 +1032,10 @@ where { let len_byte_count = data_type.get_len_byte_count(); let length = match len_byte_count { - 1 => reader.read_byte().await? as usize, - 2 => reader.read_uint16().await? as usize, + 1 => fast::read_byte(reader).await? as usize, + 2 => fast::read_uint16(reader).await? as usize, 4 => { - let len_i32 = reader.read_int32().await?; + let len_i32 = fast::read_int32(reader).await?; // Negative values indicate invalid protocol data and should error out // to prevent capacity overflow from casting negative i32 to huge usize values if len_i32 < 0 { diff --git a/mssql-tds/src/io/packet_reader.rs b/mssql-tds/src/io/packet_reader.rs index 63698ad6..4bd66f20 100644 --- a/mssql-tds/src/io/packet_reader.rs +++ b/mssql-tds/src/io/packet_reader.rs @@ -45,6 +45,18 @@ pub(crate) trait TdsPacketReader { async fn skip_bytes(&mut self, skip_count: usize) -> TdsResult<()>; async fn cancel_read_stream(&mut self) -> TdsResult<()>; fn reset_reader(&mut self); + + /// Non-async fast path for fixed-width reads of at most 8 bytes. + /// + /// Returns the bytes left-aligned in the array and consumes them when the + /// current packet buffer already holds them; returns None when a refill is + /// needed so the caller falls back to the async method. Bypassing the + /// `#[async_trait]` methods here avoids a boxed future allocation for every + /// scalar read on the row-decode hot path. + fn try_take_fixed(&mut self, n: usize) -> Option<[u8; 8]> { + let _ = n; + None + } } /// Low-level TDS packet reading operations (public under `fuzzing` cfg). @@ -76,6 +88,60 @@ pub trait TdsPacketReader { async fn skip_bytes(&mut self, skip_count: usize) -> TdsResult<()>; async fn cancel_read_stream(&mut self) -> TdsResult<()>; fn reset_reader(&mut self); + + /// Non-async fast path for fixed-width reads of at most 8 bytes. + /// + /// Returns the bytes left-aligned in the array and consumes them when the + /// current packet buffer already holds them; returns None when a refill is + /// needed so the caller falls back to the async method. Bypassing the + /// `#[async_trait]` methods here avoids a boxed future allocation for every + /// scalar read on the row-decode hot path. + fn try_take_fixed(&mut self, n: usize) -> Option<[u8; 8]> { + let _ = n; + None + } +} + +/// Zero-allocation fixed-width reads. +/// +/// Each helper first tries [`TdsPacketReader::try_take_fixed`], which serves the +/// value straight from the packet buffer without constructing the boxed future +/// that `#[async_trait]` allocates for every trait method call. Only a buffer +/// refill falls back to the async method. +pub(crate) mod fast { + use super::{TdsPacketReader, TdsResult}; + use byteorder::{ByteOrder, LittleEndian}; + + macro_rules! fast_read { + ($name:ident, $ty:ty, $n:expr, $slow:ident, $decode:expr) => { + #[inline] + pub(crate) async fn $name(reader: &mut R) -> TdsResult<$ty> + where + R: TdsPacketReader + Send + Sync + ?Sized, + { + #[allow(clippy::redundant_closure_call)] + match reader.try_take_fixed($n) { + Some(bytes) => Ok($decode(&bytes[..$n])), + None => reader.$slow().await, + } + } + }; + } + + fast_read!(read_byte, u8, 1, read_byte, |b: &[u8]| b[0]); + fast_read!(read_int16, i16, 2, read_int16, LittleEndian::read_i16); + fast_read!(read_uint16, u16, 2, read_uint16, LittleEndian::read_u16); + fast_read!(read_int32, i32, 4, read_int32, LittleEndian::read_i32); + fast_read!(read_uint32, u32, 4, read_uint32, LittleEndian::read_u32); + fast_read!(read_int64, i64, 8, read_int64, LittleEndian::read_i64); + fast_read!(read_float32, f32, 4, read_float32, LittleEndian::read_f32); + fast_read!(read_float64, f64, 8, read_float64, LittleEndian::read_f64); + fast_read!(read_uint24, u32, 3, read_uint24, |b: &[u8]| { + LittleEndian::read_uint(b, 3) as u32 + }); + fast_read!(read_uint40, u64, 5, read_uint40, |b: &[u8]| { + LittleEndian::read_uint(b, 5) + }); } /// Buffered reader that reassembles TDS packets from the network stream. @@ -263,6 +329,22 @@ impl TdsPacketReader for PacketReader<'_> { Ok(result) } + fn try_take_fixed(&mut self, n: usize) -> Option<[u8; 8]> { + debug_assert!(n <= 8); + if n > 8 || !self.do_we_have_enough_data(n) { + return None; + } + let mut out = [0u8; 8]; + out[..n] + .copy_from_slice(&self.working_buffer[self.buffer_position..self.buffer_position + n]); + self.buffer_position += n; + if self.buffer_length == self.buffer_position { + self.buffer_length = 0; + self.buffer_position = 0; + } + Some(out) + } + async fn read_int16_big_endian(&mut self) -> TdsResult { if !self.do_we_have_enough_data(2) { self.read_tds_packet().await?; @@ -502,6 +584,10 @@ impl TdsPacketReader for Box { (**self).read_byte().await } + fn try_take_fixed(&mut self, n: usize) -> Option<[u8; 8]> { + (**self).try_take_fixed(n) + } + async fn read_int16_big_endian(&mut self) -> TdsResult { (**self).read_int16_big_endian().await } diff --git a/mssql-tds/src/io/token_stream.rs b/mssql-tds/src/io/token_stream.rs index 368cf39c..a20555ec 100644 --- a/mssql-tds/src/io/token_stream.rs +++ b/mssql-tds/src/io/token_stream.rs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use crate::io::packet_reader::fast; + use crate::core::{CancelHandle, TdsResult}; use crate::datatypes::decoder::{GenericDecoder, PlpColumnStream, decrypt_encrypted_column}; use crate::datatypes::row_writer::{RowWriter, write_column_value}; @@ -111,6 +113,14 @@ impl PlpPauseState { } } +/// Outcome of a batched row read: the number of rows decoded, plus the +/// non-row result that ended the batch (a token, or a mid-row pause) when the +/// batch stopped before reaching its row limit. +pub(crate) struct BatchRowsResult { + pub(crate) rows: usize, + pub(crate) stopped_at: Option, +} + #[async_trait] #[cfg(not(fuzzing))] pub(crate) trait TdsTokenStreamReader { @@ -129,6 +139,45 @@ pub(crate) trait TdsTokenStreamReader { writer: &mut (dyn RowWriter + Send), ) -> TdsResult; + /// Decodes up to `max_rows` consecutive ROW/NBCROW tokens in one call, + /// calling [`RowWriter::end_row`] after each. + /// + /// Cursor drivers use this so the cost of the boxed trait future and the + /// surrounding cancellation and timeout machinery is paid once per rowset + /// rather than once per row. The default implementation preserves exact + /// per-row behavior; transports on the hot path override it. + async fn receive_rows_into( + &mut self, + context: &ParserContext, + remaining_request_timeout: Option, + cancel_handle: Option<&CancelHandle>, + writer: &mut (dyn RowWriter + Send), + max_rows: usize, + ) -> TdsResult { + let mut rows = 0; + while rows < max_rows { + match self + .receive_row_into(context, remaining_request_timeout, cancel_handle, writer) + .await? + { + RowReadResult::RowWritten => { + writer.end_row(); + rows += 1; + } + stopped_at => { + return Ok(BatchRowsResult { + rows, + stopped_at: Some(stopped_at), + }); + } + } + } + Ok(BatchRowsResult { + rows, + stopped_at: None, + }) + } + /// Resume a paused row decode from the column after the one that triggered /// [`pause_after_column`](RowWriter::pause_after_column). /// @@ -171,6 +220,38 @@ pub trait TdsTokenStreamReader { writer: &mut (dyn RowWriter + Send), ) -> TdsResult; + async fn receive_rows_into( + &mut self, + context: &ParserContext, + remaining_request_timeout: Option, + cancel_handle: Option<&CancelHandle>, + writer: &mut (dyn RowWriter + Send), + max_rows: usize, + ) -> TdsResult { + let mut rows = 0; + while rows < max_rows { + match self + .receive_row_into(context, remaining_request_timeout, cancel_handle, writer) + .await? + { + RowReadResult::RowWritten => { + writer.end_row(); + rows += 1; + } + stopped_at => { + return Ok(BatchRowsResult { + rows, + stopped_at: Some(stopped_at), + }); + } + } + } + Ok(BatchRowsResult { + rows, + stopped_at: None, + }) + } + async fn resume_row_into( &mut self, pause_state: RowPauseState, @@ -209,7 +290,7 @@ type RowDecodeContext<'a> = (&'a [ColumnMetadata], Option<&'a Arc( registry: &impl TokenParserRegistry, context: &ParserContext, ) -> TdsResult { - let token_type_byte = reader.read_byte().await?; + let token_type_byte = fast::read_byte(reader).await?; let token_type: TokenType = token_type_byte.try_into()?; debug!( "Received token type: {:?} ({})", @@ -327,10 +408,11 @@ async fn decode_row_columns( writer: &mut (dyn RowWriter + Send), ) -> TdsResult { let decoder = GenericDecoder::default(); + let may_pause = writer.may_pause(); for (col, meta) in columns.iter().enumerate().skip(start_col) { // For PLP target columns, pause before payload consumption so callers // can stream SQLGetData-style chunks from wire. - if meta.is_plp() && writer.pause_after_column(col) { + if may_pause && meta.is_plp() && writer.pause_after_column(col) { // TODO: Add AE-aware PLP streaming path for paused row reads. // Until then, fail fast to avoid streaming ciphertext bytes to callers. if meta.crypto_metadata.is_some() { @@ -342,7 +424,7 @@ async fn decode_row_columns( ), }); } - match PlpColumnStream::begin(meta, reader).await? { + match Box::pin(PlpColumnStream::begin(meta, reader)).await? { None => { writer.write_null(col); if col + 1 < columns.len() { @@ -370,7 +452,7 @@ async fn decode_row_columns( } decode_or_decrypt_column(&decoder, reader, meta, decryptor, col, writer).await?; - if writer.pause_after_column(col) && col + 1 < columns.len() { + if may_pause && writer.pause_after_column(col) && col + 1 < columns.len() { return Ok(RowReadResult::RowPaused(RowPauseState { next_column_index: col + 1, columns: columns.to_vec(), @@ -392,11 +474,12 @@ async fn decode_nbcrow_columns( writer: &mut (dyn RowWriter + Send), ) -> TdsResult { let decoder = GenericDecoder::default(); + let may_pause = writer.may_pause(); for (col, meta) in columns.iter().enumerate().skip(start_col) { if bitmap[col / 8] & (1 << (col % 8)) != 0 { writer.write_null(col); } else { - if meta.is_plp() && writer.pause_after_column(col) { + if may_pause && meta.is_plp() && writer.pause_after_column(col) { // TODO: Add AE-aware PLP streaming path for paused row reads. // Until then, fail fast to avoid streaming ciphertext bytes to callers. if meta.crypto_metadata.is_some() { @@ -408,7 +491,7 @@ async fn decode_nbcrow_columns( ), }); } - match PlpColumnStream::begin(meta, reader).await? { + match Box::pin(PlpColumnStream::begin(meta, reader)).await? { None => { writer.write_null(col); if col + 1 < columns.len() { @@ -437,7 +520,7 @@ async fn decode_nbcrow_columns( decode_or_decrypt_column(&decoder, reader, meta, decryptor, col, writer).await?; } - if writer.pause_after_column(col) && col + 1 < columns.len() { + if may_pause && writer.pause_after_column(col) && col + 1 < columns.len() { return Ok(RowReadResult::RowPaused(RowPauseState { next_column_index: col + 1, columns: columns.to_vec(), @@ -457,24 +540,22 @@ async fn decode_or_decrypt_column( col: usize, writer: &mut (dyn RowWriter + Send), ) -> TdsResult<()> { - match (meta.crypto_metadata.is_some(), decryptor) { - (true, Some(dec)) => { - let value = decrypt_encrypted_column(decoder, reader, meta, dec).await?; + if meta.crypto_metadata.is_some() { + if let Some(dec) = decryptor { + // Boxed: Always Encrypted is a cold path but its state machine would + // otherwise be baked into this per-column future. + let value = Box::pin(decrypt_encrypted_column(decoder, reader, meta, dec)).await?; write_column_value(writer, col, value); + return Ok(()); } - (true, None) => { - tracing::info!( - column = %meta.column_name, - "Encrypted column has no column-encryption decryptor available \ - (Always Encrypted disabled for this command, or no key-store \ - provider registered); returning the raw ciphertext varbinary" - ); - decoder.decode_into(reader, meta, col, writer).await?; - } - (false, _) => { - decoder.decode_into(reader, meta, col, writer).await?; - } + tracing::info!( + column = %meta.column_name, + "Encrypted column has no column-encryption decryptor available \ + (Always Encrypted disabled for this command, or no key-store \ + provider registered); returning the raw ciphertext varbinary" + ); } + decoder.decode_into(reader, meta, col, writer).await?; Ok(()) } @@ -484,7 +565,7 @@ pub(crate) async fn receive_row_into_internal( context: &ParserContext, writer: &mut (dyn RowWriter + Send), ) -> TdsResult { - let token_type_byte = reader.read_byte().await?; + let token_type_byte = fast::read_byte(reader).await?; let token_type: TokenType = token_type_byte.try_into()?; debug!("Parsing token type: {:?}", &token_type); @@ -501,15 +582,48 @@ pub(crate) async fn receive_row_into_internal( decode_nbcrow_columns(reader, columns, decryptor, &bitmap, 0, writer).await } _ => { - let token = dispatch_token(reader, registry, token_type, context).await?; + // Boxed on purpose: `dispatch_token` inlines every token parser's + // state machine, and leaving it unboxed would make this function's + // future — constructed for every row — as large as the largest + // parser. Non-row tokens appear once per result set, so the + // allocation is negligible there. + let token = Box::pin(dispatch_token(reader, registry, token_type, context)).await?; Ok(RowReadResult::Token(token)) } } } +/// Decodes up to `max_rows` consecutive rows without leaving this future, +/// so the caller pays for the surrounding async machinery once per batch. +pub(crate) async fn receive_rows_into_internal( + reader: &mut R, + registry: &impl TokenParserRegistry, + context: &ParserContext, + writer: &mut (dyn RowWriter + Send), + max_rows: usize, +) -> TdsResult { + let mut rows = 0; + while rows < max_rows { + match receive_row_into_internal(reader, registry, context, writer).await? { + RowReadResult::RowWritten => { + writer.end_row(); + rows += 1; + } + stopped_at => { + return Ok(BatchRowsResult { + rows, + stopped_at: Some(stopped_at), + }); + } + } + } + Ok(BatchRowsResult { + rows, + stopped_at: None, + }) +} + /// Resumes a paused row decode from `pause_state.next_column_index`. -/// -/// Does not read a token-type byte — the token has already been consumed. pub(crate) async fn resume_row_into_internal( reader: &mut R, pause_state: RowPauseState,