sqlx-postgres: Startup with background worker

joeydewaal · joeydewaal · commit cb193199f99a · 2025-07-02T15:10:43.000+02:00
diff --git a/sqlx-postgres/src/connection/establish.rs b/sqlx-postgres/src/connection/establish.rs
@@ -1,23 +1,28 @@
-use crate::HashMap;
-
-use crate::common::StatementCache;
 use crate::connection::{sasl, stream::PgStream};
 use crate::error::Error;
-use crate::io::StatementId;
 use crate::message::{
     Authentication, BackendKeyData, BackendMessageFormat, Password, ReadyForQuery, Startup,
 };
 use crate::{PgConnectOptions, PgConnection};
+use futures_channel::mpsc::unbounded;
 
-use super::PgConnectionInner;
+use super::worker::Worker;
 
 // https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.3
 // https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.11
 
 impl PgConnection {
     pub(crate) async fn establish(options: &PgConnectOptions) -> Result<Self, Error> {
         // Upgrade to TLS if we were asked to and the server supports it
-        let mut stream = PgStream::connect(options).await?;
+        let pg_stream = PgStream::connect(options).await?;
+
+        let stream = PgStream::connect(options).await?;
+
+        let (notif_tx, notif_rx) = unbounded();
+
+        let x = Worker::spawn(stream.into_inner(), notif_tx);
+
+        let mut conn = PgConnection::new(pg_stream, options, x, notif_rx);
 
         // To begin a session, a frontend opens a connection to the server
         // and sends a startup message.
@@ -45,14 +50,14 @@ impl PgConnection {
             params.push(("options", options));
         }
 
-        stream.write(Startup {
-            username: Some(&options.username),
-            database: options.database.as_deref(),
-            params: &params,
+        let mut pipe = conn.pipe(|buf| {
+            buf.write(Startup {
+                username: Some(&options.username),
+                database: options.database.as_deref(),
+                params: &params,
+            })
         })?;
 
-        stream.flush().await?;
-
         // The server then uses this information and the contents of
         // its configuration files (such as pg_hba.conf) to determine whether the connection is
         // provisionally acceptable, and what additional
@@ -63,7 +68,7 @@ impl PgConnection {
         let transaction_status;
 
         loop {
-            let message = stream.recv().await?;
+            let message = pipe.recv().await?;
             match message.format {
                 BackendMessageFormat::Authentication => match message.decode()? {
                     Authentication::Ok => {
@@ -75,11 +80,9 @@ impl PgConnection {
                         // The frontend must now send a [PasswordMessage] containing the
                         // password in clear-text form.
 
-                        stream
-                            .send(Password::Cleartext(
-                                options.password.as_deref().unwrap_or_default(),
-                            ))
-                            .await?;
+                        conn.pipe_and_forget(Password::Cleartext(
+                            options.password.as_deref().unwrap_or_default(),
+                        ))?;
                     }
 
                     Authentication::Md5Password(body) => {
@@ -88,17 +91,15 @@ impl PgConnection {
                         // using the 4-byte random salt specified in the
                         // [AuthenticationMD5Password] message.
 
-                        stream
-                            .send(Password::Md5 {
-                                username: &options.username,
-                                password: options.password.as_deref().unwrap_or_default(),
-                                salt: body.salt,
-                            })
-                            .await?;
+                        conn.pipe_and_forget(Password::Md5 {
+                            username: &options.username,
+                            password: options.password.as_deref().unwrap_or_default(),
+                            salt: body.salt,
+                        })?;
                     }
 
                     Authentication::Sasl(body) => {
-                        sasl::authenticate(&mut stream, options, body).await?;
+                        sasl::authenticate(&conn, &mut pipe, options, body).await?;
                     }
 
                     method => {
@@ -135,22 +136,10 @@ impl PgConnection {
             }
         }
 
-        Ok(PgConnection {
-            inner: Box::new(PgConnectionInner {
-                stream,
-                process_id,
-                secret_key,
-                transaction_status,
-                transaction_depth: 0,
-                pending_ready_for_query_count: 0,
-                next_statement_id: StatementId::NAMED_START,
-                cache_statement: StatementCache::new(options.statement_cache_capacity),
-                cache_type_oid: HashMap::new(),
-                cache_type_info: HashMap::new(),
-                cache_elem_type_to_array: HashMap::new(),
-                cache_table_to_column_names: HashMap::new(),
-                log_settings: options.log_settings.clone(),
-            }),
-        })
+        conn.inner.transaction_status = transaction_status;
+        conn.inner.secret_key = secret_key;
+        conn.inner.process_id = process_id;
+
+        Ok(conn)
     }
 }
diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs
@@ -4,16 +4,19 @@ use std::fmt::{self, Debug, Formatter};
 use std::sync::Arc;
 
 use crate::HashMap;
+use futures_channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
 use futures_core::future::BoxFuture;
 use futures_util::FutureExt;
+use pipe::Pipe;
+use request::{IoRequest, MessageBuf};
 
 use crate::common::StatementCache;
 use crate::error::Error;
 use crate::ext::ustr::UStr;
 use crate::io::StatementId;
 use crate::message::{
-    BackendMessageFormat, Close, Query, ReadyForQuery, ReceivedMessage, Terminate,
-    TransactionStatus,
+    BackendMessageFormat, Close, FrontendMessage, Notification, Query, ReadyForQuery,
+    ReceivedMessage, Terminate, TransactionStatus,
 };
 use crate::statement::PgStatementMetadata;
 use crate::transaction::Transaction;
@@ -47,6 +50,10 @@ pub struct PgConnectionInner {
     // wrapped in a buffered stream
     pub(crate) stream: PgStream,
 
+    chan: UnboundedSender<IoRequest>,
+
+    notifications: UnboundedReceiver<Notification>,
+
     // process id of this backend
     // used to send cancel requests
     #[allow(dead_code)]
@@ -148,6 +155,84 @@ impl PgConnection {
             TransactionStatus::Error | TransactionStatus::Idle => false,
         }
     }
+
+    fn new(
+        stream: PgStream,
+        options: &PgConnectOptions,
+        chan: UnboundedSender<IoRequest>,
+        notifications: UnboundedReceiver<Notification>,
+    ) -> Self {
+        Self {
+            inner: Box::new(PgConnectionInner {
+                chan,
+                notifications,
+                log_settings: options.log_settings.clone(),
+                process_id: 0,
+                secret_key: 0,
+                next_statement_id: StatementId::NAMED_START,
+                cache_statement: StatementCache::new(options.statement_cache_capacity),
+                cache_type_info: HashMap::new(),
+                cache_type_oid: HashMap::new(),
+                cache_elem_type_to_array: HashMap::new(),
+                cache_table_to_column_names: HashMap::new(),
+                transaction_depth: 0,
+                stream,
+                pending_ready_for_query_count: 0,
+                transaction_status: TransactionStatus::Idle,
+            }),
+        }
+    }
+
+    fn create_request<F>(&self, callback: F) -> sqlx_core::Result<IoRequest>
+    where
+        F: FnOnce(&mut MessageBuf) -> sqlx_core::Result<()>,
+    {
+        let mut buffer = MessageBuf::new();
+        (callback)(&mut buffer)?;
+        Ok(buffer.finish())
+    }
+
+    fn send_request(&self, request: IoRequest) -> sqlx_core::Result<()> {
+        self.inner
+            .chan
+            .unbounded_send(request)
+            .map_err(|_| sqlx_core::Error::WorkerCrashed)
+    }
+
+    pub(crate) fn pipe<F>(&self, callback: F) -> sqlx_core::Result<Pipe>
+    where
+        F: FnOnce(&mut MessageBuf) -> sqlx_core::Result<()>,
+    {
+        let mut req = self.create_request(callback)?;
+        let (tx, rx) = unbounded();
+        req.chan = Some(tx);
+
+        self.send_request(req)?;
+        Ok(Pipe::new(rx))
+    }
+
+    pub(crate) fn pipe_and_forget<T>(&self, value: T) -> sqlx_core::Result<()>
+    where
+        T: FrontendMessage,
+    {
+        let req = self.create_request(|buf| buf.write_msg(value))?;
+        self.send_request(req)
+    }
+
+    pub(crate) async fn start_pipe_async<F, R>(&self, callback: F) -> sqlx_core::Result<(R, Pipe)>
+    where
+        F: AsyncFnOnce(&mut MessageBuf) -> sqlx_core::Result<R>,
+    {
+        let mut buffer = MessageBuf::new();
+        let result = (callback)(&mut buffer).await?;
+        let mut req = buffer.finish();
+        let (tx, rx) = unbounded();
+        req.chan = Some(tx);
+
+        self.send_request(req)?;
+
+        Ok((result, Pipe::new(rx)))
+    }
 }
 
 impl Debug for PgConnection {
diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs
@@ -1,4 +1,3 @@
-use crate::connection::stream::PgStream;
 use crate::error::Error;
 use crate::message::{Authentication, AuthenticationSasl, SaslInitialResponse, SaslResponse};
 use crate::PgConnectOptions;
@@ -9,14 +8,18 @@ use stringprep::saslprep;
 
 use base64::prelude::{Engine as _, BASE64_STANDARD};
 
+use super::pipe::Pipe;
+use super::PgConnection;
+
 const GS2_HEADER: &str = "n,,";
 const CHANNEL_ATTR: &str = "c";
 const USERNAME_ATTR: &str = "n";
 const CLIENT_PROOF_ATTR: &str = "p";
 const NONCE_ATTR: &str = "r";
 
 pub(crate) async fn authenticate(
-    stream: &mut PgStream,
+    conn: &PgConnection,
+    pipe: &mut Pipe,
     options: &PgConnectOptions,
     data: AuthenticationSasl,
 ) -> Result<(), Error> {
@@ -67,14 +70,12 @@ pub(crate) async fn authenticate(
 
     let client_first_message = format!("{GS2_HEADER}{client_first_message_bare}");
 
-    stream
-        .send(SaslInitialResponse {
-            response: &client_first_message,
-            plus: false,
-        })
-        .await?;
+    conn.pipe_and_forget(SaslInitialResponse {
+        response: &client_first_message,
+        plus: false,
+    })?;
 
-    let cont = match stream.recv_expect().await? {
+    let cont = match pipe.recv_expect().await? {
         Authentication::SaslContinue(data) => data,
 
         auth => {
@@ -143,9 +144,9 @@ pub(crate) async fn authenticate(
     let mut client_final_message = format!("{client_final_message_wo_proof},{CLIENT_PROOF_ATTR}=");
     BASE64_STANDARD.encode_string(client_proof, &mut client_final_message);
 
-    stream.send(SaslResponse(&client_final_message)).await?;
+    conn.pipe_and_forget(SaslResponse(&client_final_message))?;
 
-    let data = match stream.recv_expect().await? {
+    let data = match pipe.recv_expect().await? {
         Authentication::SaslFinal(data) => data,
 
         auth => {
diff --git a/sqlx-postgres/src/connection/stream.rs b/sqlx-postgres/src/connection/stream.rs
@@ -41,6 +41,10 @@ pub struct PgStream {
 }
 
 impl PgStream {
+    pub fn into_inner(self) -> BufferedSocket<Box<dyn Socket>> {
+        self.inner
+    }
+
     pub(super) async fn connect(options: &PgConnectOptions) -> Result<Self, Error> {
         let socket_result = match options.fetch_socket() {
             Some(ref path) => net::connect_uds(path, MaybeUpgradeTls(options)).await?,
diff --git a/sqlx-postgres/src/connection/worker.rs b/sqlx-postgres/src/connection/worker.rs
@@ -7,7 +7,7 @@ use std::{
 };
 
 use crate::message::{
-    BackendMessageFormat, FrontendMessage, Notification, ReadyForQuery, ReceivedMessage, Terminate,
+    BackendMessageFormat, FrontendMessage, Notification, ReceivedMessage, Terminate,
 };
 use futures_channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
 use futures_util::{SinkExt, StreamExt};
@@ -129,7 +129,6 @@ impl Worker {
         while let Poll::Ready(response) = self.poll_next_message(cx)? {
             match response.format {
                 BackendMessageFormat::ReadyForQuery => {
-                    let rfq: ReadyForQuery = response.clone().decode()?;
                     self.send_back(response)?;
                     // Remove from the backlog so we dont send more responses back.
                     let _ = self.back_log.pop_front();
