Skip to content

Perf trial #1370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ygxio wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Perf trial #1370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 72 additions & 48 deletions pgdog/src/backend/prepared_statements.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ProtocolMessage>),
Prepend(Box<Prepare>),
PrependRewrite {
prepend: Prepare,
rewrite: ProtocolMessage,
prepend: Box<Prepare>,
rewrite: Box<ProtocolMessage>,
},
}

Expand Down Expand Up @@ -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)));
}
}

Expand All @@ -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,
))));
}
}
}
Expand Down Expand Up @@ -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)));
}
}

Expand All @@ -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),
)));
}
}
Expand Down Expand Up @@ -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,
))));
}
}

Expand Down Expand Up @@ -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::<HandleResult>();
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)
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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");
Expand All @@ -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]
Expand All @@ -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");
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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)))
);
}

Expand Down
6 changes: 5 additions & 1 deletion pgdog/src/net/messages/sync.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand All @@ -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());
}
}
Loading