diff --git a/Cargo.toml b/Cargo.toml index 080a9866..4d748640 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -150,8 +150,6 @@ time = { version = "0.3", optional = true } chrono = { version = "0.4", optional = true, features = ["serde"] } bstr = { version = "1.11.0", default-features = false } quanta = { version = "0.12", optional = true } -polonius-the-crab = "0.5.0" - bnum = "0.13.0" [dev-dependencies] @@ -161,8 +159,7 @@ serde = { version = "1.0.106", features = ["derive"] } tokio = { version = "1.0.1", features = ["full", "test-util", "io-util"] } hyper = { version = "1.1", features = ["server"] } indexmap = { version = "2.10.0", features = ["serde"] } -linked-hash-map = { version = "0.5.6", features = ["serde_impl"] } -fxhash = { version = "0.2.1" } +rustc-hash = "2" serde_bytes = "0.11.4" serde_json = "1" serde_repr = "0.1.7" diff --git a/rustfmt.toml b/rustfmt.toml index ef4162c2..33a75456 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,2 +1,2 @@ -edition = "2021" +edition = "2024" merge_derives = false diff --git a/src/cursors/row.rs b/src/cursors/row.rs index ea622b5c..4d9fb48a 100644 --- a/src/cursors/row.rs +++ b/src/cursors/row.rs @@ -12,7 +12,6 @@ use crate::{ use bytes::Buf; use clickhouse_types::error::TypesError; use clickhouse_types::parse_rbwnat_columns_header; -use polonius_the_crab::prelude::*; use std::marker::PhantomData; use std::pin::Pin; use std::task::{Context, Poll, ready}; @@ -100,6 +99,35 @@ impl RowCursor { Next::new(self).await } + // Why the unsafe reborrow? + // + // NLL (the current borrow checker) cannot see that `bytes` is dead in the + // NotEnoughData branch of this loop. The returned value borrows from + // `bytes`, so NLL extends that borrow to the function's return lifetime, + // which blocks the `bytes.extend()` call that only runs when no value + // exists. This is a known Polonius limitation: + // https://github.com/rust-lang/rust/issues/51132 + // + // This used to use the `polonius-the-crab` crate, which wraps the same + // raw-pointer reborrow behind a macro. It was dropped because the crate + // and its dependency tree are unmaintained: + // - `paste` transitive dep: RUSTSEC-2024-0436 (unmaintained) + // - `polonius-the-crab`: no meaningful commits in 12+ months + // - `higher-kinded-types`, `macro_rules_attribute`: same status + // Four stagnant crates and two RustSec advisories for a macro that + // expands to one line of unsafe is a poor trade-off. + // + // Alternatives considered: + // - TryRow enum: borrow still escapes via the return type, same error. + // - async-only next() + poll_next_owned for Stream: same NLL issue. + // - interior mutability in BytesExt via UnsafeCell: roughly 3x the + // diff for the same amount of unsafe, just hidden. + // - double deserialisation / probe-then-extract: roughly 2x deser cost + // on the happy path, unacceptable for a perf-sensitive cursor. + // None compiled without unsafe somewhere, or had unacceptable costs. + // + // Once Polonius lands in stable rustc, this can be removed. + #[inline] fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll>>> where @@ -110,27 +138,34 @@ impl RowCursor { debug_assert!(self.row_metadata.is_some()); } - let mut bytes = &mut self.bytes; + let bytes = &mut self.bytes; loop { - polonius!(|bytes| -> Poll>>> { - if bytes.remaining() > 0 { - let mut slice = bytes.slice(); - let result = rowbinary::deserialize_row::>( - &mut slice, - self.row_metadata.as_ref(), - ); - - match result { - Ok(value) => { - bytes.set_remaining(slice.len()); - polonius_return!(Poll::Ready(Ok(Some(value)))) - } - Err(Error::NotEnoughData) => {} - Err(err) => polonius_return!(Poll::Ready(Err(err))), + // SAFETY: we create a second &mut to `bytes` via raw pointer so the + // borrow checker releases the original. This is sound because: + // - On Ok: we return immediately, only one &mut is live. + // - On NotEnoughData: the deserialised value does not exist, the + // reborrow is dead, and we fall through to extend(). + // - On Err: we return immediately. + // Polonius would prove this automatically; NLL cannot yet. + let reborrowed = unsafe { &mut *(bytes as *mut BytesExt) }; + + if reborrowed.remaining() > 0 { + let mut slice = reborrowed.slice(); + let result = rowbinary::deserialize_row::>( + &mut slice, + self.row_metadata.as_ref(), + ); + + match result { + Ok(value) => { + reborrowed.set_remaining(slice.len()); + return Poll::Ready(Ok(Some(value))); } + Err(Error::NotEnoughData) => {} + Err(err) => return Poll::Ready(Err(err)), } - }); + } match ready!(self.raw.poll_next(cx))? { Some(chunk) => bytes.extend(chunk), @@ -196,18 +231,22 @@ where #[inline] fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - // Temporarily take the cursor out in order for `cursor.poll_next` to return a value with - // the correct lifetime `'a` rather than the unnamed lifetime of `&mut self`. - let mut cursor = self.cursor.take().expect("Future polled after completion"); - - polonius!(|cursor| -> Poll>>> { - match cursor.poll_next(cx) { - Poll::Ready(value) => polonius_return!(Poll::Ready(value)), - Poll::Pending => {} + // Take cursor out so poll_next's return value gets lifetime 'a + // (not the anonymous reborrow lifetime of &mut self). + let cursor = self.cursor.take().expect("Future polled after completion"); + + // SAFETY: same pattern as poll_next above. We create a second &mut + // via raw pointer. On Ready the reborrow escapes via the return value + // and cursor is consumed. On Pending the reborrow is dead and we put + // cursor back. Sound for the same reasons; Polonius would accept this. + let reborrowed = unsafe { &mut *(cursor as *mut RowCursor) }; + + match reborrowed.poll_next(cx) { + Poll::Ready(value) => Poll::Ready(value), + Poll::Pending => { + self.cursor = Some(cursor); + Poll::Pending } - }); - - self.cursor = Some(cursor); - Poll::Pending + } } } diff --git a/tests/it/cursor_reborrow.rs b/tests/it/cursor_reborrow.rs new file mode 100644 index 00000000..40b73024 --- /dev/null +++ b/tests/it/cursor_reborrow.rs @@ -0,0 +1,252 @@ +// Tests for the unsafe reborrow in RowCursor::poll_next and Next::poll. +// +// These specifically exercise the code paths that previously used +// polonius-the-crab and now use a manual unsafe reborrow. The key +// scenarios are the get-or-retry loop (NotEnoughData -> extend -> retry) +// and borrowed deserialization (T::Value<'_> borrowing from the buffer). + +#![cfg(feature = "test-util")] + +use clickhouse::{Client, Row, test}; +use serde::{Deserialize, Serialize}; + +// -- Mock-based tests (no ClickHouse needed) -------------------------------- + +#[tokio::test] +async fn cursor_single_row() { + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + x: u32, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + mock.add(test::handlers::provide([R { x: 42 }])); + + let mut cursor = client.query("SELECT x").fetch::().unwrap(); + assert_eq!(cursor.next().await.unwrap(), Some(R { x: 42 })); + assert_eq!(cursor.next().await.unwrap(), None); +} + +#[tokio::test] +async fn cursor_multiple_rows() { + // The loop in poll_next is the bit that needs the reborrow. Multiple + // rows means the loop iterates, which exercises extend() after a + // successful deserialisation on the previous iteration. + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + id: u64, + data: String, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + let rows: Vec = (0..100) + .map(|i| R { + id: i, + data: format!("row-{i}"), + }) + .collect(); + mock.add(test::handlers::provide(rows.clone())); + + let mut cursor = client.query("SELECT id, data").fetch::().unwrap(); + let mut got = Vec::new(); + while let Some(row) = cursor.next().await.unwrap() { + got.push(row); + } + assert_eq!(got, rows); +} + +#[tokio::test] +async fn cursor_empty_result() { + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + x: u32, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + mock.add(test::handlers::provide(Vec::::new())); + + let mut cursor = client.query("SELECT x").fetch::().unwrap(); + assert_eq!(cursor.next().await.unwrap(), None); +} + +#[tokio::test] +async fn cursor_fetch_all_and_fetch_one() { + // fetch_all and fetch_one both go through poll_next internally. + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + v: String, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + + let rows = vec![ + R { + v: "aaa".to_string(), + }, + R { + v: "bbb".to_string(), + }, + R { + v: "ccc".to_string(), + }, + ]; + mock.add(test::handlers::provide(rows.clone())); + let got = client + .query("SELECT v") + .fetch_all::() + .await + .unwrap(); + assert_eq!(got, rows); + + mock.add(test::handlers::provide([R { + v: "one".to_string(), + }])); + let got = client + .query("SELECT v") + .fetch_one::() + .await + .unwrap(); + assert_eq!(got, R { + v: "one".to_string(), + }); +} + +// -- Integration tests (need a real ClickHouse) ----------------------------- + +#[tokio::test] +async fn cursor_large_result_spanning_chunks() { + // Large enough to span multiple HTTP response chunks, exercising the + // NotEnoughData -> raw.poll_next -> extend -> retry path in the loop. + // This is the core path the unsafe reborrow protects. + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + id: u64, + payload: String, + } + + let client = prepare_database!(); + client + .query( + "CREATE TABLE test (id UInt64, payload String) \ + ENGINE = MergeTree ORDER BY id", + ) + .execute() + .await + .unwrap(); + + // 500 rows with ~200 bytes each = ~100KB, enough to span chunks. + let expected: Vec = (0..500) + .map(|i| R { + id: i, + payload: format!("{i:0>200}"), + }) + .collect(); + + let mut insert = client.insert::("test").await.unwrap(); + for row in &expected { + insert.write(row).await.unwrap(); + } + insert.end().await.unwrap(); + + let mut cursor = client + .query("SELECT id, payload FROM test ORDER BY id") + .fetch::() + .unwrap(); + + let mut got = Vec::new(); + while let Some(row) = cursor.next().await.unwrap() { + got.push(row); + } + assert_eq!(got.len(), expected.len()); + assert_eq!(got, expected); +} + +#[tokio::test] +async fn cursor_borrowed_rows() { + // Borrowed deserialization is the reason the unsafe exists — the + // returned T::Value<'_> borrows from the cursor's internal buffer. + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct Borrowed<'a> { + id: u64, + data: &'a str, + } + + let client = prepare_database!(); + crate::create_simple_table(&client, "test").await; + + let mut insert = client.insert::>("test").await.unwrap(); + insert + .write(&Borrowed { id: 1, data: "one" }) + .await + .unwrap(); + insert + .write(&Borrowed { + id: 2, + data: "two", + }) + .await + .unwrap(); + insert + .write(&Borrowed { + id: 3, + data: "three", + }) + .await + .unwrap(); + insert.end().await.unwrap(); + + let mut cursor = client + .query("SELECT id, data FROM test ORDER BY id") + .fetch::>() + .unwrap(); + + let row = cursor.next().await.unwrap().unwrap(); + assert_eq!(row, Borrowed { id: 1, data: "one" }); + let row = cursor.next().await.unwrap().unwrap(); + assert_eq!(row, Borrowed { + id: 2, + data: "two", + }); + let row = cursor.next().await.unwrap().unwrap(); + assert_eq!( + row, + Borrowed { + id: 3, + data: "three" + } + ); + assert_eq!(cursor.next().await.unwrap(), None); +} + +#[tokio::test] +async fn cursor_small_block_size() { + // Force ClickHouse to send one row per chunk. This maximises the + // number of extend() calls per row, hammering the reborrow path. + let client = prepare_database!(); + crate::create_simple_table(&client, "test").await; + + let mut insert = client.insert::("test").await.unwrap(); + for i in 0..50 { + insert + .write(&crate::SimpleRow::new(i, format!("val-{i}"))) + .await + .unwrap(); + } + insert.end().await.unwrap(); + + let mut cursor = client + .with_option("max_block_size", "1") + .query("SELECT ?fields FROM test ORDER BY id") + .fetch::() + .unwrap(); + + let mut count = 0u64; + while cursor.next().await.unwrap().is_some() { + count += 1; + } + assert_eq!(count, 50); +} diff --git a/tests/it/main.rs b/tests/it/main.rs index a137381d..f4843be6 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -250,6 +250,7 @@ mod chrono; mod cloud_jwt; mod compression; mod cursor_error; +mod cursor_reborrow; mod cursor_stats; mod fetch_bytes; mod https_errors; diff --git a/tests/it/rbwnat_smoke.rs b/tests/it/rbwnat_smoke.rs index 77ae3005..2bffd153 100644 --- a/tests/it/rbwnat_smoke.rs +++ b/tests/it/rbwnat_smoke.rs @@ -3,9 +3,8 @@ use crate::geo_types::{LineString, MultiLineString, MultiPolygon, Point, Polygon use crate::{SimpleRow, create_simple_table, execute_statements, get_client, insert_and_select}; use clickhouse::Row; use clickhouse::sql::Identifier; -use fxhash::FxHashMap; use indexmap::IndexMap; -use linked_hash_map::LinkedHashMap; +use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use serde_repr::{Deserialize_repr, Serialize_repr}; use std::collections::HashMap; @@ -392,9 +391,9 @@ async fn maps_third_party() { #[derive(Clone, Debug, Row, Serialize, Deserialize, PartialEq)] struct Data { im: IndexMap, - lhm: LinkedHashMap, + lhm: IndexMap, fx: FxHashMap, - weird_but_ok: LinkedHashMap>>>, + weird_but_ok: IndexMap>>>, } let client = prepare_database!(); @@ -417,9 +416,9 @@ async fn maps_third_party() { let rows = vec![Data { im: IndexMap::from_iter(vec![(1, "one".to_string()), (2, "two".to_string())]), - lhm: LinkedHashMap::from_iter(vec![(3, "three".to_string()), (4, "four".to_string())]), + lhm: IndexMap::from_iter(vec![(3, "three".to_string()), (4, "four".to_string())]), fx: FxHashMap::from_iter(vec![(5, "five".to_string()), (6, "six".to_string())]), - weird_but_ok: LinkedHashMap::from_iter(vec![( + weird_but_ok: IndexMap::from_iter(vec![( 7u128, IndexMap::from_iter(vec![( -8i8,