diff --git a/sqlx-mysql/src/connection/executor.rs b/sqlx-mysql/src/connection/executor.rs index ee59d03d0a..88c5b1e034 100644 --- a/sqlx-mysql/src/connection/executor.rs +++ b/sqlx-mysql/src/connection/executor.rs @@ -1,5 +1,5 @@ use super::MySqlStream; -use crate::connection::stream::Waiting; +use crate::connection::stream::{PrepareProgress, Waiting}; use crate::error::Error; use crate::executor::{Execute, Executor}; use crate::ext::ustr::UStr; @@ -33,6 +33,13 @@ impl MySqlConnection { // https://dev.mysql.com/doc/internals/en/com-stmt-prepare.html // https://dev.mysql.com/doc/internals/en/com-stmt-prepare-response.html#packet-COM_STMT_PREPARE_OK + // Queued before the request, so the drain can finish and close a + // response a cancelled prepare left behind. + self.inner + .stream + .waiting + .push_back(Waiting::Prepare(PrepareProgress::Header)); + self.inner .stream .send_packet(Prepare { query: sql }) @@ -70,6 +77,10 @@ impl MySqlConnection { column_names: Arc::new(column_names), }; + // Queued so the drain closes the statement if the caller is + // dropped before caching or closing it. + self.inner.stream.close_pending.push(id); + Ok((id, metadata)) } @@ -85,11 +96,15 @@ impl MySqlConnection { let (id, metadata) = self.prepare_statement(sql).await?; // in case of the cache being full, close the least recently used statement - if let Some((id, _)) = self + let evicted = self .inner .cache_statement - .insert(sql, (id, metadata.clone())) - { + .insert(sql, (id, metadata.clone())); + + // Cached, so no longer closed by the drain. + self.inner.stream.close_pending.pop(); + + if let Some((id, _)) = evicted { self.inner .stream .send_packet(StmtClose { statement: id }) @@ -110,7 +125,6 @@ impl MySqlConnection { let mut logger = QueryLogger::new(sql, self.inner.log_settings.clone()); self.inner.stream.wait_until_ready().await?; - self.inner.stream.waiting.push_back(Waiting::Result); Ok(try_stream! { let sql = logger.sql().as_str(); @@ -136,6 +150,10 @@ impl MySqlConnection { ); } + // Queued only once the request is sent, so a cancelled + // prepare or a failed parameter-count check leaves none. + self.inner.stream.waiting.push_back(Waiting::Result); + // https://dev.mysql.com/doc/internals/en/com-stmt-execute.html self.inner.stream .send_packet(StatementExecute { @@ -151,6 +169,7 @@ impl MySqlConnection { .await?; if arguments.types.len() != metadata.parameters { + // Not cached and not closed here; the drain closes it. return Err( err_protocol!( "prepared statement expected {} parameters but {} parameters were provided", @@ -160,6 +179,8 @@ impl MySqlConnection { ); } + self.inner.stream.waiting.push_back(Waiting::Result); + // https://dev.mysql.com/doc/internals/en/com-stmt-execute.html self.inner.stream .send_packet(StatementExecute { @@ -168,11 +189,14 @@ impl MySqlConnection { }) .await?; + self.inner.stream.close_pending.pop(); self.inner.stream.send_packet(StmtClose { statement: id }).await?; MySqlValueFormat::Binary } } else { + self.inner.stream.waiting.push_back(Waiting::Result); + // https://dev.mysql.com/doc/internals/en/com-query.html self.inner.stream.send_packet(Query(sql)).await?; @@ -335,6 +359,7 @@ impl<'c> Executor<'c> for &'c mut MySqlConnection { } else { let (id, metadata) = self.prepare_statement(sql.as_str()).await?; + self.inner.stream.close_pending.pop(); self.inner .stream .send_packet(StmtClose { statement: id }) @@ -365,6 +390,7 @@ impl<'c> Executor<'c> for &'c mut MySqlConnection { let (id, metadata) = self.prepare_statement(sql.as_str()).await?; + self.inner.stream.close_pending.pop(); self.inner .stream .send_packet(StmtClose { statement: id }) diff --git a/sqlx-mysql/src/connection/stream.rs b/sqlx-mysql/src/connection/stream.rs index 25c9f3c7ed..f9e1fe4fb9 100644 --- a/sqlx-mysql/src/connection/stream.rs +++ b/sqlx-mysql/src/connection/stream.rs @@ -1,13 +1,14 @@ use std::collections::VecDeque; -use std::ops::{Deref, DerefMut}; +use std::ops::{ControlFlow, Deref, DerefMut}; -use bytes::{Buf, Bytes, BytesMut}; +use bytes::{Bytes, BytesMut}; use crate::error::Error; use crate::io::MySqlBufExt; use crate::io::{ProtocolDecode, ProtocolEncode}; use crate::net::{BufferedSocket, Socket}; use crate::protocol::response::{EofPacket, ErrPacket, OkPacket, Status}; +use crate::protocol::statement::{PrepareOk, StmtClose}; use crate::protocol::{Capabilities, Packet}; use crate::{MySqlConnectOptions, MySqlDatabaseError}; @@ -18,6 +19,8 @@ pub struct MySqlStream> { pub(super) capabilities: Capabilities, pub(crate) sequence_id: u8, pub(crate) waiting: VecDeque, + // statements a dropped operation left open; closed by the drain + pub(crate) close_pending: Vec, pub(crate) is_tls: bool, } @@ -28,6 +31,21 @@ pub(crate) enum Waiting { // waiting for a row within a result set Row, + + // waiting for (the rest of) a COM_STMT_PREPARE response + Prepare(PrepareProgress), +} + +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum PrepareProgress { + // the stmt-prepare-ok header has not been read yet + Header, + + // the header was read; this many definition and EOF packets follow + Tail { + statement_id: u32, + packets_left: u32, + }, } impl MySqlStream { @@ -54,6 +72,7 @@ impl MySqlStream { Self { waiting: VecDeque::new(), + close_pending: Vec::new(), capabilities, server_version: (0, 0, 0), sequence_id: 0, @@ -63,11 +82,30 @@ impl MySqlStream { } pub(crate) async fn wait_until_ready(&mut self) -> Result<(), Error> { + // Close the statements a dropped operation left open. + for statement in std::mem::take(&mut self.close_pending) { + self.sequence_id = 0; + self.write_packet(StmtClose { statement })?; + } + if !self.socket.write_buffer().is_empty() { self.socket.flush().await?; } while !self.waiting.is_empty() { + while matches!(self.waiting.front(), Some(Waiting::Prepare(_))) { + // The rest of a response a cancelled prepare left behind. + let (_, prepared) = self.recv_packet_tracked().await?; + + // Close the statement in the poll that completed it, so a + // cancellation cannot lose the id. + if let Some(statement) = prepared { + self.sequence_id = 0; + self.write_packet(StmtClose { statement })?; + self.socket.flush().await?; + } + } + while self.waiting.front() == Some(&Waiting::Row) { let packet = self.recv_packet().await?; @@ -123,16 +161,37 @@ impl MySqlStream { // https://dev.mysql.com/doc/dev/mysql-server/8.0.12/page_protocol_basic_packets.html // https://mariadb.com/kb/en/library/0-packet/#standard-packet - let mut header: Bytes = self.socket.read(4).await?; + // One `try_read` takes header and payload, so a cancelled read leaves + // the stream at a part boundary. Payloads of 0xFFFFFF bytes or more + // span several parts. + const HEADER_LEN: usize = 4; + + let (sequence_id, payload) = self + .socket + .try_read(|buf| { + if buf.len() < HEADER_LEN { + return Ok(ControlFlow::Continue(HEADER_LEN)); + } - // cannot overflow - #[allow(clippy::cast_possible_truncation)] - let packet_size = header.get_uint_le(3) as usize; - let sequence_id = header.get_u8(); + // cannot overflow + #[allow(clippy::cast_possible_truncation)] + let packet_size = u32::from_le_bytes([buf[0], buf[1], buf[2], 0]) as usize; - self.sequence_id = sequence_id.wrapping_add(1); + let frame_len = HEADER_LEN + packet_size; + + if buf.len() < frame_len { + return Ok(ControlFlow::Continue(frame_len)); + } + + let mut frame = buf.split_to(frame_len); + let sequence_id = frame[3]; + let payload = frame.split_off(HEADER_LEN).freeze(); - let payload: Bytes = self.socket.read(packet_size).await?; + Ok(ControlFlow::Break((sequence_id, payload))) + }) + .await?; + + self.sequence_id = sequence_id.wrapping_add(1); // TODO: packet compression @@ -142,6 +201,14 @@ impl MySqlStream { // receive the next packet from the database server // may block (async) on more data from the server pub(crate) async fn recv_packet(&mut self) -> Result, Error> { + let (packet, _) = self.recv_packet_tracked().await?; + + Ok(packet) + } + + /// Like `recv_packet`, plus the statement id of a COM_STMT_PREPARE + /// response this packet completed. + async fn recv_packet_tracked(&mut self) -> Result<(Packet, Option), Error> { let payload = self.recv_packet_part().await?; let payload = if payload.len() < 0xFF_FF_FF { payload @@ -174,7 +241,53 @@ impl MySqlStream { ); } - Ok(Packet(payload)) + let prepared = self.note_prepare_response_packet(&payload)?; + + Ok((Packet(payload), prepared)) + } + + /// Counts the packets of a pending COM_STMT_PREPARE response and returns + /// its statement id once complete. Only stmt-prepare-ok says how many + /// packets follow, so every read path counts here. + fn note_prepare_response_packet(&mut self, payload: &Bytes) -> Result, Error> { + let capabilities = self.capabilities; + + let Some(Waiting::Prepare(progress)) = self.waiting.front_mut() else { + return Ok(None); + }; + + let (statement_id, packets_left) = match *progress { + PrepareProgress::Header => { + let ok = PrepareOk::decode_with(payload.clone(), capabilities)?; + + let eof_packets = if capabilities.contains(Capabilities::DEPRECATE_EOF) { + 0 + } else { + u32::from(ok.params > 0) + u32::from(ok.columns > 0) + }; + + let packets_left = u32::from(ok.params) + u32::from(ok.columns) + eof_packets; + + (ok.statement_id, packets_left) + } + PrepareProgress::Tail { + statement_id, + packets_left, + } => (statement_id, packets_left - 1), + }; + + if packets_left == 0 { + self.waiting.pop_front(); + + return Ok(Some(statement_id)); + } + + *progress = PrepareProgress::Tail { + statement_id, + packets_left, + }; + + Ok(None) } pub(crate) async fn recv<'de, T>(&mut self) -> Result @@ -215,6 +328,7 @@ impl MySqlStream { capabilities: self.capabilities, sequence_id: self.sequence_id, waiting: self.waiting, + close_pending: self.close_pending, is_tls: self.is_tls, } } diff --git a/sqlx-mysql/src/connection/tls.rs b/sqlx-mysql/src/connection/tls.rs index dddb4c6aa5..359a5d852f 100644 --- a/sqlx-mysql/src/connection/tls.rs +++ b/sqlx-mysql/src/connection/tls.rs @@ -108,6 +108,7 @@ impl WithSocket for MapStream { capabilities: self.capabilities, sequence_id: self.sequence_id, waiting: self.waiting, + close_pending: Vec::new(), is_tls: true, } } diff --git a/tests/mysql/error.rs b/tests/mysql/error.rs index bd2cbb9ec0..66696fabcb 100644 --- a/tests/mysql/error.rs +++ b/tests/mysql/error.rs @@ -1,4 +1,9 @@ -use sqlx::{error::ErrorKind, mysql::MySql, Connection, Error}; +use anyhow::Context; +use sqlx::{ + error::ErrorKind, + mysql::{MySql, MySqlConnection}, + Connection, Error, Executor, Row, +}; use sqlx_test::new; fn mysql_supports_check_constraints(version: &str) -> bool { @@ -174,3 +179,59 @@ async fn it_fails_with_parameter_count_mismatch_zero_expected() -> anyhow::Resul Ok(()) } + +#[sqlx_macros::test] +async fn it_stays_usable_after_parameter_count_mismatch() -> anyhow::Result<()> { + let mut conn = new::().await?; + + let (prepared_before, closed_before) = stmt_prepare_close_counts(&mut conn).await?; + + for persistent in [true, false] { + let res = sqlx::query("SELECT * FROM tweet WHERE id = ?") + .persistent(persistent) + .execute(&mut conn) + .await; + + assert!(matches!(res, Err(Error::Protocol(_))), "{res:?}"); + + // The check fails before COM_STMT_EXECUTE, so no response is pending. + sqlx_core::rt::timeout(std::time::Duration::from_secs(5), conn.ping()) + .await + .with_context(|| format!("ping hangs, persistent({persistent})"))? + .with_context(|| format!("ping fails, persistent({persistent})"))?; + } + + // The persistent statement stays in the cache, and the other one is + // closed. + let (prepared_after, closed_after) = stmt_prepare_close_counts(&mut conn).await?; + let prepared = prepared_after - prepared_before; + let closed = closed_after - closed_before; + let cached = conn.cached_statements_size(); + assert_eq!( + (prepared, closed, cached), + (2, 1, 1), + "(prepared, closed, cached)" + ); + + Ok(()) +} + +// The session's Com_stmt_prepare and Com_stmt_close counts. +async fn stmt_prepare_close_counts(conn: &mut MySqlConnection) -> anyhow::Result<(u64, u64)> { + let rows = conn + .fetch_all( + "SHOW SESSION STATUS WHERE Variable_name IN ('Com_stmt_prepare', 'Com_stmt_close')", + ) + .await?; + + let counter = |name: &str| -> anyhow::Result { + let row = rows + .iter() + .find(|row| row.get::<&str, _>(0).eq_ignore_ascii_case(name)) + .with_context(|| format!("{name} is missing from SHOW SESSION STATUS"))?; + + Ok(row.try_get::<&str, _>(1)?.parse()?) + }; + + Ok((counter("Com_stmt_prepare")?, counter("Com_stmt_close")?)) +} diff --git a/tests/mysql/mysql.rs b/tests/mysql/mysql.rs index 78d3e3a761..51c85d182f 100644 --- a/tests/mysql/mysql.rs +++ b/tests/mysql/mysql.rs @@ -468,6 +468,233 @@ async fn test_issue_622() -> anyhow::Result<()> { Ok(()) } +// A query cancelled mid-read must leave the connection usable. +// Tokio only, since the test cancels through `tokio::select!`. +#[cfg(feature = "_rt-tokio")] +#[tokio::test] +async fn test_cancelled_mid_read_leaves_connection_usable() -> anyhow::Result<()> { + use std::time::Duration; + + setup_if_needed(); + + let mut conn = MySqlConnection::connect(&env::var("DATABASE_URL")?).await?; + + // Rows large enough to span several socket reads, so reads await + // mid-packet. Built by doubling to keep the SQL short. + conn.execute("DROP TABLE IF EXISTS sqlx_test_cancel_probe") + .await?; + conn.execute( + "CREATE TABLE sqlx_test_cancel_probe ( + id INT AUTO_INCREMENT PRIMARY KEY, + filler MEDIUMTEXT NOT NULL + )", + ) + .await?; + conn.execute("INSERT INTO sqlx_test_cancel_probe (filler) VALUES (REPEAT('x', 100000))") + .await?; + for _ in 0..5 { + conn.execute( + "INSERT INTO sqlx_test_cancel_probe (filler) \ + SELECT filler FROM sqlx_test_cancel_probe", + ) + .await?; + } + + let query = "SELECT id, filler FROM sqlx_test_cancel_probe"; + + for i in 0u64..300 { + // Tokio rounds sleeps up to whole milliseconds: a 0 to 3 ms race. + let cancel_after = Duration::from_micros((i % 23) * 90); + + tokio::select! { + rows = conn.fetch_all(query) => { + rows?; + } + _ = tokio::time::sleep(cancel_after) => {} + } + + // `ping` drains what the cancelled read left; a misaligned stream + // fails or hangs here. + tokio::time::timeout(Duration::from_secs(30), conn.ping()) + .await + .map_err(|_| anyhow::anyhow!("connection wedged after cancelled read (iteration {i})"))? + .with_context(|| format!("iteration {i}"))?; + } + + conn.execute("DROP TABLE sqlx_test_cancel_probe").await?; + conn.close().await?; + + Ok(()) +} + +// A prepare cancelled before its response is read must leave the +// connection usable, and every prepared statement closed or cached. +#[cfg(feature = "_rt-tokio")] +#[tokio::test] +async fn test_cancelled_mid_prepare_leaves_connection_usable() -> anyhow::Result<()> { + use futures_util::FutureExt; + use std::time::Duration; + + setup_if_needed(); + + let mut conn = MySqlConnection::connect(&env::var("DATABASE_URL")?).await?; + + conn.execute("DROP TABLE IF EXISTS sqlx_test_cancel_probe_prepared") + .await?; + conn.execute( + "CREATE TABLE sqlx_test_cancel_probe_prepared ( + id INT AUTO_INCREMENT PRIMARY KEY, + filler MEDIUMTEXT NOT NULL + )", + ) + .await?; + conn.execute( + "INSERT INTO sqlx_test_cancel_probe_prepared (filler) VALUES (REPEAT('x', 100000))", + ) + .await?; + for _ in 0..5 { + conn.execute( + "INSERT INTO sqlx_test_cancel_probe_prepared (filler) \ + SELECT filler FROM sqlx_test_cancel_probe_prepared", + ) + .await?; + } + + let query = "SELECT id, filler FROM sqlx_test_cancel_probe_prepared"; + + let (p, c, r) = stmt_prepare_close_counts(&mut conn).await?; + let mut counts = (p, c, 0, r); + + // Polled once, a fetch sends COM_STMT_PREPARE and is dropped, leaving + // the response to the drain. The statements cover columns only, + // parameters and columns, parameters only, and neither. + let shapes = [ + (query, false), + ( + "SELECT id FROM sqlx_test_cancel_probe_prepared WHERE id > ?", + true, + ), + ("DO ?", true), + ("DO 1", false), + ]; + + // Poll-once drops are deterministic: the first poll always ends at the + // prepare response read, so the balance below can only fail if the + // drain or the close tracking is wrong. Checked per persistent setting. + for persistent in [false, true] { + for (sql, with_param) in shapes { + for i in 0..100 { + let mut polled = sqlx::query(sql).persistent(persistent); + + if with_param { + polled = polled.bind(0); + } + + if let Some(rows) = polled.fetch_all(&mut conn).now_or_never() { + rows?; + } + + tokio::time::timeout(Duration::from_secs(30), conn.ping()) + .await + .map_err(|_| { + anyhow::anyhow!( + "connection wedged after dropped prepare \ + ({sql:?}, persistent({persistent}), iteration {i})" + ) + })? + .with_context(|| format!("{sql:?}, persistent({persistent}), iteration {i}"))?; + } + } + + counts = assert_statement_balance( + &mut conn, + counts, + &format!("dropped prepares, persistent({persistent})"), + ) + .await?; + } + + // Mostly cancels inside binary row reads; the loop above covers the + // prepare drain. No counter balance here: a cancel lands wherever the + // timer catches the fetch, so the balance depends on timing the code + // under test does not control. The pings cover the recovery. + for i in 0u64..300 { + // A 0 to 3 ms race, as above. + let cancel_after = Duration::from_micros((i % 23) * 90); + + tokio::select! { + rows = sqlx::query(query).persistent(false).fetch_all(&mut conn) => { + rows?; + } + _ = tokio::time::sleep(cancel_after) => {} + } + + tokio::time::timeout(Duration::from_secs(30), conn.ping()) + .await + .map_err(|_| { + anyhow::anyhow!("connection wedged after cancelled prepare/execute (iteration {i})") + })? + .with_context(|| format!("iteration {i}"))?; + } + + conn.execute("DROP TABLE sqlx_test_cancel_probe_prepared") + .await?; + conn.close().await?; + + Ok(()) +} + +// The Com_stmt_prepare/Com_stmt_close deltas since `before` must balance +// against the statement cache delta, with server-side re-prepares taken out +// of the prepare side: an execute that finds the statement's metadata stale +// re-prepares it on the server, which counts as Com_stmt_prepare without +// ever getting a COM_STMT_CLOSE. The phase name pins down a leak. +#[cfg(feature = "_rt-tokio")] +async fn assert_statement_balance( + conn: &mut MySqlConnection, + before: (u64, u64, u64, u64), + phase: &str, +) -> anyhow::Result<(u64, u64, u64, u64)> { + let (prepared, closed, reprepared) = stmt_prepare_close_counts(conn).await?; + let cached = conn.cached_statements_size() as u64; + + assert_eq!( + (closed - before.1) + (cached - before.2), + (prepared - before.0) - (reprepared - before.3), + "{phase}: {} prepared ({} reprepared), {} closed, cached {cached}", + prepared - before.0, + reprepared - before.3, + closed - before.1 + ); + + Ok((prepared, closed, cached, reprepared)) +} + +// The session's Com_stmt_prepare, Com_stmt_close and Com_stmt_reprepare counts. +#[cfg(feature = "_rt-tokio")] +async fn stmt_prepare_close_counts(conn: &mut MySqlConnection) -> anyhow::Result<(u64, u64, u64)> { + let rows = conn + .fetch_all( + "SHOW SESSION STATUS WHERE Variable_name IN ('Com_stmt_prepare', 'Com_stmt_close', 'Com_stmt_reprepare')", + ) + .await?; + + let counter = |name: &str| -> anyhow::Result { + let row = rows + .iter() + .find(|row| row.get::<&str, _>(0).eq_ignore_ascii_case(name)) + .with_context(|| format!("{name} is missing from SHOW SESSION STATUS"))?; + + Ok(row.try_get::<&str, _>(1)?.parse()?) + }; + + Ok(( + counter("Com_stmt_prepare")?, + counter("Com_stmt_close")?, + counter("Com_stmt_reprepare")?, + )) +} + #[sqlx_macros::test] async fn it_can_work_with_transactions() -> anyhow::Result<()> { let mut conn = new::().await?;