Skip to content

Add experimental TDS performance baseline #303

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Saurabh Singh (saurabh500) wants to merge 1 commit into from
Draft

Add experimental TDS performance baseline #303

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
195 changes: 136 additions & 59 deletions mssql-tds/src/connection/tds_client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,8 @@ use crate::{
message::{batch::SqlBatch, messages::Request},
token::tokens::{ColMetadataToken, CurrentCommand, DoneStatus, EnvChangeTokenSubType, Tokens},
};
use async_trait::async_trait;
use std::collections::HashMap;
use std::future::Future;
use std::num::NonZeroU32;
use tracing::{debug, error, info, instrument, warn};

Expand Down Expand Up @@ -2781,15 +2781,16 @@ 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")]
// Not instrumented: the span pushes ResultSet::next_row over the 4 KiB
// hot-path future budget. Successful rows still emit `Row Received`.
pub(crate) async fn get_next_row(&mut self) -> TdsResult<Option<Vec<ColumnValues>>> {
let col_count = self
.current_metadata
.as_ref()
.map(|m| m.columns.len())
.unwrap_or(0);
let mut writer = DefaultRowWriter::new(col_count);
if self.get_next_row_into(&mut writer).await? {
if self.next_row_into(&mut writer).await? {
Ok(Some(writer.take_row()))
} else {
Ok(None)
Expand Down Expand Up @@ -3587,24 +3588,25 @@ 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),
) -> TdsResult<bool> {
/// Concrete writers stay concrete through the production transport and
/// decode chain. [`ResultSet::next_row_into`] provides the same operation
/// through statically dispatched trait calls.
// `#[instrument]` adds enough state to exceed the 4096 B budget once the
// lazy timeout future is inlined. Successful rows still emit `Row Received`.
pub async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult<bool> {
// Every error return below must abort the pending prepare capture. A
// wrapper that centralizes this cleanup exceeds the row-future size budget.
// End-of-set reads are idempotent even after advancing clears metadata.
if self.current_result_set_has_been_read_till_end {
return Ok(false);
}
if self.current_metadata.is_none() {
self.abort_pending_prepare_capture();
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(),
));
}

// Idempotent at end-of-set: after the terminating DONE, calling again
// must report exhaustion instead of blocking on a wire read for a
// packet the server will never send (until the caller advances).
if self.current_result_set_has_been_read_till_end {
return Ok(false);
}

// The push path decodes whole rows and never pauses, so it must not run
// while the pull cursor has a row *partially* read: silently draining it
// would discard that row and return the *next* one, mapping the caller's
Expand All @@ -3615,21 +3617,26 @@ impl TdsClient {
self.active_row_read_state,
ActiveRowReadState::RowPaused(_) | ActiveRowReadState::PlpPaused(_)
) {
self.abort_pending_prepare_capture();
return Err(UsageError(
"get_next_row_into called while a pull-cursor row is still active; \
"next_row_into called while a pull-cursor row is still active; \
advance the cursor with next_row_cursor before using the push row API"
.to_string(),
));
}

self.drain_active_row().await?;

let metadata = Arc::clone(self.current_metadata.as_ref().unwrap());
let decryptor = self.resolve_cell_decryptor(&metadata).await?;
let decryptor = match self.resolve_cell_decryptor(&metadata).await {
Ok(decryptor) => decryptor,
Err(error) => {
self.abort_pending_prepare_capture();
return Err(error);
}
};
let parser_context = ParserContext::ColumnMetadata(metadata, decryptor);
loop {
let start = Instant::now();
let result = self
let result = match self
.transport
.receive_row_into(
&parser_context,
Expand All @@ -3638,7 +3645,14 @@ impl TdsClient {
ColumnPolicy::DecodeAll,
writer,
)
.await?;
.await
{
Ok(result) => result,
Err(error) => {
self.abort_pending_prepare_capture();
return Err(error);
}
};
self.update_remaining_timeout(start);

match result {
Expand All @@ -3649,13 +3663,21 @@ impl TdsClient {
}
RowReadResult::RowPaused(_) | RowReadResult::PlpPaused(_) => {
// DecodeAll never pauses; a pause here is a protocol/logic error.
self.abort_pending_prepare_capture();
return Err(crate::error::Error::ProtocolError(
"Unexpected pause while decoding a full row (ColumnPolicy::DecodeAll)"
.to_string(),
));
}
RowReadResult::Token(token) => {
if let Some(has_row) = self.handle_row_read_token(token).await? {
let handled = match Box::pin(self.handle_row_read_token(token)).await {
Ok(handled) => handled,
Err(error) => {
self.abort_pending_prepare_capture();
return Err(error);
}
};
if let Some(has_row) = handled {
return Ok(has_row);
}
}
Expand Down Expand Up @@ -3870,7 +3892,7 @@ impl TdsClient {
/// The scratch buffer is heap-allocated rather than a stack array: it is live
/// across the await below, so a stack array would be stored inline in this
/// future and propagate into every caller that awaits it — `read_row_column`
/// directly, plus `drain_rows`, `get_next_row_into` and `next_row_cursor` via
/// directly, plus `drain_rows`, `next_row_into` and `next_row_cursor` via
/// `drain_active_row`. Abandoning a partially read PLP column is rare and
/// already network-bound, so one allocation there is negligible; an 8 KiB
/// per-row state machine is not.
Expand Down Expand Up @@ -4450,7 +4472,6 @@ impl TdsClient {
}
}

#[async_trait]
impl ResultSet for TdsClient {
fn get_metadata(&self) -> &Vec<ColumnMetadata> {
// If no metadata is available, return an empty vector
Expand All @@ -4462,39 +4483,23 @@ impl ResultSet for TdsClient {
.unwrap_or(&self.empty_metadata)
}

#[instrument(skip(self), level = "info")]
async fn next_row(&mut self) -> TdsResult<Option<Vec<ColumnValues>>> {
let result = if self.maybe_has_unread_rows() {
self.get_next_row().await
} else {
Ok(None)
};
if result.is_err() {
self.abort_pending_prepare_capture();
}
result
fn next_row(&mut self) -> impl Future<Output = TdsResult<Option<Vec<ColumnValues>>>> + Send {
self.get_next_row()
}

#[instrument(skip(self, writer), level = "info")]
async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult<bool> {
let result = if self.maybe_has_unread_rows() {
self.get_next_row_into(writer).await
} else {
Ok(false)
};
if result.is_err() {
self.abort_pending_prepare_capture();
}
result
fn next_row_into(
&mut self,
writer: &mut (dyn RowWriter + Send),
) -> impl Future<Output = TdsResult<bool>> + Send {
TdsClient::next_row_into(self, writer)
}

fn maybe_has_unread_rows(&self) -> bool {
!self.current_result_set_has_been_read_till_end
}

#[instrument(skip(self), level = "info")]
async fn close(&mut self) -> TdsResult<()> {
self.close_query().await
fn close(&mut self) -> impl Future<Output = TdsResult<()>> + Send {
self.close_query()
}
}

Expand Down Expand Up @@ -4686,16 +4691,26 @@ enum ResultBoundaryKind {
End,
}

/// Async result set iteration.
#[async_trait]
/// Async result set iteration through statically dispatched futures.
///
/// The returned futures are native, unboxed futures with an explicit [`Send`]
/// guarantee.
///
/// # Dyn compatibility
///
/// This trait is intentionally not dyn-compatible and cannot be used through
/// `dyn ResultSet`. This is a breaking change for trait-object consumers and for
/// implementations written with `#[async_trait]`; concrete call sites can keep
/// awaiting the methods unchanged, while implementations must return native
/// `Send` futures.
pub trait ResultSet {
/// Returns the metadata of the result set.
/// This metadata includes information about the columns in the result set.
fn get_metadata(&self) -> &Vec<ColumnMetadata>;

/// Returns the next row of data as a vector of column values.
/// If there is no more data, it returns None.
async fn next_row(&mut self) -> TdsResult<Option<Vec<ColumnValues>>>;
fn next_row(&mut self) -> impl Future<Output = TdsResult<Option<Vec<ColumnValues>>>> + Send;

/// Decodes the next row directly into a [`RowWriter`], returning `true` if
/// a row was written or `false` when the result set is exhausted.
Expand All @@ -4710,14 +4725,17 @@ pub trait ResultSet {
/// partially read. Draining that row here would silently discard it and
/// return the *next* one, so callers must first finish the row with
/// `next_row_cursor`. A fully-consumed or absent row is fine.
async fn next_row_into(&mut self, writer: &mut (dyn RowWriter + Send)) -> TdsResult<bool>;
fn next_row_into(
&mut self,
writer: &mut (dyn RowWriter + Send),
) -> impl Future<Output = TdsResult<bool>> + Send;

/// Returns `true` if the result set may still contain unread rows.
fn maybe_has_unread_rows(&self) -> bool;

/// Iterates over the result set, and marks it as closed. After calling close, the next_row method,
/// will always return None.
async fn close(&mut self) -> TdsResult<()>;
fn close(&mut self) -> impl Future<Output = TdsResult<()>> + Send;
}

#[cfg(test)]
Expand Down Expand Up @@ -5046,22 +5064,53 @@ mod tests {
let next_row_cursor = std::mem::size_of_val(&client.next_row_cursor());
let read_row_column = std::mem::size_of_val(&client.read_row_column(0));
let drain_rows = std::mem::size_of_val(&client.drain_rows());
let get_next_row_into = std::mem::size_of_val(&client.get_next_row_into(&mut sink));
let next_row_into = std::mem::size_of_val(&client.next_row_into(&mut sink));
let next_row_into_dyn =
std::mem::size_of_val(&client.next_row_into(&mut sink as &mut (dyn RowWriter + Send)));
let read_active_plp_chunk =
std::mem::size_of_val(&client.read_active_plp_chunk(&mut plp_out));

for (name, size) in [
("next_row_cursor", next_row_cursor),
("read_row_column", read_row_column),
("drain_rows", drain_rows),
("get_next_row_into", get_next_row_into),
("next_row_into", next_row_into),
("next_row_into (dyn writer)", next_row_into_dyn),
("read_active_plp_chunk", read_active_plp_chunk),
] {
assert!(
size <= MAX,
"{name} future is {size} B, expected <= {MAX} B"
);
}

let native_next_row = std::mem::size_of_val(&client.get_next_row());
let result_set_next_row = std::mem::size_of_val(&ResultSet::next_row(&mut client));
let native_next_row_into =
std::mem::size_of_val(&client.next_row_into(&mut sink as &mut (dyn RowWriter + Send)));
let result_set_next_row_into =
std::mem::size_of_val(&ResultSet::next_row_into(&mut client, &mut sink));

assert_eq!(
result_set_next_row, native_next_row,
"ResultSet::next_row must forward the native future without boxing"
);
assert_eq!(
result_set_next_row_into, native_next_row_into,
"ResultSet::next_row_into must forward the native future without boxing"
);

// `close` is not checked against the per-row budget because it runs only
// once per result set; its larger future does not affect row iteration.
for (name, size) in [
("ResultSet::next_row", result_set_next_row),
("ResultSet::next_row_into", result_set_next_row_into),
] {
assert!(
size <= MAX,
"{name} future is {size} B, expected <= {MAX} B"
);
}
}

#[test]
Expand Down Expand Up @@ -5448,7 +5497,7 @@ mod tests {
}

#[tokio::test]
async fn get_next_row_into_rejects_active_pull_cursor_row() {
async fn next_row_into_rejects_active_pull_cursor_row() {
// A row parked by the pull cursor (`next_row_cursor`) must not be
// silently drained by the push path. Mixing the two would discard the
// parked row and hand back the *next* one, so the earlier
Expand All @@ -5465,7 +5514,7 @@ mod tests {

let mut sink = DiscardRowWriter;
let err = client
.get_next_row_into(&mut sink)
.next_row_into(&mut sink)
.await
.expect_err("push path must reject a parked pull-cursor row");
assert!(
Expand Down Expand Up @@ -5497,6 +5546,19 @@ mod tests {
assert_eq!(TdsClient::normalize_param_name(""), "");
}

#[tokio::test]
async fn next_row_into_is_idempotent_after_end_without_metadata() {
let mut client = create_test_client();
client.current_metadata = None;
client.current_result_set_has_been_read_till_end = true;
let mut sink = DiscardRowWriter;

assert!(
!ResultSet::next_row_into(&mut client, &mut sink)
.await
.unwrap()
);
}
#[tokio::test]
async fn consume_done_token_captures_all_info_tokens() {
let mut client = create_test_client_with_tokens(vec![
Expand Down Expand Up @@ -6686,7 +6748,22 @@ mod tests {
client.current_result_set_has_been_read_till_end = false;
client.pending_capture = Some(sid(1));

assert!(client.next_row().await.is_err());
assert!(ResultSet::next_row(&mut client).await.is_err());
assert!(client.pending_capture.is_none());
}

#[tokio::test]
async fn next_row_into_error_aborts_pending_prepare_capture() {
let mut client = create_test_client();
client.current_result_set_has_been_read_till_end = false;
client.pending_capture = Some(sid(1));
let mut writer = DiscardRowWriter;

assert!(
ResultSet::next_row_into(&mut client, &mut writer)
.await
.is_err()
);
assert!(client.pending_capture.is_none());
}

Expand Down
Loading
Loading