From 4419321733ab266def2c3fb04ff708b45e2a59e6 Mon Sep 17 00:00:00 2001 From: Yogesh Singla Date: Sun, 16 Aug 2026 23:01:04 +0530 Subject: [PATCH] a trail for perf --- pgdog/src/backend/prepared_statements.rs | 120 ++++++++++++++--------- pgdog/src/net/messages/sync.rs | 6 +- 2 files changed, 77 insertions(+), 49 deletions(-) diff --git a/pgdog/src/backend/prepared_statements.rs b/pgdog/src/backend/prepared_statements.rs index 17665d0bd..b4496aff2 100644 --- a/pgdog/src/backend/prepared_statements.rs +++ b/pgdog/src/backend/prepared_statements.rs @@ -76,15 +76,19 @@ impl Prepare { } } +/// Payloads are boxed so the enum stays pointer-sized. It is returned from +/// [`PreparedStatements::handle`] for every protocol message, and the answer is +/// `Forward` for anything that isn't a named prepared statement; carrying the +/// variants inline made that 528 bytes to move per message. #[derive(Debug, Clone, PartialEq)] pub(super) enum HandleResult { Drop, Forward, - Rewrite(ProtocolMessage), - Prepend(Prepare), + Rewrite(Box), + Prepend(Box), PrependRewrite { - prepend: Prepare, - rewrite: ProtocolMessage, + prepend: Box, + rewrite: Box, }, } @@ -179,11 +183,11 @@ impl PreparedStatements { let mut bind = bind.clone(); bind.anonymize(); return Ok(HandleResult::PrependRewrite { - prepend: message, - rewrite: ProtocolMessage::Bind(bind), + prepend: Box::new(message), + rewrite: Box::new(ProtocolMessage::Bind(bind)), }); } else { - return Ok(HandleResult::Prepend(message)); + return Ok(HandleResult::Prepend(Box::new(message))); } } @@ -192,7 +196,9 @@ impl PreparedStatements { if self.config.level.rewrite_anonymous() { let mut bind = bind.clone(); bind.anonymize(); - return Ok(HandleResult::Rewrite(ProtocolMessage::Bind(bind))); + return Ok(HandleResult::Rewrite(Box::new(ProtocolMessage::Bind( + bind, + )))); } } } @@ -224,11 +230,11 @@ impl PreparedStatements { let mut describe = describe.clone(); describe.anonymize(); return Ok(HandleResult::PrependRewrite { - prepend: message, - rewrite: ProtocolMessage::Describe(describe), + prepend: Box::new(message), + rewrite: Box::new(ProtocolMessage::Describe(describe)), }); } else { - return Ok(HandleResult::Prepend(message)); + return Ok(HandleResult::Prepend(Box::new(message))); } } @@ -240,8 +246,8 @@ impl PreparedStatements { if self.config.level.rewrite_anonymous() { let mut describe = describe.clone(); describe.anonymize(); - return Ok(HandleResult::Rewrite(ProtocolMessage::Describe( - describe, + return Ok(HandleResult::Rewrite(Box::new( + ProtocolMessage::Describe(describe), ))); } } @@ -288,7 +294,9 @@ impl PreparedStatements { self.state.add('1'); if rewritten { - return Ok(HandleResult::Rewrite(ProtocolMessage::Parse(parse))); + return Ok(HandleResult::Rewrite(Box::new(ProtocolMessage::Parse( + parse, + )))); } } @@ -633,6 +641,17 @@ pub(crate) mod test { }; use pgdog_config::PreparedStatements as PreparedStatementsLevel; + /// Returned by value for every protocol message, so it is moved on the + /// hot path whether or not it carries a payload. It was 528 bytes when + /// the variants held `Prepare`/`ProtocolMessage` inline, which was + /// visible as `memcpy` under `Server::send_one` in profiles. + #[test] + fn test_handle_result_stays_small() { + let size = size_of::(); + println!("HandleResult = {} bytes", size); + assert!(size <= 32, "HandleResult grew to {} bytes", size); + } + /// Build a PreparedStatements instance configured for ExtendedAnonymous mode. fn new_extended_anonymous() -> PreparedStatements { new_with_level(PreparedStatementsLevel::ExtendedAnonymous) @@ -882,13 +901,14 @@ pub(crate) mod test { let mut ps = new_extended_anonymous(); let parse = Parse::named("stmt1", "SELECT 1"); let result = ps.handle(&ProtocolMessage::Parse(parse)).unwrap(); - match result { - HandleResult::Rewrite(ProtocolMessage::Parse(p)) => { - assert!(p.anonymous(), "Parse should be anonymized"); - assert_eq!(p.query(), "SELECT 1"); - } - other => panic!("expected Rewrite(Parse), got {:?}", other), - } + let HandleResult::Rewrite(msg) = &result else { + panic!("expected Rewrite(Parse), got {:?}", result); + }; + let ProtocolMessage::Parse(p) = &**msg else { + panic!("expected Rewrite(Parse), got {:?}", result); + }; + assert!(p.anonymous(), "Parse should be anonymized"); + assert_eq!(p.query(), "SELECT 1"); } #[test] @@ -938,12 +958,13 @@ pub(crate) mod test { let mut ps = new_extended_anonymous(); let bind = Bind::new_statement("stmt1"); let result = ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); - match result { - HandleResult::Rewrite(ProtocolMessage::Bind(b)) => { - assert!(b.anonymous(), "Bind should be anonymized"); - } - other => panic!("expected Rewrite(Bind), got {:?}", other), - } + let HandleResult::Rewrite(msg) = &result else { + panic!("expected Rewrite(Bind), got {:?}", result); + }; + let ProtocolMessage::Bind(b) = &**msg else { + panic!("expected Rewrite(Bind), got {:?}", result); + }; + assert!(b.anonymous(), "Bind should be anonymized"); } #[test] @@ -988,7 +1009,7 @@ pub(crate) mod test { panic!("expected prepend to be Parse"); } // The rewritten Bind should be anonymized. - if let ProtocolMessage::Bind(b) = &rewrite { + if let ProtocolMessage::Bind(b) = &*rewrite { assert!(b.anonymous(), "rewritten Bind should be anonymous"); } else { panic!("expected rewrite to be Bind"); @@ -1007,12 +1028,13 @@ pub(crate) mod test { let mut ps = new_extended_anonymous(); let describe = Describe::new_statement("stmt1"); let result = ps.handle(&ProtocolMessage::Describe(describe)).unwrap(); - match result { - HandleResult::Rewrite(ProtocolMessage::Describe(d)) => { - assert!(d.anonymous(), "Describe should be anonymized"); - } - other => panic!("expected Rewrite(Describe), got {:?}", other), - } + let HandleResult::Rewrite(msg) = &result else { + panic!("expected Rewrite(Describe), got {:?}", result); + }; + let ProtocolMessage::Describe(d) = &**msg else { + panic!("expected Rewrite(Describe), got {:?}", result); + }; + assert!(d.anonymous(), "Describe should be anonymized"); } #[test] @@ -1028,7 +1050,7 @@ pub(crate) mod test { } else { panic!("expected prepend to be Parse"); } - if let ProtocolMessage::Describe(d) = &rewrite { + if let ProtocolMessage::Describe(d) = &*rewrite { assert!(d.anonymous(), "rewritten Describe should be anonymous"); } else { panic!("expected rewrite to be Describe"); @@ -1123,12 +1145,13 @@ pub(crate) mod test { // Parse let parse = Parse::named("stmt1", "SELECT $1"); let result = ps.handle(&ProtocolMessage::Parse(parse)).unwrap(); - match &result { - HandleResult::Rewrite(ProtocolMessage::Parse(p)) => { - assert!(p.anonymous()); - } - other => panic!("expected Rewrite(Parse), got {:?}", other), - } + let HandleResult::Rewrite(msg) = &result else { + panic!("expected Rewrite(Parse), got {:?}", result); + }; + let ProtocolMessage::Parse(p) = &**msg else { + panic!("expected Rewrite(Parse), got {:?}", result); + }; + assert!(p.anonymous()); // Bind let bind = Bind::new_params( @@ -1139,12 +1162,13 @@ pub(crate) mod test { }], ); let result = ps.handle(&ProtocolMessage::Bind(bind)).unwrap(); - match &result { - HandleResult::Rewrite(ProtocolMessage::Bind(b)) => { - assert!(b.anonymous()); - } - other => panic!("expected Rewrite(Bind), got {:?}", other), - } + let HandleResult::Rewrite(msg) = &result else { + panic!("expected Rewrite(Bind), got {:?}", result); + }; + let ProtocolMessage::Bind(b) = &**msg else { + panic!("expected Rewrite(Bind), got {:?}", result); + }; + assert!(b.anonymous()); // Execute let result = ps @@ -1195,7 +1219,7 @@ pub(crate) mod test { let expected = parse.with_data_types(&[10001, 10002]); assert_eq!( result, - HandleResult::Rewrite(ProtocolMessage::Parse(expected)) + HandleResult::Rewrite(Box::new(ProtocolMessage::Parse(expected))) ); } diff --git a/pgdog/src/net/messages/sync.rs b/pgdog/src/net/messages/sync.rs index c55c7b4f0..8afb5a2aa 100644 --- a/pgdog/src/net/messages/sync.rs +++ b/pgdog/src/net/messages/sync.rs @@ -27,9 +27,12 @@ impl FromBytes for Sync { } } +/// Sync has no payload, so its encoding is the constant 'S' + length 4. +static ENCODED: &[u8] = &[b'S', 0, 0, 0, 4]; + impl ToBytes for Sync { fn to_bytes(&self) -> Bytes { - Payload::named('S').freeze() + Bytes::from_static(ENCODED) } } @@ -46,5 +49,6 @@ mod test { #[test] fn test_sync() { assert_eq!(Sync.len(), Sync.to_bytes().len()); + assert_eq!(Sync.to_bytes(), Payload::named('S').freeze()); } }