Skip to content

Commit

Permalink
proxy: lazily parse startup pg params (#7905)
Browse files Browse the repository at this point in the history
## Problem

proxy params being a `HashMap<String,String>` when it contains just
```
application_name: psql
database: neondb
user: neondb_owner
```
is quite wasteful allocation wise.

## Summary of changes

Keep the params in the wire protocol form, eg:
```
application_name\0psql\0database\0neondb\0user\0neondb_owner\0
```

Using a linear search for the map is fast enough at small sizes, which
is the normal case.
  • Loading branch information
conradludgate authored May 30, 2024
1 parent fddd11d commit 9a081c2
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 36 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions libs/pq_proto/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ license.workspace = true
[dependencies]
bytes.workspace = true
byteorder.workspace = true
itertools.workspace = true
pin-project-lite.workspace = true
postgres-protocol.workspace = true
rand.workspace = true
Expand Down
76 changes: 42 additions & 34 deletions libs/pq_proto/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,9 @@ pub mod framed;

use byteorder::{BigEndian, ReadBytesExt};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use itertools::Itertools;
use serde::{Deserialize, Serialize};
use std::{borrow::Cow, collections::HashMap, fmt, io, str};
use std::{borrow::Cow, fmt, io, str};

// re-export for use in utils pageserver_feedback.rs
pub use postgres_protocol::PG_EPOCH;
Expand Down Expand Up @@ -51,19 +52,36 @@ pub enum FeStartupPacket {
}

#[derive(Debug, Clone, Default)]
pub struct StartupMessageParams {
params: HashMap<String, String>,
pub struct StartupMessageParamsBuilder {
params: BytesMut,
}

impl StartupMessageParams {
impl StartupMessageParamsBuilder {
/// Set parameter's value by its name.
/// name and value must not contain a \0 byte
pub fn insert(&mut self, name: &str, value: &str) {
self.params.insert(name.to_owned(), value.to_owned());
self.params.put(name.as_bytes());
self.params.put(&b"\0"[..]);
self.params.put(value.as_bytes());
self.params.put(&b"\0"[..]);
}

pub fn freeze(self) -> StartupMessageParams {
StartupMessageParams {
params: self.params.freeze(),
}
}
}

#[derive(Debug, Clone, Default)]
pub struct StartupMessageParams {
params: Bytes,
}

impl StartupMessageParams {
/// Get parameter's value by its name.
pub fn get(&self, name: &str) -> Option<&str> {
self.params.get(name).map(|s| s.as_str())
self.iter().find_map(|(k, v)| (k == name).then_some(v))
}

/// Split command-line options according to PostgreSQL's logic,
Expand Down Expand Up @@ -117,15 +135,19 @@ impl StartupMessageParams {

/// Iterate through key-value pairs in an arbitrary order.
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
self.params.iter().map(|(k, v)| (k.as_str(), v.as_str()))
let params =
std::str::from_utf8(&self.params).expect("should be validated as utf8 already");
params.split_terminator('\0').tuples()
}

// This function is mostly useful in tests.
#[doc(hidden)]
pub fn new<'a, const N: usize>(pairs: [(&'a str, &'a str); N]) -> Self {
Self {
params: pairs.map(|(k, v)| (k.to_owned(), v.to_owned())).into(),
let mut b = StartupMessageParamsBuilder::default();
for (k, v) in pairs {
b.insert(k, v)
}
b.freeze()
}
}

Expand Down Expand Up @@ -350,35 +372,21 @@ impl FeStartupPacket {
(major_version, minor_version) => {
// StartupMessage

// Parse pairs of null-terminated strings (key, value).
// See `postgres: ProcessStartupPacket, build_startup_packet`.
let mut tokens = str::from_utf8(&msg)
.map_err(|_e| {
ProtocolError::BadMessage("StartupMessage params: invalid utf-8".to_owned())
})?
.strip_suffix('\0') // drop packet's own null
.ok_or_else(|| {
ProtocolError::Protocol(
"StartupMessage params: missing null terminator".to_string(),
)
})?
.split_terminator('\0');

let mut params = HashMap::new();
while let Some(name) = tokens.next() {
let value = tokens.next().ok_or_else(|| {
ProtocolError::Protocol(
"StartupMessage params: key without value".to_string(),
)
})?;

params.insert(name.to_owned(), value.to_owned());
}
let s = str::from_utf8(&msg).map_err(|_e| {
ProtocolError::BadMessage("StartupMessage params: invalid utf-8".to_owned())
})?;
let s = s.strip_suffix('\0').ok_or_else(|| {
ProtocolError::Protocol(
"StartupMessage params: missing null terminator".to_string(),
)
})?;

FeStartupPacket::StartupMessage {
major_version,
minor_version,
params: StartupMessageParams { params },
params: StartupMessageParams {
params: msg.slice_ref(s.as_bytes()),
},
}
}
};
Expand Down
4 changes: 2 additions & 2 deletions proxy/src/serverless/sql_over_http.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use hyper1::http::HeaderValue;
use hyper1::Response;
use hyper1::StatusCode;
use hyper1::{HeaderMap, Request};
use pq_proto::StartupMessageParams;
use pq_proto::StartupMessageParamsBuilder;
use serde_json::json;
use serde_json::Value;
use tokio::time;
Expand Down Expand Up @@ -193,7 +193,7 @@ fn get_conn_info(

let mut options = Option::None;

let mut params = StartupMessageParams::default();
let mut params = StartupMessageParamsBuilder::default();
params.insert("user", &username);
params.insert("database", &dbname);
for (key, value) in pairs {
Expand Down

1 comment on commit 9a081c2

@github-actions
Copy link

@github-actions github-actions bot commented on 9a081c2 May 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

3262 tests run: 3117 passed, 0 failed, 145 skipped (full report)


Flaky tests (6)

Postgres 15

  • test_timeline_deletion_with_files_stuck_in_upload_queue: debug
  • test_vm_bit_clear_on_heap_lock: release, debug

Postgres 14

  • test_pageserver_restarts_under_worload: release
  • test_timeline_deletion_with_files_stuck_in_upload_queue: debug
  • test_replication_start: release

Code coverage* (full report)

  • functions: 31.4% (6485 of 20666 functions)
  • lines: 48.4% (50201 of 103758 lines)

* collected from Rust tests only


The comment gets automatically updated with the latest test results
9a081c2 at 2024-05-30T12:54:49.205Z :recycle:

Please sign in to comment.