From 2fa4fa3366bd776cdace83cbbb252323d11043d3 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Tue, 21 Jul 2026 14:17:31 -0600 Subject: [PATCH] Skip an allocation in our hot path Profiling showed this allocation was as much as .5% of our wall time. This function unfortunately does need to allocate, as one of the strategies involves randomizing the elements and another involves sorting them. Neither of those are possible without allocation. We also can't do any funky shit like re-using a thread local allocation, as `candidates` does need to be held past an await point. That leaves us with stack allocation as our only remaining option. This will always take up 256 bytes on the stack, and only heap allocate when we have more elements than that. The largest deployment that we know of will have 14 targets here, so this is potentially twice as large as it needs to be. But we can start by giving ourselves some wiggle room and we can always shrink this down later. The cost here is that the size of the future being returned by this function will be 222 bytes larger, so copying the future itself becomes a more expensive operation. In practice this should happen exactly once, as it gets moved to the heap when spawned, so it should be a net improvement over a random heap allocation. But profiling should be done to confirm. --- pgdog/src/backend/pool/lb/mod.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index dd6f7942e..755c82c65 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -348,8 +348,9 @@ impl LoadBalancer { async fn get_internal(&self, request: &Request) -> Result { use LoadBalancingStrategy::*; use ReadWriteSplit::*; + use smallvec::SmallVec; - let mut candidates: Vec<&Target> = self + let mut candidates: SmallVec<[&Target; 32]> = self .targets .iter() .filter(|target| !target.pool.config().resharding_only) // Don't let reads on resharding-only replicas. @@ -384,7 +385,7 @@ impl LoadBalancer { Random => candidates.shuffle(&mut rand::rng()), RoundRobin => { let first = self.round_robin.fetch_add(1, Ordering::Relaxed) % candidates.len(); - let mut reshuffled = vec![]; + let mut reshuffled = SmallVec::with_capacity(candidates.len()); reshuffled.extend_from_slice(&candidates[first..]); reshuffled.extend_from_slice(&candidates[..first]); candidates = reshuffled;