From 3df111f89182769b86e71d1c86040dd625a32e08 Mon Sep 17 00:00:00 2001 From: meskill <8974488+meskill@users.noreply.github.com> Date: Thu, 13 Aug 2026 10:07:36 +0000 Subject: [PATCH 1/2] fix(pool): extended pipeline mixed messages for multi-shard --- .../go/go_pgx/pipeline_sharded_test.go | 126 ++++ .../frontend/client/query_engine/test/mod.rs | 1 + .../query_engine/test/pipeline_execution.rs | 564 ++++++++++++++++++ pgdog/src/frontend/client/test/test_client.rs | 24 +- .../src/net/messages/parameter_description.rs | 5 + 5 files changed, 710 insertions(+), 10 deletions(-) create mode 100644 integration/go/go_pgx/pipeline_sharded_test.go create mode 100644 pgdog/src/frontend/client/query_engine/test/pipeline_execution.rs diff --git a/integration/go/go_pgx/pipeline_sharded_test.go b/integration/go/go_pgx/pipeline_sharded_test.go new file mode 100644 index 000000000..f8ad60d8f --- /dev/null +++ b/integration/go/go_pgx/pipeline_sharded_test.go @@ -0,0 +1,126 @@ +package main + +import ( + "context" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + textOID = uint32(25) + bigintOID = uint32(20) +) + +func fieldNames(description *pgconn.StatementDescription) []string { + names := []string{} + for _, field := range description.Fields { + names = append(names, field.Name) + } + return names +} + +// pgx folds ParseComplete, ParameterDescription and RowDescription into one +// StatementDescription, so dropping any of the three desynchronises the read. +func nextStatement(t *testing.T, pipeline *pgconn.Pipeline) *pgconn.StatementDescription { + t.Helper() + + result, err := pipeline.GetResults() + require.NoError(t, err) + + description, ok := result.(*pgconn.StatementDescription) + require.Truef(t, ok, "expected a statement description, got %T", result) + return description +} + +// A pipelined prepare phase describes several statements before one Sync, and +// pgx reads the replies positionally. The aggregate and the ORDER BY send the +// request to every shard, which is what puts it through the multi-shard path. +func TestPipelinedPrepareDescribesEveryStatement(t *testing.T) { + ctx := context.Background() + + conn, err := pgx.Connect(ctx, testConnStr) + require.NoError(t, err) + defer conn.Close(ctx) + + pipeline := conn.PgConn().StartPipeline(ctx) + + pipeline.SendPrepare("pipe1", "SELECT count(*) FROM sharded", nil) + pipeline.SendPrepare("pipe2", "SELECT id, $1::text AS second FROM sharded ORDER BY id", nil) + pipeline.SendPrepare( + "pipe3", + "SELECT id, $1::text AS second, $2::bigint AS third FROM sharded ORDER BY id", + nil, + ) + // Binary results, as every pgx query asks for. + pipeline.SendQueryPrepared("pipe1", nil, nil, []int16{1}) + require.NoError(t, pipeline.Sync()) + + first := nextStatement(t, pipeline) + assert.Empty(t, first.ParamOIDs) + assert.Equal(t, []string{"count"}, fieldNames(first)) + + second := nextStatement(t, pipeline) + assert.Equal(t, []uint32{textOID}, second.ParamOIDs) + assert.Equal(t, []string{"id", "second"}, fieldNames(second)) + + third := nextStatement(t, pipeline) + assert.Equal(t, []uint32{textOID, bigintOID}, third.ParamOIDs) + assert.Equal(t, []string{"id", "second", "third"}, fieldNames(third)) + + result, err := pipeline.GetResults() + require.NoError(t, err) + reader, ok := result.(*pgconn.ResultReader) + require.Truef(t, ok, "expected a result reader, got %T", result) + require.NoError(t, reader.Read().Err) + + result, err = pipeline.GetResults() + require.NoError(t, err) + require.IsType(t, &pgconn.PipelineSync{}, result) + + require.NoError(t, pipeline.Close()) +} + +// pgx asks for binary results, and pgdog decodes the rows itself to merge them. +// Read as text these ids sort with 10 before 2. +func TestShardedOrderByMergesBinaryRows(t *testing.T) { + ctx := context.Background() + + conn, err := pgx.Connect(ctx, testConnStr) + require.NoError(t, err) + defer conn.Close(ctx) + + _, err = conn.Exec(ctx, "TRUNCATE TABLE sharded") + require.NoError(t, err) + + want := []int64{} + for id := int64(1); id <= 20; id++ { + _, err = conn.Exec(ctx, "INSERT INTO sharded (id, value) VALUES ($1, $2)", id, "row") + require.NoError(t, err) + want = append(want, id) + } + + // Nothing is merged unless both shards hold rows. + assert.NotZero(t, countOnShardByComment(t, conn, 0, 1), "shard 0 holds no rows") + assert.NotZero(t, countOnShardByComment(t, conn, 1, 11), "shard 1 holds no rows") + + rows, err := conn.Query(ctx, "SELECT id FROM sharded ORDER BY id") + require.NoError(t, err) + + got := []int64{} + for rows.Next() { + var id int64 + require.NoError(t, rows.Scan(&id)) + got = append(got, id) + } + require.NoError(t, rows.Err()) + rows.Close() + + assert.Equal(t, want, got, "rows must arrive merge-sorted across shards") + + _, err = conn.Exec(ctx, "TRUNCATE TABLE sharded") + require.NoError(t, err) +} diff --git a/pgdog/src/frontend/client/query_engine/test/mod.rs b/pgdog/src/frontend/client/query_engine/test/mod.rs index e7f42285b..efea13559 100644 --- a/pgdog/src/frontend/client/query_engine/test/mod.rs +++ b/pgdog/src/frontend/client/query_engine/test/mod.rs @@ -22,6 +22,7 @@ mod lock_session; mod manual_lock; mod multi_binding; mod omni; +mod pipeline_execution; pub mod prelude; mod prepared_syntax_error; mod replicas; diff --git a/pgdog/src/frontend/client/query_engine/test/pipeline_execution.rs b/pgdog/src/frontend/client/query_engine/test/pipeline_execution.rs new file mode 100644 index 000000000..24f6d63bc --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/pipeline_execution.rs @@ -0,0 +1,564 @@ +use crate::expect_message; +use crate::net::{ + BindComplete, CommandComplete, DataRow, DataType, Message, ParameterDescription, Parameters, + ParseComplete, Protocol, ReadyForQuery, RowDescription, +}; + +use super::prelude::*; + +/// Assert the next reply message +macro_rules! assert_message { + ($messages:expr, $ty:ty) => {{ + let message = $messages + .next() + .unwrap_or_else(|| panic!("expected {}, got no more replies", stringify!($ty))); + + expect_message!(message, $ty) + }}; +} + +/// Send the request with message and wait for the whole stream of response messages +async fn send_and_wait( + client: &mut TestClient, + request: impl IntoIterator, +) -> Vec { + for message in request { + client.send(message).await; + } + client.try_process().await.unwrap(); + + client.read_until('Z').await.unwrap() +} + +/// Same as [`send_and_wait`] but won't fail on the first error +/// and will collect all the responses +async fn send_fail_and_wait( + client: &mut TestClient, + request: impl IntoIterator, +) -> Vec { + for message in request { + client.send(message).await; + } + let _ = client.try_process().await; + + let mut replies = vec![]; + loop { + let message = client.read().await; + let last = message.code() == 'Z'; + replies.push(message); + + if last { + return replies; + } + } +} + +/// The id in the next row. +fn next_id(messages: &mut impl Iterator, text: bool) -> i64 { + let row = assert_message!(messages, DataRow); + assert_eq!(row.len(), 1); + + row.get_int(0, text).unwrap() +} + +/// Parameter types a ParameterDescription must carry, in order. +fn assert_params(description: ParameterDescription, types: &[DataType]) { + let params = description + .params() + .iter() + .map(|oid| DataType::from_oid(*oid)) + .collect::>(); + + assert_eq!(params, types); +} + +/// Column names a RowDescription must carry, in order. +fn assert_columns(description: RowDescription, names: &[&str]) { + let columns = description + .fields + .iter() + .map(|field| field.name.as_str()) + .collect::>(); + + assert_eq!(columns, names); +} + +/// Two ids, one per shard, so a merged sort has to interleave them. The one on +/// shard 1 is the smaller number and the longer string, so numeric and text +/// order disagree. +fn ids_across_shards(client: &mut TestClient) -> (i64, i64) { + loop { + let shard_0 = client.random_id_for_shard(0); + let shard_1 = client.random_id_for_shard(1); + + if shard_1 < shard_0 && shard_1.to_string() > shard_0.to_string() { + return (shard_0, shard_1); + } + } +} + +/// One row per shard, under random ids, so a test that dies before it cleans +/// up cannot collide with another test. +async fn seed(client: &mut TestClient, (shard_0, shard_1): (i64, i64)) { + client + .send_simple(Query::new(format!( + "INSERT INTO sharded (id) VALUES ({shard_0}), ({shard_1})" + ))) + .await; + client.read_until('Z').await.unwrap(); +} + +async fn cleanup(client: &mut TestClient, (shard_0, shard_1): (i64, i64)) { + client + .send_simple(Query::new(format!( + "DELETE FROM sharded WHERE id IN ({shard_0}, {shard_1})" + ))) + .await; + client.read_until('Z').await.unwrap(); +} + +fn test_sql_sort(ids: (i64, i64)) -> String { + format!( + "SELECT id FROM sharded WHERE id IN ({}, {}) ORDER BY id", + ids.0, ids.1 + ) +} + +fn test_sql_sort_desc(ids: (i64, i64)) -> String { + format!("{} DESC", test_sql_sort(ids)) +} + +/// Two Parse and Describe pairs, then Sync. Nothing executes, so the router +/// sends the request to a single shard. +#[tokio::test] +async fn test_pipelined_parse_describe_sync_only() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let mut messages = send_and_wait( + &mut client, + vec![ + Parse::named("s1", "SELECT id FROM sharded WHERE 1 = 0").into(), + Describe::new_statement("s1").into(), + Parse::named( + "s2", + "SELECT id, $1::text AS second FROM sharded WHERE 1 = 0", + ) + .into(), + Describe::new_statement("s2").into(), + Sync.into(), + ], + ) + .await + .into_iter(); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, ParseComplete); + assert_params( + assert_message!(messages, ParameterDescription), + &[DataType::Text], + ); + assert_columns(assert_message!(messages, RowDescription), &["id", "second"]); + + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// The same shape, plus Bind and Execute. The request is executable, so it +/// stays on the cross-shard route of the first statement. +#[tokio::test] +async fn test_pipelined_parse_describe() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let mut messages = send_and_wait( + &mut client, + vec![ + Parse::named("c1", "SELECT id FROM sharded WHERE 1 = 0").into(), + Describe::new_statement("c1").into(), + Parse::named( + "c2", + "SELECT id, $1::text AS second FROM sharded WHERE 1 = 0", + ) + .into(), + Describe::new_statement("c2").into(), + Bind::new_statement("c1").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await + .into_iter(); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, ParseComplete); + assert_params( + assert_message!(messages, ParameterDescription), + &[DataType::Text], + ); + assert_columns(assert_message!(messages, RowDescription), &["id", "second"]); + + assert_message!(messages, BindComplete); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 0" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// Three pairs in a row, each describing different parameters and columns. +#[tokio::test] +async fn test_pipelined_parse_describe_three_pairs() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let mut messages = send_and_wait( + &mut client, + vec![ + Parse::named("c1", "SELECT id FROM sharded WHERE 1 = 0").into(), + Describe::new_statement("c1").into(), + Parse::named( + "c2", + "SELECT id, $1::text AS second FROM sharded WHERE 1 = 0", + ) + .into(), + Describe::new_statement("c2").into(), + Parse::named( + "c3", + "SELECT id, $1::text AS second, $2::bigint AS third FROM sharded WHERE 1 = 0", + ) + .into(), + Describe::new_statement("c3").into(), + Bind::new_statement("c1").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await + .into_iter(); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, ParseComplete); + assert_params( + assert_message!(messages, ParameterDescription), + &[DataType::Text], + ); + assert_columns(assert_message!(messages, RowDescription), &["id", "second"]); + + assert_message!(messages, ParseComplete); + assert_params( + assert_message!(messages, ParameterDescription), + &[DataType::Text, DataType::Bigint], + ); + assert_columns( + assert_message!(messages, RowDescription), + &["id", "second", "third"], + ); + + assert_message!(messages, BindComplete); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 0" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// A sorted cross-shard query buffers its rows, so pgdog decodes them. The +/// other Describe in the exchange names no sort column, and must not take the +/// decoder over. +#[tokio::test] +async fn test_pipelined_parse_describe_sorted() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let ids = ids_across_shards(&mut client); + seed(&mut client, ids).await; + + let replies = send_and_wait( + &mut client, + [ + ProtocolMessage::from(Parse::named("o1", test_sql_sort(ids))), + Describe::new_statement("o1").into(), + Parse::named("o2", "SELECT $1::text AS second FROM sharded").into(), + Describe::new_statement("o2").into(), + Bind::new_statement("o1").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + cleanup(&mut client, ids).await; + + let mut messages = replies.into_iter(); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, ParseComplete); + assert_params( + assert_message!(messages, ParameterDescription), + &[DataType::Text], + ); + assert_columns(assert_message!(messages, RowDescription), &["second"]); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, true), ids.1); + assert_eq!(next_id(&mut messages, true), ids.0); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// The other Describe names the sort column too, with another type. The rows +/// belong to the statement that executes, so they sort as bigint. +#[tokio::test] +async fn test_pipelined_parse_describe_shadowed_column_type() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let ids = ids_across_shards(&mut client); + seed(&mut client, ids).await; + + let replies = send_and_wait( + &mut client, + [ + ProtocolMessage::from(Parse::named("t1", test_sql_sort(ids))), + Describe::new_statement("t1").into(), + Parse::named("t2", "SELECT 'shadow'::text AS id FROM sharded").into(), + Describe::new_statement("t2").into(), + Bind::new_statement("t1").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + cleanup(&mut client, ids).await; + + let mut messages = replies.into_iter(); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, true), ids.1); + assert_eq!(next_id(&mut messages, true), ids.0); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// Two prepared statements execute in one exchange, as pgx does after its +/// prepare phase. Each sorted result set must reach the client whole, on its +/// own side of the CommandComplete that ends it. +#[tokio::test] +async fn test_pipelined_two_executes_keep_every_row() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let ids = ids_across_shards(&mut client); + seed(&mut client, ids).await; + + send_and_wait( + &mut client, + [ + ProtocolMessage::from(Parse::named("b1", test_sql_sort(ids))), + Describe::new_statement("b1").into(), + Parse::named("b2", test_sql_sort_desc(ids)).into(), + Describe::new_statement("b2").into(), + Sync.into(), + ], + ) + .await; + + // Both portals run before the client reads anything. + let replies = send_and_wait( + &mut client, + [ + ProtocolMessage::from(Bind::new_statement("b1")), + Execute::new().into(), + Bind::new_statement("b2").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + cleanup(&mut client, ids).await; + + let mut messages = replies.into_iter(); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, true), ids.1); + assert_eq!(next_id(&mut messages, true), ids.0); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, true), ids.0); + assert_eq!(next_id(&mut messages, true), ids.1); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// The client asks for binary results, as tokio-postgres does. The Bind is the +/// only thing telling pgdog these buffered rows are binary. +#[tokio::test] +async fn test_pipelined_binary_results_sort_as_binary() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let ids = ids_across_shards(&mut client); + seed(&mut client, ids).await; + + // Naming the statement caches its RowDescription, so the execute phase + // sends no Parse and no Describe. + send_and_wait( + &mut client, + [ + ProtocolMessage::from(Parse::named("y1", test_sql_sort(ids))), + Describe::new_statement("y1").into(), + Sync.into(), + ], + ) + .await; + + let replies = send_and_wait( + &mut client, + [ + ProtocolMessage::from(Bind::new_params_codes_results("y1", &[], &[], &[1])), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + cleanup(&mut client, ids).await; + + let mut messages = replies.into_iter(); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, false), ids.1); + assert_eq!(next_id(&mut messages, false), ids.0); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// An anonymous Bind has no cached RowDescription, so the Describe of the same +/// statement is the only thing that can type its rows. This is the shape pgdog +/// sends when it rewrites a sharding key update. +#[tokio::test] +async fn test_pipelined_anonymous_bind_sorts() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let ids = ids_across_shards(&mut client); + seed(&mut client, ids).await; + + let replies = send_and_wait( + &mut client, + [ + ProtocolMessage::from(Parse::new_anonymous(&test_sql_sort(ids))), + Describe::new_statement("").into(), + Bind::new_statement("").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + cleanup(&mut client, ids).await; + + let mut messages = replies.into_iter(); + + assert_message!(messages, ParseComplete); + assert_params(assert_message!(messages, ParameterDescription), &[]); + assert_columns(assert_message!(messages, RowDescription), &["id"]); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, true), ids.1); + assert_eq!(next_id(&mut messages, true), ids.0); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} + +/// An Execute that fails at run time abandons the rest of the exchange, and +/// the session still sorts the next one. +#[tokio::test] +async fn test_pipelined_error_does_not_leak_into_the_next_exchange() { + let mut client = TestClient::new_sharded(Parameters::default()).await; + let ids = ids_across_shards(&mut client); + seed(&mut client, ids).await; + + let failing = format!( + "SELECT id / 0 FROM sharded WHERE id IN ({}, {})", + ids.0, ids.1 + ); + + let replies = send_fail_and_wait( + &mut client, + [ + ProtocolMessage::from(Parse::named("f1", failing)), + Parse::named("f2", test_sql_sort(ids)).into(), + Describe::new_statement("f2").into(), + Bind::new_statement("f1").into(), + Execute::new().into(), + Bind::new_statement("f2").into(), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + // Every shard reports the error, so the count is not pinned here. + assert!(replies.iter().any(|message| message.code() == 'E')); + assert_eq!(replies.last().unwrap().code(), 'Z'); + assert!(!replies.iter().any(|message| message.code() == 'D')); + + let replies = send_and_wait( + &mut client, + [ + ProtocolMessage::from(Bind::new_statement("f2")), + Execute::new().into(), + Sync.into(), + ], + ) + .await; + + cleanup(&mut client, ids).await; + + let mut messages = replies.into_iter(); + + assert_message!(messages, BindComplete); + assert_eq!(next_id(&mut messages, true), ids.1); + assert_eq!(next_id(&mut messages, true), ids.0); + assert_eq!( + assert_message!(messages, CommandComplete).command(), + "SELECT 2" + ); + assert_eq!(assert_message!(messages, ReadyForQuery).status, 'I'); + assert!(messages.next().is_none()); +} diff --git a/pgdog/src/frontend/client/test/test_client.rs b/pgdog/src/frontend/client/test/test_client.rs index 5ddae8f1d..f026ddb9c 100644 --- a/pgdog/src/frontend/client/test/test_client.rs +++ b/pgdog/src/frontend/client/test/test_client.rs @@ -260,21 +260,25 @@ impl TestClient { pid } + /// The shard an ID lands on. + pub(crate) fn shard_for_id(&mut self, id: i64) -> Shard { + let cluster = self.engine.backend().cluster().unwrap(); + + ContextBuilder::new(cluster.sharded_tables().first().unwrap()) + .data(id) + .shards(cluster.shards().len()) + .build() + .unwrap() + .apply() + .unwrap() + } + /// Generate a random ID for a given shard. pub(crate) fn random_id_for_shard(&mut self, shard: usize) -> i64 { - let cluster = self.engine.backend().cluster().unwrap().clone(); - loop { let id: i64 = rng().random(); - let calc = ContextBuilder::new(cluster.sharded_tables().first().unwrap()) - .data(id) - .shards(cluster.shards().len()) - .build() - .unwrap() - .apply() - .unwrap(); - if calc == Shard::Direct(shard) { + if self.shard_for_id(id) == Shard::Direct(shard) { return id; } } diff --git a/pgdog/src/net/messages/parameter_description.rs b/pgdog/src/net/messages/parameter_description.rs index 8229cadd6..3224f69a8 100644 --- a/pgdog/src/net/messages/parameter_description.rs +++ b/pgdog/src/net/messages/parameter_description.rs @@ -49,6 +49,11 @@ impl ParameterDescription { Self { params } } + /// Type OIDs of the parameters, in order. + pub fn params(&self) -> &[i32] { + &self.params + } + pub(crate) fn rewrite_data_types(&mut self, mapping: &HashMap) { for param in &mut self.params { if let Some(&canonical) = mapping.get(&(*param as u32)) { From 8281a572e4f3056a1e43c07800d7e3e9e2990fd7 Mon Sep 17 00:00:00 2001 From: meskill <8974488+meskill@users.noreply.github.com> Date: Fri, 14 Aug 2026 11:58:42 +0000 Subject: [PATCH 2/2] the fix --- .../src/backend/pool/connection/aggregate.rs | 10 +- pgdog/src/backend/pool/connection/binding.rs | 4 +- pgdog/src/backend/pool/connection/buffer.rs | 22 +-- pgdog/src/backend/pool/connection/mod.rs | 2 +- .../pool/connection/multi_shard/context.rs | 19 --- .../pool/connection/multi_shard/mod.rs | 61 ++++---- .../pool/connection/multi_shard/test.rs | 99 ++++++++++++- .../pool/connection/multi_shard/validator.rs | 8 +- pgdog/src/net/decoder.rs | 130 ++++++++++++++---- pgdog/src/net/messages/bind.rs | 11 ++ pgdog/src/net/messages/data_row.rs | 4 +- 11 files changed, 266 insertions(+), 104 deletions(-) delete mode 100644 pgdog/src/backend/pool/connection/multi_shard/context.rs diff --git a/pgdog/src/backend/pool/connection/aggregate.rs b/pgdog/src/backend/pool/connection/aggregate.rs index da06817c0..878b38fff 100644 --- a/pgdog/src/backend/pool/connection/aggregate.rs +++ b/pgdog/src/backend/pool/connection/aggregate.rs @@ -515,7 +515,7 @@ mod test { let aggregate = parse("SELECT COUNT(*)::int FROM users"); let rd = RowDescription::new(&[integer_field("count")]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let mut rows = VecDeque::new(); let mut shard0 = DataRow::new(); @@ -544,7 +544,7 @@ mod test { let aggregate = parse("SELECT AVG(price) FROM menu"); let rd = RowDescription::new(&[Field::double("avg")]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let mut rows = VecDeque::new(); let mut shard0 = DataRow::new(); @@ -578,7 +578,7 @@ mod test { let aggregate = parse("SELECT price, SUM(quantity) FROM menu GROUP BY 1"); let rd = RowDescription::new(&[Field::double("price"), Field::bigint("sum")]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let mut rows = VecDeque::new(); let mut shard0 = DataRow::new(); @@ -620,7 +620,7 @@ mod test { let aggregate = parse("SELECT matrix, COUNT(*) FROM samples GROUP BY 1"); let rd = RowDescription::new(&[integer_array_field("matrix"), Field::bigint("count")]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let mut rows = VecDeque::new(); @@ -670,7 +670,7 @@ mod test { interval_array_field("sample_interval_array"), Field::bigint("count"), ]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let input = Bytes::from_static(br#"{"1 year 2 mons 1 day 04:05:06.7"}"#); diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index faa5fe28a..0abff75f3 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -299,7 +299,9 @@ impl Binding { match self { Binding::Admin(admin) => !admin.done(), Binding::Direct(server, ..) => server.has_more_messages(), - Binding::MultiShard(servers, _state) => servers.iter().any(|s| s.has_more_messages()), + Binding::MultiShard(servers, state) => { + state.has_more_messages() || servers.iter().any(|s| s.has_more_messages()) + } _ => false, } } diff --git a/pgdog/src/backend/pool/connection/buffer.rs b/pgdog/src/backend/pool/connection/buffer.rs index edea8fd85..2fc71cfe5 100644 --- a/pgdog/src/backend/pool/connection/buffer.rs +++ b/pgdog/src/backend/pool/connection/buffer.rs @@ -216,6 +216,11 @@ impl Buffer { } } + /// The buffer holds rows the client can read now. + pub(super) fn can_take(&self) -> bool { + self.full && !self.is_empty() + } + /// Execute LIMIT ... OFFSET ... pub(super) fn limit(&mut self, limit: &Limit) { let offset = limit.offset.unwrap_or(0); @@ -230,7 +235,6 @@ impl Buffer { self.buffer.len() } - #[allow(dead_code)] pub(super) fn is_empty(&self) -> bool { self.len() == 0 } @@ -254,7 +258,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); buf.sort(&columns, &decoder); buf.full(); @@ -284,7 +288,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - buf.aggregate(&agg, &Decoder::from(&rd), &AggregateRewritePlan::default()) + buf.aggregate(&agg, &Decoder::from(rd), &AggregateRewritePlan::default()) .unwrap(); buf.full(); @@ -311,7 +315,7 @@ mod test { } } - buf.aggregate(&agg, &Decoder::from(&rd), &AggregateRewritePlan::default()) + buf.aggregate(&agg, &Decoder::from(rd), &AggregateRewritePlan::default()) .unwrap(); buf.full(); @@ -345,7 +349,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); buf.sort(&columns, &decoder); buf.full(); @@ -385,7 +389,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); buf.sort(&columns, &decoder); buf.full(); @@ -441,7 +445,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); buf.sort(&columns, &decoder); buf.full(); @@ -485,7 +489,7 @@ mod test { buf.add(dr.message().unwrap()).unwrap(); } - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); buf.sort(&columns, &decoder); buf.full(); @@ -577,7 +581,7 @@ mod test { fn test_distinct() { let mut buf = Buffer::default(); let rd = RowDescription::new(&[Field::bigint("id"), Field::text("email")]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); for email in ["test@test.com", "apples@test.com", "domain@test.com"] { for i in 0..5 { diff --git a/pgdog/src/backend/pool/connection/mod.rs b/pgdog/src/backend/pool/connection/mod.rs index c41e3c045..24067640f 100644 --- a/pgdog/src/backend/pool/connection/mod.rs +++ b/pgdog/src/backend/pool/connection/mod.rs @@ -414,7 +414,7 @@ impl Connection { pub(crate) fn bind(&mut self, bind: &Bind) -> Result<(), Error> { match self.binding { Binding::MultiShard(_, ref mut state) => { - state.set_context(bind); + state.set_bind_context(bind); Ok(()) } diff --git a/pgdog/src/backend/pool/connection/multi_shard/context.rs b/pgdog/src/backend/pool/connection/multi_shard/context.rs deleted file mode 100644 index fa40667e0..000000000 --- a/pgdog/src/backend/pool/connection/multi_shard/context.rs +++ /dev/null @@ -1,19 +0,0 @@ -use crate::net::{Bind, RowDescription}; - -#[derive(Debug, Clone)] -pub enum Context<'a> { - Bind(&'a Bind), - RowDescription(&'a RowDescription), -} - -impl<'a> From<&'a RowDescription> for Context<'a> { - fn from(value: &'a RowDescription) -> Self { - Context::RowDescription(value) - } -} - -impl<'a> From<&'a Bind> for Context<'a> { - fn from(value: &'a Bind) -> Self { - Context::Bind(value) - } -} diff --git a/pgdog/src/backend/pool/connection/multi_shard/mod.rs b/pgdog/src/backend/pool/connection/multi_shard/mod.rs index 8c00ee2a2..829c929cd 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/mod.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/mod.rs @@ -1,11 +1,11 @@ //! Multi-shard connection state. -use context::Context; +use std::collections::VecDeque; use crate::{ - frontend::{PreparedStatements, router::Route}, + frontend::router::Route, net::{ - BackendPid, Decoder, ReadyForQuery, + BackendPid, Bind, Decoder, ReadyForQuery, messages::{ DataRow, FromBytes, Message, Protocol, RowDescription, ToBytes, command_complete::CommandComplete, @@ -15,7 +15,6 @@ use crate::{ use super::buffer::Buffer; -mod context; mod error; #[cfg(test)] mod test; @@ -65,6 +64,8 @@ pub struct MultiShard { decoder: Decoder, /// Row consistency validator. validator: Validator, + /// Binds waiting for their BindComplete, in the order they were sent. + bound_statements: VecDeque, } impl MultiShard { @@ -110,6 +111,7 @@ impl MultiShard { // 1. Route to keep routing decision // 2. Number of shards // 3. Decoder + // 4. Pending Binds, pushed before this runs } /// Check if the message should be sent to the client, skipped, @@ -125,6 +127,8 @@ impl MultiShard { } forward = if self.counters.ready_for_query.is_multiple_of(self.shards) { + self.bound_statements.clear(); + if self.counters.transaction_error { Some(ReadyForQuery::error().message()?) } else { @@ -194,16 +198,17 @@ impl MultiShard { self.counters.row_description += 1; let rd = RowDescription::from_bytes(message.to_bytes())?; - // Validate row description consistency - let is_first = self.validator.validate_row_description(&rd)?; + // Validate row description consistency inside the group. + // we reset after shards processed the group of the same RD + let is_first_in_group = self.validator.validate_row_description(&rd)?; - // Set row description info as soon as we have it, - // so it's available to the aggregator and sorter. - if is_first { - self.decoder.row_description(&rd); + // Set it as soon as we have it, so the aggregator and sorter + // can use it. + if is_first_in_group { + self.decoder.set_row_description(rd.clone()); } - if self.counters.row_description == self.shards { + if self.counters.row_description.is_multiple_of(self.shards) { // Only send it to the client once all shards sent it, // so we don't get early requests from clients. let plan = self.route.aggregate_rewrite_plan(); @@ -213,6 +218,9 @@ impl MultiShard { let client_rd = rd.drop_columns(plan.drop_columns()); forward = Some(client_rd.message()?); } + + // The next statement describes a different result set. + self.validator.reset(); } } @@ -289,6 +297,10 @@ impl MultiShard { if self.counters.bind_complete.is_multiple_of(self.shards) { forward = Some(message); + + if let Some(bind) = self.bound_statements.pop_front() { + self.decoder.bind(&bind); + } } } @@ -346,25 +358,12 @@ impl MultiShard { } } - pub(super) fn set_context<'a>(&mut self, message: impl Into>) { - let context = message.into(); - match context { - Context::Bind(bind) => { - if self.decoder.rd().fields.is_empty() - && !bind.anonymous() - && let Some(rd) = PreparedStatements::global() - .read() - .row_description(bind.statement()) - { - self.decoder.row_description(&rd); - self.validator.set_row_description(&rd); - } - self.decoder.bind(bind); - } - Context::RowDescription(rd) => { - self.decoder.row_description(rd); - self.validator.set_row_description(rd); - } - } + /// Sorted rows or a merged CommandComplete are waiting for the client. + pub(super) fn has_more_messages(&self) -> bool { + self.buffer.can_take() || self.counters.command_complete.is_some() + } + + pub(super) fn set_bind_context(&mut self, bind: &Bind) { + self.bound_statements.push_back(bind.clone()); } } diff --git a/pgdog/src/backend/pool/connection/multi_shard/test.rs b/pgdog/src/backend/pool/connection/multi_shard/test.rs index 65eb3fe32..afb72bedd 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/test.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/test.rs @@ -1,6 +1,6 @@ use crate::{ frontend::router::parser::{Shard, ShardWithPriority}, - net::{DataRow, Field}, + net::{BindComplete, DataRow, Field, Format}, }; use super::*; @@ -267,3 +267,100 @@ fn test_omni_data_rows_only_from_first_server() { .unwrap(); assert!(result.is_some()); // Should be forwarded } + +/// Statements pipelined in one exchange each get their own RowDescription, +/// and they may describe different result sets. +#[test] +fn test_pipelined_describe_forwards_every_group() { + for shards in [1, 2] { + let mut multi_shard = MultiShard::new( + (0..shards).collect(), + &Route::read(ShardWithPriority::new_default_unset(Shard::All)), + ); + + let id = RowDescription::new(&[Field::bigint("id")]); + let name = RowDescription::new(&[Field::text("name"), Field::bigint("id")]); + let mut forwarded = vec![]; + + for description in [&id, &name] { + for _ in 0..shards { + if let Some(message) = multi_shard.forward(description.message().unwrap()).unwrap() + { + forwarded.push(message); + } + } + } + + assert_eq!( + forwarded, + vec![id.message().unwrap(), name.message().unwrap()], + "{shards} shard(s)", + ); + } +} + +/// Each Bind in a pipelined exchange decides the wire format of its own rows, +/// and a server RowDescription must not take that over. +#[test] +fn test_bind_result_formats_apply_per_statement() { + let mut multi_shard = MultiShard::new( + vec![0, 1], + &Route::read(ShardWithPriority::new_default_unset(Shard::All)), + ); + + let binary = Bind::new_params_codes_results("b1", &[], &[], &[1]); + let text = Bind::new_statement("b2"); + let rd = RowDescription::new(&[Field::bigint("id")]); + + multi_shard.set_bind_context(&binary); + multi_shard.set_bind_context(&text); + + for _ in 0..2 { + multi_shard + .forward(BindComplete.message().unwrap()) + .unwrap(); + } + for _ in 0..2 { + multi_shard.forward(rd.message().unwrap()).unwrap(); + } + assert_eq!(multi_shard.decoder.format(0), Format::Binary); + + // The second statement asked for no formats. + for _ in 0..2 { + multi_shard + .forward(BindComplete.message().unwrap()) + .unwrap(); + } + for _ in 0..2 { + multi_shard.forward(rd.message().unwrap()).unwrap(); + } + assert_eq!(multi_shard.decoder.format(0), Format::Text); +} + +/// A Bind that never completed is dropped at ReadyForQuery, so the next +/// exchange cannot pop it and type its rows after the wrong statement. +#[test] +fn test_ready_for_query_drops_pending_binds() { + let mut multi_shard = MultiShard::new( + vec![0, 1], + &Route::read(ShardWithPriority::new_default_unset(Shard::All)), + ); + + multi_shard.set_bind_context(&Bind::new_statement("b1")); + multi_shard.set_bind_context(&Bind::new_statement("b2")); + + // Only the first statement binds. The second is abandoned. + for _ in 0..2 { + multi_shard + .forward(BindComplete.message().unwrap()) + .unwrap(); + } + assert_eq!(multi_shard.bound_statements.len(), 1); + + for _ in 0..2 { + multi_shard + .forward(ReadyForQuery::idle().message().unwrap()) + .unwrap(); + } + assert!(multi_shard.bound_statements.is_empty()); +} diff --git a/pgdog/src/backend/pool/connection/multi_shard/validator.rs b/pgdog/src/backend/pool/connection/multi_shard/validator.rs index 95b6cfd00..44933c08e 100644 --- a/pgdog/src/backend/pool/connection/multi_shard/validator.rs +++ b/pgdog/src/backend/pool/connection/multi_shard/validator.rs @@ -14,17 +14,13 @@ pub(super) struct Validator { } impl Validator { - /// Reset the validator state. + /// Reset the validator state, so the next statement in the exchange starts + /// a new comparison. pub(super) fn reset(&mut self) { self.first_row_description = None; self.expected_column_count = None; } - /// Set the row description. - pub(super) fn set_row_description(&mut self, rd: &RowDescription) { - self.first_row_description = Some(rd.clone()); - } - /// Validate a row description against the first one received. /// Returns true if this is the first row description, false if it's a duplicate that matches. pub(super) fn validate_row_description(&mut self, rd: &RowDescription) -> Result { diff --git a/pgdog/src/net/decoder.rs b/pgdog/src/net/decoder.rs index 44169ab8f..fcb8da94f 100644 --- a/pgdog/src/net/decoder.rs +++ b/pgdog/src/net/decoder.rs @@ -2,24 +2,13 @@ use crate::frontend::PreparedStatements; use super::{Bind, Format, RowDescription}; -impl From<&Bind> for Decoder { - fn from(value: &Bind) -> Self { - let mut decoder = Decoder::new(); - decoder.bind(value); - decoder - } -} - -impl From<&RowDescription> for Decoder { - fn from(value: &RowDescription) -> Self { - let mut decoder = Decoder::new(); - decoder.row_description(value); - decoder - } -} - +/// Decodes result columns. +/// +/// The server owns the column names and types, the client's Bind owns the +/// wire formats. Neither may overwrite the other. #[derive(Debug, Clone, Default)] pub struct Decoder { + /// Formats requested by the client's Bind. formats: Vec, rd: RowDescription, } @@ -30,15 +19,14 @@ impl Decoder { Self::default() } - /// Infer types from Bind, if any provided. + /// Describe the rows a Bind is about to produce. Only a statement the + /// cache knows replaces the columns, because clearing them would leave + /// the sorter with no types at all. pub fn bind(&mut self, bind: &Bind) { - // Only override RowDescription formats if - // Bind specifies formats. - if !bind.codes().is_empty() { - self.formats = bind.codes().to_vec(); - } + self.formats.clear(); + self.formats.extend(bind.result_formats()); - if self.rd.is_empty() + if !bind.anonymous() && let Some(rd) = PreparedStatements::global() .read() .row_description(bind.statement()) @@ -47,17 +35,19 @@ impl Decoder { } } - /// Infer types from RowDescription, if any. - pub fn row_description(&mut self, rd: &RowDescription) { - let formats = rd.fields.iter().map(|f| f.format()).collect(); - self.formats = formats; - self.rd = rd.clone(); + /// Describe the rows the server just announced. + pub fn set_row_description(&mut self, rd: RowDescription) { + self.rd = rd; } /// Get format used for column at position. pub fn format(&self, position: usize) -> Format { match self.formats.len() { - 0 => Format::Text, + 0 => self + .rd + .field(position) + .map(|field| field.format()) + .unwrap_or(Format::Text), 1 => self.formats[0], _ => self.formats.get(position).copied().unwrap_or(Format::Text), } @@ -67,3 +57,85 @@ impl Decoder { &self.rd } } + +#[cfg(test)] +mod test_impls { + use super::*; + + impl From for Decoder { + fn from(value: RowDescription) -> Self { + let mut decoder = Decoder::new(); + decoder.set_row_description(value); + decoder + } + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::net::messages::Field; + + fn text_rd() -> RowDescription { + RowDescription::new(&[Field::bigint("id"), Field::text("name")]) + } + + #[test] + fn test_row_description_decides_without_bind() { + let mut decoder = Decoder::new(); + decoder.set_row_description(text_rd()); + + assert_eq!(decoder.format(0), Format::Text); + assert_eq!(decoder.format(1), Format::Text); + } + + #[test] + fn test_bind_result_formats_survive_row_description() { + let mut decoder = Decoder::new(); + let bind = Bind::new_params_codes_results("s1", &[], &[], &[1, 0]); + + decoder.bind(&bind); + decoder.set_row_description(text_rd()); + + assert_eq!(decoder.format(0), Format::Binary); + assert_eq!(decoder.format(1), Format::Text); + } + + #[test] + fn test_row_description_before_bind_gives_the_same_answer() { + let mut decoder = Decoder::new(); + let bind = Bind::new_params_codes_results("s1", &[], &[], &[1, 0]); + + decoder.set_row_description(text_rd()); + decoder.bind(&bind); + + assert_eq!(decoder.format(0), Format::Binary); + assert_eq!(decoder.format(1), Format::Text); + } + + #[test] + fn test_one_result_format_applies_to_every_column() { + let mut decoder = Decoder::new(); + decoder.set_row_description(text_rd()); + decoder.bind(&Bind::new_params_codes_results("s1", &[], &[], &[1])); + + assert_eq!(decoder.format(0), Format::Binary); + assert_eq!(decoder.format(1), Format::Binary); + } + + #[test] + fn test_bind_keeps_the_server_description_when_the_cache_is_empty() { + let mut decoder = Decoder::new(); + decoder.set_row_description(text_rd()); + decoder.bind(&Bind::new_params_codes_results("s1", &[], &[], &[1])); + + // Nothing described s1, so the only columns we have are the ones the + // server announced. The formats still come from the Bind. + assert_eq!(decoder.rd().fields.len(), 2); + assert_eq!(decoder.format(0), Format::Binary); + + decoder.bind(&Bind::new_statement("s2")); + + assert_eq!(decoder.format(0), Format::Text); + } +} diff --git a/pgdog/src/net/messages/bind.rs b/pgdog/src/net/messages/bind.rs index ebe0d6863..acd6145c3 100644 --- a/pgdog/src/net/messages/bind.rs +++ b/pgdog/src/net/messages/bind.rs @@ -204,6 +204,17 @@ impl Bind { unsafe { from_utf8_unchecked(&self.statement[0..self.statement.len() - 1]) } } + /// Format the client asked each result column to be returned in. + pub fn result_formats(&self) -> impl ExactSizeIterator + '_ { + self.results.chunks_exact(2).map(|code| { + if i16::from_be_bytes([code[0], code[1]]) == 0 { + Format::Text + } else { + Format::Binary + } + }) + } + /// Format codes, if any. pub fn codes(&self) -> &[Format] { &self.codes diff --git a/pgdog/src/net/messages/data_row.rs b/pgdog/src/net/messages/data_row.rs index dee6df550..b33112f31 100644 --- a/pgdog/src/net/messages/data_row.rs +++ b/pgdog/src/net/messages/data_row.rs @@ -306,7 +306,7 @@ mod test { type_modifier: -1, format: 0, }]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let row = DataRow::from_columns(vec![Bytes::from_static(b"{1,2,3}")]); let column = row.get_column(0, &decoder).unwrap().unwrap(); @@ -321,7 +321,7 @@ mod test { #[test] fn get_column_checked_returns_err_on_missing_field() { let rd = RowDescription::new(&[Field::bigint("stuff")]); - let decoder = Decoder::from(&rd); + let decoder = Decoder::from(rd); let row = DataRow::from_columns(vec![Bytes::from_static(b"1")]); let column = row.get_column_checked(0, &decoder);