From 24357a093c776913225714dccb9c1103458b3064 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Thu, 23 Jul 2026 17:43:57 -0700 Subject: [PATCH] clippy + docs --- pgdog/src/frontend/client/mod.rs | 47 ++++++++++++++++++++++---------- pgdog/src/frontend/comms.rs | 39 ++++++++++++++++++++++---- 2 files changed, 67 insertions(+), 19 deletions(-) diff --git a/pgdog/src/frontend/client/mod.rs b/pgdog/src/frontend/client/mod.rs index 56f9004ed..cd18cc618 100644 --- a/pgdog/src/frontend/client/mod.rs +++ b/pgdog/src/frontend/client/mod.rs @@ -531,43 +531,46 @@ impl Client { /// Handle client messages. async fn client_messages(&mut self, query_engine: &mut QueryEngine) -> Result<(), Error> { - // Check maintenance mode. - if !self.in_transaction() - && !self.admin - && let Some(waiter) = maintenance_mode::waiter(&self.database) - { - let state = query_engine.get_state(); - query_engine.set_state(State::Waiting); - waiter.await; - query_engine.set_state(state); - } - // If client sent multiple requests, split them up and execute individually. let spliced = self.client_request.spliced()?; if spliced.is_empty() { + // Execute the original request as-is. let mut context = QueryEngineContext::new(self); query_engine.handle(&mut context).await?; self.transaction = context.transaction(); } else { + // Execute the spliced requests one at a time. let total = spliced.len(); let mut reqs = spliced.into_iter().enumerate(); + + // Pipeliend requests are executed inside the same Postgres transaction. self.transaction.get_or_insert(TransactionType::Implicit); + while let Some((num, mut req)) = reqs.next() { debug!("processing spliced request {}/{}", num + 1, total); + let mut context = QueryEngineContext::new(self).spliced(&mut req, reqs.len()); query_engine.handle(&mut context).await?; + + // Save transaction state after each request. self.transaction = context.transaction(); // If pipeline is aborted due to error, skip to Sync to complete the pipeline. - // Postgres ignores all commands after an error until it receives Sync. + // Postgres ignores all commands after an error until it receives a Sync. if query_engine.out_of_sync() && !req.is_sync_only() { debug!("pipeline aborted, skipping to Sync"); + for (_, mut next_req) in reqs.by_ref() { + // The splicer (called above) leaves + // the Sync in its own request. This is by design + // to handle this exact scenario. if next_req.is_sync_only() { debug!("processing Sync to complete aborted pipeline"); + let mut ctx = QueryEngineContext::new(self).spliced(&mut next_req, 0); query_engine.handle(&mut ctx).await?; self.transaction = ctx.transaction(); + break; } } @@ -582,13 +585,29 @@ impl Client { Ok(()) } - /// Buffer extended protocol messages until client requests a sync. + /// Buffer protocol messages until client requests a sync or a flush. /// /// This ensures we don't check out a connection from the pool until the client - /// sent a complete request. + /// sends a complete request. + /// async fn buffer(&mut self, state: State) -> Result { self.client_request.clear(); + // Check maintenance mode. + // + // This will pause the client until the maintenance mode is off + // without checking any timeouts or reading data off of the socket + // so if the client disconnects while we wait, we won't send + // its pending query to the database. + if !self.in_transaction() + && !self.admin + && let Some(waiter) = maintenance_mode::waiter(&self.database) + { + self.comms.set_state(State::Waiting); + waiter.await; + self.comms.set_state(state); + } + // Only start timer once we receive the first message. let mut timer = None; diff --git a/pgdog/src/frontend/comms.rs b/pgdog/src/frontend/comms.rs index 4c8e081c0..8eb1bac92 100644 --- a/pgdog/src/frontend/comms.rs +++ b/pgdog/src/frontend/comms.rs @@ -15,6 +15,7 @@ use tokio_util::task::TaskTracker; use crate::net::Parameters; use crate::net::messages::{BackendKeyData, FrontendPid}; +use crate::state::State; use super::{ConnectedClient, Stats}; @@ -155,23 +156,51 @@ impl Deref for ClientComms { } impl ClientComms { - pub fn disconnect(&self) { + /// Client has disconnected, remove it from the global state. + pub(crate) fn disconnect(&self) { self.comms.disconnect(self.id); } - pub fn update_stats(&self, stats: Stats) { + /// Update global client stats with the latest stats the caller knows about. + /// The caller must maintain up-to-date stats, or this will overwrite with bad + /// info. + pub(crate) fn update_stats(&self, stats: Stats) { self.comms.update_stats(self.id, stats); } - pub fn new(id: FrontendPid) -> Self { + /// Update the client state quickly. + /// + /// Use only when you're not updating other stats. If updating other stats, + /// use [`Self::update_stats`] instead. + /// + pub(crate) fn set_state(&self, state: State) { + if let Some(mut client) = self.comms.global.clients.get_mut(&self.id) { + client.stats.state = state; + } + } + + /// Create new comms reference for a new client. + /// + /// Makes it easy to update globals without constantly reaching for the key. + pub(crate) fn new(id: FrontendPid) -> Self { Self { id, comms: comms() } } - pub fn connect(&self, key: BackendKeyData, addr: SocketAddr, params: &Parameters) { + /// Mark the client as connected to the proxy + /// and save some of its identifying information. + /// + /// # Parameters + /// + /// - `key`: The query cancellation key we gave to the client. + /// - `addr`: The client's IP address. + /// - `params`: The client's startup parameters. + /// + pub(crate) fn connect(&self, key: BackendKeyData, addr: SocketAddr, params: &Parameters) { self.comms.connect(key, addr, params) } - pub fn update_params(&self, params: &Parameters) { + /// Update the params the client has in its session state. + pub(crate) fn update_params(&self, params: &Parameters) { self.comms.update_params(self.id, params.clone()); } }