|
| 1 | +use std::time::Duration; |
| 2 | + |
| 3 | +use bytes::{BufMut, BytesMut}; |
| 4 | +use rust::setup::{admin_tokio, connection_sqlx_direct}; |
| 5 | +use sqlx::PgPool; |
| 6 | +use tokio::{io::AsyncWriteExt, net::TcpStream, task::JoinHandle, time::timeout}; |
| 7 | +use tokio_postgres::{CancelToken, Error as PgError, NoTls, SimpleQueryMessage}; |
| 8 | + |
| 9 | +/// Returns whether `pid` has an active `pg_sleep` query visible in `pg_stat_activity`. |
| 10 | +/// Uses a direct PostgreSQL connection so the result bypasses pgdog completely. |
| 11 | +async fn is_sleeping(direct: &PgPool, pid: i32) -> bool { |
| 12 | + let count: i64 = sqlx::query_scalar( |
| 13 | + "SELECT COUNT(*) \ |
| 14 | + FROM pg_stat_activity \ |
| 15 | + WHERE pid = $1 \ |
| 16 | + AND state = 'active' \ |
| 17 | + AND query LIKE '%pg_sleep%'", |
| 18 | + ) |
| 19 | + .bind(pid) |
| 20 | + .fetch_one(direct) |
| 21 | + .await |
| 22 | + .unwrap(); |
| 23 | + count == 1 |
| 24 | +} |
| 25 | + |
| 26 | +/// Connect to pgdog, pin to a specific PG backend via BEGIN, capture the backend pid |
| 27 | +/// via `pg_backend_pid()`, and launch `SELECT pg_sleep(60)` in a background task. |
| 28 | +/// |
| 29 | +/// `application_name` is embedded in the connection string so the caller can identify |
| 30 | +/// this connection in `SHOW CLIENTS` if needed. |
| 31 | +/// |
| 32 | +/// Returns `(backend_pid, cancel_token, query_handle)`. The caller owns `cancel_token` |
| 33 | +/// and `query_handle`; both must be driven to completion to keep the test clean. |
| 34 | +async fn start_sleeping_connection( |
| 35 | + application_name: &str, |
| 36 | +) -> ( |
| 37 | + i32, |
| 38 | + CancelToken, |
| 39 | + JoinHandle<Result<Vec<SimpleQueryMessage>, PgError>>, |
| 40 | +) { |
| 41 | + let (client, connection) = tokio_postgres::connect( |
| 42 | + &format!( |
| 43 | + "host=127.0.0.1 user=pgdog dbname=pgdog password=pgdog port=6432 application_name={application_name}" |
| 44 | + ), |
| 45 | + NoTls, |
| 46 | + ) |
| 47 | + .await |
| 48 | + .unwrap(); |
| 49 | + |
| 50 | + tokio::spawn(async move { |
| 51 | + if let Err(e) = connection.await { |
| 52 | + eprintln!("pgdog connection error: {}", e); |
| 53 | + } |
| 54 | + }); |
| 55 | + |
| 56 | + let cancel_token = client.cancel_token(); |
| 57 | + |
| 58 | + // BEGIN pins the client to one backend for the duration of the transaction. |
| 59 | + // Without this, transaction-mode pooling may assign a different backend to |
| 60 | + // pg_sleep than the one whose pid we captured. |
| 61 | + client.simple_query("BEGIN").await.unwrap(); |
| 62 | + |
| 63 | + let row = client |
| 64 | + .query_one("SELECT pg_backend_pid()", &[]) |
| 65 | + .await |
| 66 | + .unwrap(); |
| 67 | + let backend_pid: i32 = row.get(0); |
| 68 | + |
| 69 | + let handle = tokio::spawn(async move { client.simple_query("SELECT pg_sleep(60)").await }); |
| 70 | + |
| 71 | + (backend_pid, cancel_token, handle) |
| 72 | +} |
| 73 | + |
| 74 | +/// Assert that a query handle returned by `start_sleeping_connection` was cancelled: |
| 75 | +/// it must resolve to SQLSTATE 57014 (canceling statement due to user request). |
| 76 | +async fn assert_cancelled( |
| 77 | + handle: JoinHandle<Result<Vec<SimpleQueryMessage>, PgError>>, |
| 78 | + label: &str, |
| 79 | +) { |
| 80 | + let result = timeout(Duration::from_secs(5), handle) |
| 81 | + .await |
| 82 | + .expect(&format!( |
| 83 | + "{label}: cancelled query did not unblock within 5 seconds" |
| 84 | + )) |
| 85 | + .expect(&format!("{label}: task panicked")); |
| 86 | + |
| 87 | + let err = result.expect_err(&format!( |
| 88 | + "{label}: query should have been cancelled, but it succeeded" |
| 89 | + )); |
| 90 | + let db_err = err.as_db_error().expect(&format!( |
| 91 | + "{label}: expected a PostgreSQL error, not a network error" |
| 92 | + )); |
| 93 | + |
| 94 | + assert_eq!( |
| 95 | + db_err.code().code(), |
| 96 | + "57014", |
| 97 | + "{label}: expected SQLSTATE 57014, got {}", |
| 98 | + db_err.code().code() |
| 99 | + ); |
| 100 | +} |
| 101 | + |
| 102 | +/// Verify that cancellation is precise: two independent connections both run a long |
| 103 | +/// query and each cancel request stops exactly one of them. |
| 104 | +/// |
| 105 | +/// Steps: |
| 106 | +/// 1. Two clients connect through pgdog; each starts `SELECT pg_sleep(60)`. |
| 107 | +/// 2. Both queries are confirmed active on specific PG backends via `pg_stat_activity`. |
| 108 | +/// 3. Cancel connection 1 → only backend 1 stops; backend 2 remains active. |
| 109 | +/// 4. Cancel connection 2 → backend 2 stops. |
| 110 | +#[tokio::test] |
| 111 | +async fn test_cancel_query() { |
| 112 | + let direct = connection_sqlx_direct().await; |
| 113 | + |
| 114 | + let (pid1, token1, handle1) = start_sleeping_connection("cancel_test").await; |
| 115 | + let (pid2, token2, handle2) = start_sleeping_connection("cancel_test").await; |
| 116 | + |
| 117 | + // Give both queries time to reach their respective backends. |
| 118 | + tokio::time::sleep(Duration::from_millis(300)).await; |
| 119 | + |
| 120 | + assert!( |
| 121 | + is_sleeping(&direct, pid1).await, |
| 122 | + "connection 1 (backend {pid1}) should be active before any cancel" |
| 123 | + ); |
| 124 | + assert!( |
| 125 | + is_sleeping(&direct, pid2).await, |
| 126 | + "connection 2 (backend {pid2}) should be active before any cancel" |
| 127 | + ); |
| 128 | + |
| 129 | + // ── Cancel connection 1 ──────────────────────────────────────────────── |
| 130 | + token1.cancel_query(NoTls).await.unwrap(); |
| 131 | + |
| 132 | + // Wait for the client to receive the cancellation error. |
| 133 | + // By the time the handle resolves, the backend has already stopped. |
| 134 | + assert_cancelled(handle1, "connection 1").await; |
| 135 | + |
| 136 | + // Connection 1's backend is gone; connection 2 must still be running. |
| 137 | + tokio::time::sleep(Duration::from_millis(100)).await; |
| 138 | + assert!( |
| 139 | + !is_sleeping(&direct, pid1).await, |
| 140 | + "backend {pid1} should be idle after cancelling connection 1" |
| 141 | + ); |
| 142 | + assert!( |
| 143 | + is_sleeping(&direct, pid2).await, |
| 144 | + "backend {pid2} should still be active after cancelling connection 1 only" |
| 145 | + ); |
| 146 | + |
| 147 | + // ── Cancel connection 2 ──────────────────────────────────────────────── |
| 148 | + token2.cancel_query(NoTls).await.unwrap(); |
| 149 | + |
| 150 | + assert_cancelled(handle2, "connection 2").await; |
| 151 | + |
| 152 | + tokio::time::sleep(Duration::from_millis(100)).await; |
| 153 | + assert!( |
| 154 | + !is_sleeping(&direct, pid2).await, |
| 155 | + "backend {pid2} should be idle after cancelling connection 2" |
| 156 | + ); |
| 157 | +} |
| 158 | + |
| 159 | +/// Verify that a cancel request carrying a wrong pid and secret is silently rejected: |
| 160 | +/// the running query is unaffected and the client does not receive a cancellation error. |
| 161 | +/// |
| 162 | +/// pgdog's `verify_cancel` gate must reject the request before it reaches the pool, |
| 163 | +/// so the backend continues executing as if nothing happened. |
| 164 | +#[tokio::test] |
| 165 | +async fn test_cancel_query_wrong_secret() { |
| 166 | + let direct = connection_sqlx_direct().await; |
| 167 | + let app_name = "cancel_test_wrong_secret"; |
| 168 | + let (backend_pid, real_cancel_token, query_handle) = start_sleeping_connection(app_name).await; |
| 169 | + |
| 170 | + // Give the query time to reach the backend. |
| 171 | + tokio::time::sleep(Duration::from_millis(300)).await; |
| 172 | + |
| 173 | + assert!( |
| 174 | + is_sleeping(&direct, backend_pid).await, |
| 175 | + "query should be running before wrong-secret cancel" |
| 176 | + ); |
| 177 | + |
| 178 | + // Look up the pgdog client pid from the admin interface. |
| 179 | + // SHOW CLIENTS exposes the pid (the 'id' column) that pgdog assigned during login — |
| 180 | + // the same value that was sent in the K message and that verify_cancel checks against. |
| 181 | + let admin = admin_tokio().await; |
| 182 | + let messages = admin.simple_query("SHOW CLIENTS").await.unwrap(); |
| 183 | + let pgdog_pid: i32 = messages |
| 184 | + .iter() |
| 185 | + .filter_map(|m| match m { |
| 186 | + SimpleQueryMessage::Row(row) => Some(row), |
| 187 | + _ => None, |
| 188 | + }) |
| 189 | + .find(|row| row.get("application_name") == Some(app_name)) |
| 190 | + .expect("connection should appear in SHOW CLIENTS") |
| 191 | + .get("id") |
| 192 | + .expect("id column should be present") |
| 193 | + .parse() |
| 194 | + .expect("id should be a valid i32"); |
| 195 | + |
| 196 | + // Send a CancelRequest with the real pgdog client pid but a wrong secret. |
| 197 | + // pgdog will find the client in comms by pid, then reject it because |
| 198 | + // the secret doesn't match — verify_cancel returns false. |
| 199 | + let mut raw = TcpStream::connect("127.0.0.1:6432").await.unwrap(); |
| 200 | + let mut buf = BytesMut::new(); |
| 201 | + buf.put_i32(16); // total message length (including the length field) |
| 202 | + buf.put_i32(80877102); // CancelRequest magic code |
| 203 | + buf.put_i32(pgdog_pid); // correct pid |
| 204 | + buf.put_i32(0); // wrong secret |
| 205 | + raw.write_all(&buf).await.unwrap(); |
| 206 | + // pgdog closes the connection silently after processing; no response is sent. |
| 207 | + drop(raw); |
| 208 | + |
| 209 | + // Give pgdog enough time to receive and process the bogus cancel. |
| 210 | + tokio::time::sleep(Duration::from_millis(300)).await; |
| 211 | + |
| 212 | + // The query must still be running — the secret mismatch was caught by verify_cancel. |
| 213 | + assert!( |
| 214 | + is_sleeping(&direct, backend_pid).await, |
| 215 | + "query should still be running after wrong-secret cancel — verify_cancel must have rejected it" |
| 216 | + ); |
| 217 | + |
| 218 | + // Clean up: cancel for real. |
| 219 | + real_cancel_token.cancel_query(NoTls).await.unwrap(); |
| 220 | + assert_cancelled(query_handle, "wrong-secret test cleanup").await; |
| 221 | +} |
0 commit comments