Skip to content
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
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
2 changes: 2 additions & 0 deletions docs/runtime/sql.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -985,6 +985,8 @@ const sql = new SQL({
});
```

Each connection keeps up to 256 idle prepared statements cached; a statement still referenced by a query (in-flight or completed but not yet garbage-collected) is never evicted, so the count can temporarily exceed that. Past the limit, the least recently used idle statement is deallocated on the server with a `Close` message to make room. Without the cap, a long-lived connection running many distinct query strings (for example, an ORM interpolating identifiers) would accumulate named prepared statements in the server session, and their metadata in the client, until the connection closed.

When `prepare: false` is set:

Queries still use the "extended" protocol, but run as [unnamed prepared statements](https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY). An unnamed prepared statement lasts only until the next Parse statement specifying the unnamed statement as destination is issued.
Expand Down
121 changes: 110 additions & 11 deletions src/sql_jsc/postgres/PostgresSQLConnection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,12 @@ bun_core::define_scoped_log!(debug, Postgres, visible);

const MAX_PIPELINE_SIZE: usize = u16::MAX as usize; // about 64KB per connection

/// Per-connection cap on cached named prepared statements. Each cached
/// statement pins its metadata and row `Structure` on the client and a named
/// prepared statement in the server session until it is closed, so the cache
/// must be bounded and evictions must send `Close`.
const MAX_CACHED_PREPARED_STATEMENTS: usize = 256;

type PreparedStatementsMap = StringHashMap<*mut PostgresSQLStatement>;

pub mod js {
Expand Down Expand Up @@ -125,6 +131,9 @@ pub struct PostgresSQLConnection {
// so `vm_mut()`'s `&mut *as_ptr()` is sound.
pub vm: BackRef<VirtualMachine>,
pub statements: JsCell<PreparedStatementsMap>,
/// Monotonic clock for the statement cache's LRU order; bumped on every
/// cache lookup and stamped into `PostgresSQLStatement::last_used` on reuse.
statement_lru_clock: Cell<u64>,
pub prepared_statement_id: Cell<u64>,
pub pending_activity_count: AtomicU32,
// Self-wrapper back-ref (the JS object that owns this payload). Stored as a
Expand Down Expand Up @@ -1206,6 +1215,7 @@ pub(crate) fn call(global_object: &JSGlobalObject, callframe: &CallFrame) -> JsR
// canonical `VirtualMachine::as_mut()` accessor.
vm: BackRef::new_mut(vm),
statements: JsCell::new(PreparedStatementsMap::default()),
statement_lru_clock: Cell::new(0),
prepared_statement_id: Cell::new(0),
pending_activity_count: AtomicU32::new(0),
js_value: JsCell::new(crate::jsc::JsRef::empty()),
Expand Down Expand Up @@ -1634,6 +1644,102 @@ impl PostgresSQLConnection {
&& !flags.contains(ConnectionFlags::HAS_BACKPRESSURE) // dont make sense to buffer more if we have backpressure
&& (self.write_buffer.get().len() as usize) < MAX_PIPELINE_SIZE // buffer is too big need to flush before pipeline more
}

/// Statement-cache hit probe, keyed by the query signature. Stamps the LRU
/// clock on the hit so the least recently *reused* entry is the eviction
/// victim. Returns `None` on a miss; the caller then goes through
/// [`put_statement`](Self::put_statement).
pub(crate) fn lookup_statement(&self, name: &[u8]) -> Option<*mut PostgresSQLStatement> {
let stmt_ptr = self.statements.get().get(name).copied()?;
self.statement_lru_clock
.set(self.statement_lru_clock.get() + 1);
// Shared borrow through `ParentRef`: map values are live boxed
// statements (the map owns one intrusive ref on each) and `last_used`
// is a `Cell`.
ParentRef::from(NonNull::new(stmt_ptr).expect("map entries are non-null"))
.last_used
.set(self.statement_lru_clock.get());
Some(stmt_ptr)
}

/// Reserve a cache slot for a statement `lookup_statement` just missed on,
/// evicting LRU idle statements first so the cache stays under
/// [`MAX_CACHED_PREPARED_STATEMENTS`]. Returns the raw value-slot pointer;
/// the caller stores the new statement through it.
///
/// `JsCell::with_mut` scopes the `&mut PreparedStatementsMap` to the
/// `get_or_put` call (single-JS-thread; no re-entry into JS until after
/// the raw value-slot ptr is captured), so callers need no further `&mut`
/// to the map.
pub(crate) fn put_statement(
&self,
name: &[u8],
) -> Result<*mut *mut PostgresSQLStatement, bun_core::AllocError> {
self.evict_lru_statements();
self.statements.with_mut(|s| {
s.get_or_put(name)
.map(|e| core::ptr::from_mut::<*mut PostgresSQLStatement>(e.value_ptr))
})
}

/// Evict idle least-recently-used cached statements until the cache is
/// under [`MAX_CACHED_PREPARED_STATEMENTS`], writing a `Close('S', name)`
/// for each so the server drops its side of the named statement. The
/// paired CloseComplete is consumed by `on()` without touching the
/// request queue.
fn evict_lru_statements(&self) {
while self.statements.get().len() >= MAX_CACHED_PREPARED_STATEMENTS {
let mut victim: Option<NonNull<PostgresSQLStatement>> = None;
let mut oldest = u64::MAX;
for value in self.statements.get().values() {
let ptr = NonNull::new(*value).expect("map entries are non-null");
// Shared borrow; every field read below is `Cell`/`Copy`.
let stmt = ParentRef::from(ptr);
// A `PostgresSQLQuery` still holds this statement (pending,
// running, or not yet finalized): its server-side name must
// stay valid for the Bind that query may still send.
if !stmt.has_one_ref() {
continue;
}
if victim.is_none() || stmt.last_used.get() < oldest {
oldest = stmt.last_used.get();
victim = Some(ptr);
}
}
// Every cached statement is still referenced by a query; nothing
// can be released right now. The next insert tries again.
let Some(ptr) = victim else { return };
// `stmt` borrows the statement's own allocation, disjoint from the
// map's, so it stays live across the `with_mut` below.
let stmt = ParentRef::from(ptr);
// Only a statement the server acknowledged (status `Prepared`)
// owns a named server-side object to deallocate; an idle `Failed`
// one never completed a Parse (a rejected Parse is unmapped on
// ErrorResponse), so there is nothing to close for it.
if stmt.status == StatementStatus::Prepared
&& (protocol::Close {
p: protocol::PortalOrPreparedStatement::PreparedStatement(
&stmt.signature.prepared_statement_name,
),
})
.write(&mut self.writer())
.is_err()
{
// The Close could not be buffered (OOM): keep the cache
// entry so the server-side statement is not orphaned.
return;
}
Comment thread
robobun marked this conversation as resolved.
// The cache is keyed by `signature.name`, which the statement owns
// a copy of (same invariant the ErrorResponse arm relies on).
let removed = self
.statements
.with_mut(|m| m.remove(&stmt.signature.name[..]));
debug_assert!(removed.is_some(), "victim came from the map");
// SAFETY: `has_one_ref` above ⇒ the map owned the last ref;
// removing the entry transfers it to us to release here.
unsafe { PostgresSQLStatement::deref(ptr.as_ptr()) };
}
}
}

// `Writer.connection` is a
Expand Down Expand Up @@ -3006,18 +3112,11 @@ impl PostgresSQLConnection {
debug!("TODO PortalSuspended");
}
MessageType::CloseComplete => {
// The only Close this client sends is the statement cache's
// eviction (`evict_lru_statements`). It is not tied to any
// queued query, so consume the acknowledgement without
// touching the request queue.
reader.eat_message(&protocol::CLOSE_COMPLETE)?;
let request = self.current().ok_or(AnyPostgresError::ExpectedRequest)?;
if request.status.get() == QueryStatus::Fail {
return Ok(());
}
request.on_result(
b"CLOSECOMPLETE",
self.global(),
self.js_value.get().get(),
false,
);
self.update_ref();
}
MessageType::CopyInResponse => {
reader.skip_message()?;
Expand Down
25 changes: 9 additions & 16 deletions src/sql_jsc/postgres/PostgresSQLQuery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -615,14 +615,11 @@ impl PostgresSQLQuery {
.get()
.contains(ConnectionFlags::USE_UNNAMED_PREPARED_STATEMENTS)
{
// Zero-allocation hit probe: `get_or_put` below boxes the key
// Zero-allocation hit probe: `put_statement` below boxes the key
// bytes even when the entry already exists, and a hit (an
// already-prepared named statement) is the steady state.
let existing_stmt = connection
.statements
.get()
.get(&signature.name[..])
.copied();
// already-prepared named statement) is the steady state. A hit
// also stamps the entry's LRU clock.
let existing_stmt = connection.lookup_statement(&signature.name);
if let Some(stmt_ptr) = existing_stmt {
this.statement.set(Some(stmt_ptr));
// Route the `&mut` through the audited `statement_mut()`
Expand Down Expand Up @@ -688,15 +685,11 @@ impl PostgresSQLQuery {

break 'enqueue;
}
// `JsCell::with_mut` scopes the `&mut PreparedStatementsMap` to
// the `get_or_put` call (single-JS-thread; no re-entry into JS
// until after the raw value-slot ptr is captured). Extract the
// raw slot ptr while the borrow is live so the remainder of
// this block needs no further `&mut` to the map.
let entry_value_ptr = match connection.statements.with_mut(|s| {
s.get_or_put(&signature.name)
.map(|e| std::ptr::from_mut::<*mut PostgresSQLStatement>(e.value_ptr))
}) {
// `put_statement` enforces the statement-cache cap (evicting +
// closing LRU idle statements) and hands back the raw slot ptr,
// so the remainder of this block needs no further `&mut` to the
// map.
let entry_value_ptr = match connection.put_statement(&signature.name) {
Ok(v) => v,
Err(err) => {
drop(signature);
Expand Down
13 changes: 13 additions & 0 deletions src/sql_jsc/postgres/PostgresSQLStatement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@ pub struct PostgresSQLStatement {
pub error_response: Option<Error>,
pub needs_duplicate_check: bool,
pub fields_flags: DataCellFlags,
/// LRU stamp from the owning connection's statement clock, bumped each
/// time a later query reuses this cached statement; never-reused
/// statements keep 0 and are evicted first.
pub last_used: Cell<u64>,
}

impl Default for PostgresSQLStatement {
Expand All @@ -45,6 +49,7 @@ impl Default for PostgresSQLStatement {
error_response: None,
needs_duplicate_check: true,
fields_flags: DataCellFlags::default(),
last_used: Cell::new(0),
}
}
}
Expand Down Expand Up @@ -83,6 +88,14 @@ impl PostgresSQLStatement {
self.ref_count.set(n);
}

/// Whether the caller holds the only outstanding ref. The statement cache
/// only evicts statements it is the sole owner of (no `PostgresSQLQuery`
/// left that could still bind to the server-side statement name).
#[inline]
pub(crate) fn has_one_ref(&self) -> bool {
bun_ptr::CellRefCounted::ref_count(self).get() == 1
}

pub fn check_for_duplicate_fields(&mut self) {
if !self.needs_duplicate_check {
return;
Expand Down
Loading
Loading