From f177f879dfb83d8a8d8276f02ae61859691b7d11 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Tue, 13 Oct 2026 12:22:43 +0700 Subject: [PATCH] fix(cancel): wait for the attention ack across message boundaries flush_done_attention stopped at the first end-of-message. SQL Server 2022 ends a batch cancelled before it returned anything with a DONE_ERROR message and sends the DONE_ATTN acknowledgement as the next message, so cancel_query failed with "Never got a DONE token acknowledging the Attention signal." and the next request read the acknowledgement. --- src/tds/stream/token.rs | 77 +++++++++++++++++++++++++++++++++++++++-- tests/query.rs | 33 ++++++++++++++++++ 2 files changed, 107 insertions(+), 3 deletions(-) diff --git a/src/tds/stream/token.rs b/src/tds/stream/token.rs index 8896eaf4..3d8008c0 100644 --- a/src/tds/stream/token.rs +++ b/src/tds/stream/token.rs @@ -219,7 +219,8 @@ where /// received, per MS-TDS section 2.2.1.6. Returns that DONE token so the /// connection is left clean and ready for reuse. pub(crate) async fn flush_done_attention(self) -> crate::Result { - let mut stream = self.try_unfold(); + // The acknowledgement can arrive in a later message than the cancelled response. + let mut stream = self.unfold(false); loop { match stream.try_next().await? { @@ -456,6 +457,10 @@ where } pub fn try_unfold(self) -> BoxStream<'a, crate::Result> { + self.unfold(true) + } + + fn unfold(self, stop_at_end_of_message: bool) -> BoxStream<'a, crate::Result> { // Read the per-response command timeout before `self` is moved into the // unfold closure. Every command result-read path funnels through here // (query/execute/simple_query result streams, the bulk-insert server @@ -467,8 +472,8 @@ where // otherwise unreachable from `RoundTripTimeout`). let poison = self.conn.command_desync_flag(); - let stream = futures_util::stream::try_unfold(self, |mut this| async move { - if this.conn.is_eof() { + let stream = futures_util::stream::try_unfold(self, move |mut this| async move { + if stop_at_end_of_message && this.conn.is_eof() { match this.last_error { None => return Ok(None), Some(error) => return Err(error), @@ -961,6 +966,72 @@ mod tests { "with no command timeout the read should stay pending on the stalled server" ); } + + // --- (e) attention acknowledgement -------------------------------------- + + const DONE_ERROR: u16 = 0x02; + const DONE_COUNT: u16 = 0x10; + const DONE_ATTN: u16 = 0x20; + + // What SQL Server 2022 sends for a batch cancelled before it returned + // anything: a DONE_ERROR message, then the DONE_ATTN message. + fn cancelled_response_then_ack() -> Vec { + let mut script = packet_bytes(&done_token(DONE_ERROR, 0, 0)); + script.extend(packet_bytes(&done_token(DONE_ATTN, 0, 0))); + script + } + + #[tokio::test] + async fn cancel_reads_the_ack_from_the_next_message() { + let mut conn = Connection::test_over(MockStream::new(cancelled_response_then_ack()), false); + + let done = conn + .cancel_request() + .await + .expect("the acknowledgement in the second message must be read"); + + assert!(done.is_attention()); + assert!(conn.is_eof()); + } + + #[tokio::test] + async fn request_after_cancel_reads_its_own_response() { + use crate::tds::codec::BatchRequest; + + let mut script = cancelled_response_then_ack(); + script.extend(packet_bytes(&done_token(DONE_COUNT, 0, 1))); + let mut conn = Connection::test_over(MockStream::new(script), false); + + conn.cancel_request().await.expect("cancel must succeed"); + + conn.flush_stream().await.unwrap(); + let id = conn.context_mut().next_packet_id(); + let req = BatchRequest::new("SELECT 1", [0; 8]); + conn.send(PacketHeader::batch(id), req).await.unwrap(); + + let done = TokenStream::new(&mut conn) + .flush_done() + .await + .expect("the next request must read its own DONE"); + + assert!(!done.is_attention()); + assert_eq!(done.rows(), 1); + } + + #[tokio::test] + async fn cancel_fails_when_the_transport_closes_before_the_ack() { + let mut conn = conn_over(&done_token(DONE_ERROR, 0, 0)); + + let err = conn + .cancel_request() + .await + .expect_err("a closed transport must not count as an acknowledgement"); + + assert!( + matches!(err, Error::Io { kind, .. } if kind == io::ErrorKind::UnexpectedEof), + "expected UnexpectedEof, got: {err:?}" + ); + } } // Server-free tests for the `RoundTripTimeout` stream combinator that backs diff --git a/tests/query.rs b/tests/query.rs index c573e435..03ffdc3a 100644 --- a/tests/query.rs +++ b/tests/query.rs @@ -1,3 +1,4 @@ +use futures_util::future::{self, Either}; use futures_util::io::{AsyncRead, AsyncWrite}; use futures_util::stream::TryStreamExt; use names::{Generator, Name}; @@ -5,6 +6,7 @@ use once_cell::sync::Lazy; use std::cell::RefCell; use std::env; use std::sync::Once; +use std::time::{Duration, Instant}; use tiberius::FromSql; use tiberius::{numeric::Numeric, xml::XmlData, ColumnType, Query, QueryItem, Result}; @@ -3182,3 +3184,34 @@ where Ok(()) } + +#[test_on_runtimes] +async fn cancel_query_after_dropping_a_running_query(mut conn: tiberius::Client) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin + Send, +{ + let running = Box::pin(conn.simple_query("WAITFOR DELAY '00:00:30'")); + let delay = futures_timer::Delay::new(Duration::from_millis(300)); + + if let Either::Left(_) = future::select(running, delay).await { + panic!("WAITFOR DELAY '00:00:30' must still be running after 300 ms"); + } + + let started = Instant::now(); + conn.cancel_query().await?; + assert!( + started.elapsed() < Duration::from_secs(5), + "cancel_query must not wait for the WAITFOR to finish, took {:?}", + started.elapsed() + ); + + let row = conn + .simple_query("SELECT 1") + .await? + .into_row() + .await? + .unwrap(); + assert_eq!(Some(1i32), row.get(0)); + + Ok(()) +}