From 801f41392dddf5a70b918eb684b656ce90dcb751 Mon Sep 17 00:00:00 2001 From: Saurabh Singh <1623701+saurabh500@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:30:18 -0700 Subject: [PATCH] Add NBCROW row throughput benchmark Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- mssql-mock-tds/src/protocol.rs | 109 +++++++++++++-- mssql-mock-tds/src/query_response.rs | 26 +++- mssql-mock-tds/src/server.rs | 69 +++++----- mssql-tds-bench/Cargo.toml | 5 + mssql-tds-bench/README.md | 8 +- mssql-tds-bench/benches/nbcrow_rows.rs | 179 +++++++++++++++++++++++++ 6 files changed, 346 insertions(+), 50 deletions(-) create mode 100644 mssql-tds-bench/benches/nbcrow_rows.rs diff --git a/mssql-mock-tds/src/protocol.rs b/mssql-mock-tds/src/protocol.rs index 850e8bbe..9fa79c2f 100644 --- a/mssql-mock-tds/src/protocol.rs +++ b/mssql-mock-tds/src/protocol.rs @@ -150,6 +150,7 @@ impl PacketHeader { pub enum TokenType { ColMetadata = 0x81, Row = 0xD1, + NbcRow = 0xD2, Done = 0xFD, DoneProc = 0xFE, DoneInProc = 0xFF, @@ -873,7 +874,7 @@ pub fn build_query_result(response: &crate::query_response::QueryResponse) -> By // Serialize each column for col in &response.columns { result.put_u32_le(0); // UserType - result.put_u16_le(0x0000); // Flags: not nullable, no special flags + result.put_u16_le(u16::from(response.use_nbc_rows)); // nullable for NBCROW result.put_u8(col.data_type.tds_type_code()); if col.data_type == crate::query_response::SqlDataType::NVarChar { // Required to support string responses (e.g., @@USERAGENT). @@ -899,9 +900,23 @@ pub fn build_query_result(response: &crate::query_response::QueryResponse) -> By // Serialize each row for row in &response.rows { - result.put_u8(TokenType::Row as u8); - for value in &row.values { - value.write_to_buffer(&mut result); + if response.use_nbc_rows { + result.put_u8(TokenType::NbcRow as u8); + let mut bitmap = vec![0u8; response.columns.len().div_ceil(8)]; + for (index, value) in row.values.iter().enumerate() { + if value.is_null() { + bitmap[index / 8] |= 1 << (index % 8); + } + } + result.extend_from_slice(&bitmap); + for value in row.values.iter().filter(|value| !value.is_null()) { + value.write_to_buffer(&mut result); + } + } else { + result.put_u8(TokenType::Row as u8); + for value in &row.values { + value.write_to_buffer(&mut result); + } } } @@ -959,16 +974,31 @@ pub fn build_error_response(message: &str) -> BytesMut { /// Wrap token data in a TDS packet fn wrap_in_packet(packet_type: PacketType, data: BytesMut) -> BytesMut { - let total_length = (PACKET_HEADER_SIZE + data.len()) as u16; + let max_payload = MAX_PACKET_SIZE - PACKET_HEADER_SIZE; + let packet_count = data.len().max(1).div_ceil(max_payload); + let mut packet = BytesMut::with_capacity(data.len() + packet_count * PACKET_HEADER_SIZE); + + if data.is_empty() { + PacketHeader::new(packet_type, PACKET_HEADER_SIZE as u16, 1).write(&mut packet); + return packet; + } - let mut packet = BytesMut::with_capacity(total_length as usize); - let header = PacketHeader::new(packet_type, total_length, 1); - header.write(&mut packet); - packet.extend_from_slice(&data); + for (index, chunk) in data.chunks(max_payload).enumerate() { + let total_length = (PACKET_HEADER_SIZE + chunk.len()) as u16; + let mut header = + PacketHeader::new(packet_type, total_length, (index as u8).wrapping_add(1)); + if index + 1 < packet_count { + header.status = PacketStatus::normal(); + } + header.write(&mut packet); + packet.extend_from_slice(chunk); + } trace!( - "Built packet: type={:?}, length={}", - packet_type, total_length + "Built packet stream: type={:?}, packets={}, length={}", + packet_type, + packet_count, + packet.len() ); packet } @@ -1136,6 +1166,63 @@ mod tests { assert!(response.len() >= PACKET_HEADER_SIZE); } + #[test] + fn test_nbcrow_result_serializes_bitmap_and_omits_null_values() { + use crate::query_response::{ + ColumnDefinition, ColumnValue, QueryResponse, Row, SqlDataType, + }; + + let columns = (0..9) + .map(|_| ColumnDefinition::new("", SqlDataType::Int)) + .collect(); + let values = (0..9) + .map(|index| { + if index == 0 || index == 8 { + ColumnValue::Null + } else { + ColumnValue::Int(index) + } + }) + .collect(); + let response = build_query_result( + &QueryResponse::new(columns, vec![Row::new(values)]).with_nbc_rows(), + ); + + let row_offset = PACKET_HEADER_SIZE + 3 + (9 * 9); + assert_eq!(response[row_offset], TokenType::NbcRow as u8); + assert_eq!(&response[row_offset + 1..row_offset + 3], &[0x01, 0x01]); + assert_eq!(response[row_offset + 3], 4); + assert_eq!( + i32::from_le_bytes( + response[row_offset + 4..row_offset + 8] + .try_into() + .expect("first non-null integer should have four bytes") + ), + 1 + ); + } + + #[test] + fn test_large_payload_is_split_into_tds_packets() { + let response = wrap_in_packet( + PacketType::TabularResult, + BytesMut::from(&vec![0xAA; MAX_PACKET_SIZE][..]), + ); + + let mut first = &response[..PACKET_HEADER_SIZE]; + let first_header = + PacketHeader::parse(&mut first).expect("first packet header should parse"); + assert_eq!(usize::from(first_header.length), MAX_PACKET_SIZE); + assert!(!first_header.status.is_end_of_message()); + + let second_offset = usize::from(first_header.length); + let mut second = &response[second_offset..second_offset + PACKET_HEADER_SIZE]; + let second_header = + PacketHeader::parse(&mut second).expect("second packet header should parse"); + assert!(second_header.status.is_end_of_message()); + assert_eq!(usize::from(second_header.length), PACKET_HEADER_SIZE * 2); + } + #[test] fn test_routing_envchange_token() { let token = build_routing_envchange_token("sqlserver.database.windows.net", 1433); diff --git a/mssql-mock-tds/src/query_response.rs b/mssql-mock-tds/src/query_response.rs index 786d1516..03b7e5de 100644 --- a/mssql-mock-tds/src/query_response.rs +++ b/mssql-mock-tds/src/query_response.rs @@ -3,7 +3,7 @@ //! Query response definitions for the mock TDS server -use bytes::{BufMut, BytesMut}; +use bytes::{BufMut, Bytes, BytesMut}; use std::collections::HashMap; /// SQL data types supported by the mock server @@ -57,6 +57,10 @@ pub enum ColumnValue { } impl ColumnValue { + pub(crate) fn is_null(&self) -> bool { + matches!(self, Self::Null) + } + /// Get the SQL data type for this value pub fn data_type(&self) -> SqlDataType { match self { @@ -193,6 +197,7 @@ pub struct QueryResponse { /// An error emitted (with a DONE MORE token) before the result set, so the /// server keeps streaming the row set after a statement-scoped error. pub leading_error: Option, + pub(crate) use_nbc_rows: bool, } impl QueryResponse { @@ -203,6 +208,7 @@ impl QueryResponse { rows, info_tokens: Vec::new(), leading_error: None, + use_nbc_rows: false, } } @@ -218,6 +224,12 @@ impl QueryResponse { self } + /// Encode rows as NBCROW tokens with a null bitmap. + pub fn with_nbc_rows(mut self) -> Self { + self.use_nbc_rows = true; + self + } + /// Helper to create a response for SELECT 1 pub fn select_one() -> Self { Self { @@ -225,6 +237,7 @@ impl QueryResponse { rows: vec![Row::new(vec![ColumnValue::Int(1)])], info_tokens: Vec::new(), leading_error: None, + use_nbc_rows: false, } } @@ -243,6 +256,7 @@ impl QueryResponse { ])], info_tokens: Vec::new(), leading_error: None, + use_nbc_rows: false, } } } @@ -250,6 +264,7 @@ impl QueryResponse { /// Registry of query responses pub struct QueryRegistry { responses: HashMap, + wire_responses: HashMap, } impl QueryRegistry { @@ -257,6 +272,7 @@ impl QueryRegistry { pub fn new() -> Self { let mut registry = Self { responses: HashMap::new(), + wire_responses: HashMap::new(), }; // Add default responses @@ -272,13 +288,19 @@ impl QueryRegistry { /// Register a query response pub fn register(&mut self, query: impl Into, response: QueryResponse) { let query = query.into().to_uppercase(); - self.responses.insert(query, response); + let wire_response = crate::protocol::build_query_result(&response).freeze(); + self.responses.insert(query.clone(), response); + self.wire_responses.insert(query, wire_response); } /// Get a response for a query pub fn get(&self, query: &str) -> Option<&QueryResponse> { self.responses.get(&query.to_uppercase()) } + + pub(crate) fn get_wire_response(&self, query: &str) -> Option { + self.wire_responses.get(&query.to_uppercase()).cloned() + } } impl Default for QueryRegistry { diff --git a/mssql-mock-tds/src/server.rs b/mssql-mock-tds/src/server.rs index 4f3f4339..e0f520ef 100644 --- a/mssql-mock-tds/src/server.rs +++ b/mssql-mock-tds/src/server.rs @@ -7,11 +7,11 @@ use crate::protocol::{ PACKET_HEADER_SIZE, PacketHeader, PacketType, ProtocolError, build_done_token, build_error_response, build_feature_ext_ack_fedauth, build_fedauth_challenge_response, build_login_ack, build_prelogin_response, build_prelogin_response_with_fedauth, - build_query_result, build_routing_response, build_transaction_manager_response, - parse_fedauth_token, parse_login7_auth, parse_sql_batch, parse_transaction_manager_request, + build_routing_response, build_transaction_manager_response, parse_fedauth_token, + parse_login7_auth, parse_sql_batch, parse_transaction_manager_request, }; use crate::query_response::QueryRegistry; -use bytes::BytesMut; +use bytes::{Bytes, BytesMut}; use native_tls::Identity; use std::collections::BTreeMap; use std::net::SocketAddr; @@ -179,7 +179,7 @@ impl ConnectionProcessor { } /// Process a single packet from the buffer and return the response - pub async fn process_packet(&mut self) -> Result, ProtocolError> { + pub async fn process_packet(&mut self) -> Result, ProtocolError> { if self.buffer.len() < PACKET_HEADER_SIZE { return Ok(None); } @@ -241,10 +241,7 @@ impl ConnectionProcessor { "Redirecting client {} to {}:{}", self.addr, redir.redirect_host, redir.redirect_port ); - Some(build_routing_response( - &redir.redirect_host, - redir.redirect_port, - )) + Some(build_routing_response(&redir.redirect_host, redir.redirect_port).freeze()) } else if auth_info.has_fedauth && auth_info.access_token_bytes.is_none() { self.awaiting_fedauth_token = true; self.is_authenticated = false; @@ -252,10 +249,13 @@ impl ConnectionProcessor { "FedAuth login without inline token from {}; sending challenge", self.addr ); - Some(build_fedauth_challenge_response( - FEDAUTH_CHALLENGE_STS_URL, - FEDAUTH_CHALLENGE_SPN, - )) + Some( + build_fedauth_challenge_response( + FEDAUTH_CHALLENGE_STS_URL, + FEDAUTH_CHALLENGE_SPN, + ) + .freeze(), + ) } else { self.awaiting_fedauth_token = false; self.is_authenticated = true; @@ -300,7 +300,7 @@ impl ConnectionProcessor { packet.extend_from_slice(&response); self.record_to_store().await; - Some(packet) + Some(packet.freeze()) } } @@ -310,7 +310,7 @@ impl ConnectionProcessor { "Received unexpected FedAuthToken packet from {} without challenge", self.addr ); - Some(build_error_response("Unexpected FedAuth token")) + Some(build_error_response("Unexpected FedAuth token").freeze()) } else { debug!("Handling FedAuthToken from {}", self.addr); let packet_body = &packet_data[PACKET_HEADER_SIZE..]; @@ -337,12 +337,15 @@ impl ConnectionProcessor { packet.extend_from_slice(&response); self.record_to_store().await; - Some(packet) + Some(packet.freeze()) } Err(e) => { warn!("Failed to parse FedAuthToken from {}: {}", self.addr, e); self.awaiting_fedauth_token = false; - Some(build_error_response(&format!("FedAuth parse error: {}", e))) + Some( + build_error_response(&format!("FedAuth parse error: {}", e)) + .freeze(), + ) } } } @@ -354,7 +357,7 @@ impl ConnectionProcessor { "Received SQL batch from {} before authentication", self.addr ); - Some(build_error_response("Not authenticated")) + Some(build_error_response("Not authenticated").freeze()) } else { debug!("Handling SQL batch from {}", self.addr); @@ -368,10 +371,8 @@ impl ConnectionProcessor { // Look up query in registry let registry = self.query_registry.lock().await; - if let Some(response_data) = registry.get(&sql) { + if let Some(packet) = registry.get_wire_response(&sql) { info!("Found registered response for query"); - // build_query_result already wraps in a packet, so return directly - let packet = build_query_result(response_data); Some(packet) } else { info!("No registered response, returning empty result"); @@ -385,12 +386,12 @@ impl ConnectionProcessor { resp_header.write(&mut packet); packet.extend_from_slice(&response); - Some(packet) + Some(packet.freeze()) } } Err(e) => { warn!("Failed to parse SQL batch from {}: {}", self.addr, e); - Some(build_error_response(&format!("Parse error: {}", e))) + Some(build_error_response(&format!("Parse error: {}", e)).freeze()) } } } @@ -407,7 +408,7 @@ impl ConnectionProcessor { resp_header.write(&mut packet); packet.extend_from_slice(&response); - Some(packet) + Some(packet.freeze()) } PacketType::TransactionManager => { @@ -430,7 +431,7 @@ impl ConnectionProcessor { resp_header.write(&mut packet); packet.extend_from_slice(&tokens); - Some(packet) + Some(packet.freeze()) } _ => { @@ -1191,7 +1192,7 @@ async fn handle_connection( let response = match header.packet_type { PacketType::PreLogin => { debug!("Handling PreLogin"); - Some(build_prelogin_response()) + Some(build_prelogin_response().freeze()) } PacketType::Login7 => { @@ -1211,13 +1212,13 @@ async fn handle_connection( resp_header.write(&mut packet); packet.extend_from_slice(&response); - Some(packet) + Some(packet.freeze()) } PacketType::SqlBatch => { if !is_authenticated { warn!("Received SQL batch before authentication"); - Some(build_error_response("Not authenticated")) + Some(build_error_response("Not authenticated").freeze()) } else { debug!("Handling SQL batch"); @@ -1231,8 +1232,8 @@ async fn handle_connection( // Look up query in registry let registry = query_registry.lock().await; - if let Some(response) = registry.get(&sql) { - Some(build_query_result(response)) + if let Some(response) = registry.get_wire_response(&sql) { + Some(response) } else if sql.to_uppercase().starts_with("SELECT") { // Return empty result set with DONE for unknown SELECT queries let mut response = BytesMut::new(); @@ -1248,7 +1249,7 @@ async fn handle_connection( resp_header.write(&mut packet); packet.extend_from_slice(&response); - Some(packet) + Some(packet.freeze()) } else { // For other commands, just return DONE let response = build_done_token(0); @@ -1263,12 +1264,12 @@ async fn handle_connection( resp_header.write(&mut packet); packet.extend_from_slice(&response); - Some(packet) + Some(packet.freeze()) } } Err(e) => { error!("Failed to parse SQL batch: {}", e); - Some(build_error_response("Failed to parse SQL")) + Some(build_error_response("Failed to parse SQL").freeze()) } } } @@ -1285,7 +1286,7 @@ async fn handle_connection( resp_header.write(&mut packet); packet.extend_from_slice(&response); - Some(packet) + Some(packet.freeze()) } PacketType::TransactionManager => { @@ -1304,7 +1305,7 @@ async fn handle_connection( resp_header.write(&mut packet); packet.extend_from_slice(&tokens); - Some(packet) + Some(packet.freeze()) } _ => { diff --git a/mssql-tds-bench/Cargo.toml b/mssql-tds-bench/Cargo.toml index f2d0d178..4188fbb3 100644 --- a/mssql-tds-bench/Cargo.toml +++ b/mssql-tds-bench/Cargo.toml @@ -29,6 +29,7 @@ workspace = true # across both runs, so any statistically significant delta is attributable to # mssql-tds itself. mssql-tds = { path = "../mssql-tds", features = ["test-util"] } +mssql-mock-tds = { path = "../mssql-mock-tds" } async-trait = "0.1.89" criterion = "0.5" @@ -62,3 +63,7 @@ harness = false [[bench]] name = "tds_specific" harness = false + +[[bench]] +name = "nbcrow_rows" +harness = false diff --git a/mssql-tds-bench/README.md b/mssql-tds-bench/README.md index 827c7842..da753dbf 100644 --- a/mssql-tds-bench/README.md +++ b/mssql-tds-bench/README.md @@ -35,10 +35,11 @@ that file, so every move is reviewed and recorded in git history. | `datatypes` | primitives; VARCHAR/NVARCHAR; temporal types; LOB (1 MB/20 MB, byte throughput) | | `bulk` | `BulkCopy` insert across batch sizes 500/5,000 | | `tds_specific` | packet-size sensitivity (4096/8192/32768); zero-copy `next_row` row-iteration throughput | +| `nbcrow_rows` | end-to-end NBCROW decode throughput for 4/16/64-column fixed-width, text-heavy, and mixed rows with 25% NULLs | -Each benchmark creates its own session temp tables / temp procedures, so **no -server-side setup script is required** — the objects vanish when the connection -closes. +The SQL Server-backed benchmarks create their own session temp objects, while +`nbcrow_rows` starts an in-process mock server with pre-serialized responses. +No server-side setup script is required. ## Running locally @@ -81,6 +82,7 @@ Criterion tuning knobs (defaults chosen for noisy, network-bound runs): | `BENCH_NOISE` | `0.05` | noise threshold | | `BENCH_BULK_ROWS` | `10000` | rows for the bulk-insert bench | | `BENCH_ITER_ROWS` | `50000` | rows for the row-iteration bench | +| `BENCH_NBC_ROWS` | `10000` | rows per result set for each NBCROW throughput scenario | CPU pinning (used by `perf-lab/run-benchmarks.sh` when SQL Server is colocated): diff --git a/mssql-tds-bench/benches/nbcrow_rows.rs b/mssql-tds-bench/benches/nbcrow_rows.rs new file mode 100644 index 00000000..ed57c9f1 --- /dev/null +++ b/mssql-tds-bench/benches/nbcrow_rows.rs @@ -0,0 +1,179 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! End-to-end NBCROW row-decode throughput against `mssql-mock-tds`. +//! +//! The mock pre-serializes each response before the timed loop, so the measured +//! path is query dispatch, loopback transport, and complete client row decoding. + +use std::env; +use std::net::SocketAddr; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use mssql_mock_tds::{ + ColumnDefinition, ColumnValue, MockTdsServer, QueryResponse, Row, SqlDataType, +}; +use mssql_tds::{ + connection::{client_context::ClientContext, tds_client::TdsClient}, + connection_provider::tds_connection_provider::TdsConnectionProvider, + core::{EncryptionOptions, EncryptionSetting}, +}; +use mssql_tds_bench::{criterion_config, drain, runtime}; +use tokio::sync::oneshot; + +const QUERY: &str = "SELECT nbcrow_benchmark"; +const NULL_PERCENT: usize = 25; + +#[derive(Clone, Copy)] +enum RowShape { + FixedWidth, + TextHeavy, + Mixed, +} + +impl RowShape { + fn name(self) -> &'static str { + match self { + Self::FixedWidth => "fixed_width", + Self::TextHeavy => "text_heavy", + Self::Mixed => "mixed", + } + } + + fn data_type(self, column: usize) -> SqlDataType { + match self { + Self::FixedWidth => { + if column.is_multiple_of(2) { + SqlDataType::Int + } else { + SqlDataType::BigInt + } + } + Self::TextHeavy => SqlDataType::NVarChar, + Self::Mixed => match column % 4 { + 0 | 3 => SqlDataType::Int, + 1 => SqlDataType::BigInt, + 2 => SqlDataType::NVarChar, + _ => unreachable!(), + }, + } + } +} + +fn response(shape: RowShape, column_count: usize, row_count: usize) -> QueryResponse { + let columns: Vec<_> = (0..column_count) + .map(|column| ColumnDefinition::new(format!("c{column}"), shape.data_type(column))) + .collect(); + let rows = (0..row_count) + .map(|row| { + let values = (0..column_count) + .map(|column| { + if (row + column).is_multiple_of(100 / NULL_PERCENT) { + return ColumnValue::Null; + } + + match shape.data_type(column) { + SqlDataType::Int => ColumnValue::Int(row as i32 ^ column as i32), + SqlDataType::BigInt => { + ColumnValue::BigInt((row as i64) << 32 | column as i64) + } + SqlDataType::NVarChar => { + ColumnValue::NVarChar(format!("row_{row}_column_{column}")) + } + SqlDataType::TinyInt | SqlDataType::SmallInt => unreachable!(), + } + }) + .collect(); + Row::new(values) + }) + .collect(); + + QueryResponse::new(columns, rows).with_nbc_rows() +} + +async fn connect_mock(address: SocketAddr) -> TdsClient { + let mut context = ClientContext::default(); + context.user_name = "sa".to_string(); + context.password = "MockBenchmark1!".to_string(); + context.database = "master".to_string(); + context.encryption_options = EncryptionOptions { + mode: EncryptionSetting::PreferOff, + trust_server_certificate: true, + host_name_in_cert: None, + server_certificate: None, + }; + + TdsConnectionProvider {} + .create_client( + context, + &format!("tcp:{},{}", address.ip(), address.port()), + None, + ) + .await + .expect("failed to connect to mock TDS server") +} + +fn nbcrow_row_throughput(c: &mut Criterion) { + let row_count = env::var("BENCH_NBC_ROWS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(10_000usize); + let rt = runtime(); + let server = rt + .block_on(MockTdsServer::new("127.0.0.1:0")) + .expect("failed to start mock TDS server"); + let address = server.local_addr(); + let registry = server.query_registry(); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let server_handle = rt.spawn(server.run_with_shutdown(shutdown_rx)); + let mut client = rt.block_on(connect_mock(address)); + + let mut group = c.benchmark_group("nbcrow_row_throughput"); + group.throughput(Throughput::Elements(row_count as u64)); + + for shape in [RowShape::FixedWidth, RowShape::TextHeavy, RowShape::Mixed] { + for column_count in [4usize, 16, 64] { + rt.block_on(async { + registry + .lock() + .await + .register(QUERY, response(shape, column_count, row_count)); + }); + + group.bench_with_input( + BenchmarkId::new( + shape.name(), + format!("{column_count}_columns_{NULL_PERCENT}_percent_null"), + ), + &column_count, + |b, _| { + b.iter(|| { + rt.block_on(async { + client + .execute(QUERY.to_string(), ()) + .await + .expect("execute failed"); + assert_eq!(drain(&mut client).await, row_count as u64); + }); + }); + }, + ); + } + } + + group.finish(); + drop(client); + shutdown_tx + .send(()) + .expect("mock TDS server stopped before benchmark shutdown"); + rt.block_on(server_handle) + .expect("mock TDS server task failed") + .expect("mock TDS server shutdown failed"); +} + +criterion_group! { + name = benches; + config = criterion_config(); + targets = nbcrow_row_throughput +} +criterion_main!(benches);