sqlx-postgres: Move rfq handling to the bg worker

Author: joeydewaal

sqlx-postgres/src/connection/describe.rs

@@ -414,11 +414,13 @@ WHERE rngtypid = $1
 
     /// Check whether EXPLAIN statements are supported by the current connection
     fn is_explain_available(&self) -> bool {
-        let parameter_statuses = &self.inner.stream.parameter_statuses;
-        let is_cockroachdb = parameter_statuses.contains_key("crdb_version");
-        let is_materialize = parameter_statuses.contains_key("mz_version");
-        let is_questdb = parameter_statuses.contains_key("questdb_version");
-        !is_cockroachdb && !is_materialize && !is_questdb
+        self.inner.shared.with_lock(|shared| {
+            let parameter_statuses = &shared.parameter_statuses;
+            let is_cockroachdb = parameter_statuses.contains_key("crdb_version");
+            let is_materialize = parameter_statuses.contains_key("mz_version");
+            let is_questdb = parameter_statuses.contains_key("questdb_version");
+            !is_cockroachdb && !is_materialize && !is_questdb
+        })
     }
 
     pub(crate) async fn get_nullable_for_columns(

sqlx-postgres/src/connection/establish.rs

@@ -1,11 +1,10 @@
 use crate::connection::{sasl, stream::PgStream};
 use crate::error::Error;
-use crate::message::{
-    Authentication, BackendKeyData, BackendMessageFormat, Password, ReadyForQuery, Startup,
-};
+use crate::message::{Authentication, BackendKeyData, BackendMessageFormat, Password, Startup};
 use crate::{PgConnectOptions, PgConnection};
 use futures_channel::mpsc::unbounded;
 
+use super::stream::parse_server_version;
 use super::worker::Worker;
 
 // https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.5.7.3
@@ -20,9 +19,9 @@ impl PgConnection {
 
         let (notif_tx, notif_rx) = unbounded();
 
-        let x = Worker::spawn(stream.into_inner(), notif_tx);
+        let (x, shared) = Worker::spawn(stream.into_inner(), notif_tx);
 
-        let mut conn = PgConnection::new(pg_stream, options, x, notif_rx);
+        let mut conn = PgConnection::new(pg_stream, options, x, notif_rx, shared);
 
         // To begin a session, a frontend opens a connection to the server
         // and sends a startup message.
@@ -65,7 +64,6 @@ impl PgConnection {
 
         let mut process_id = 0;
         let mut secret_key = 0;
-        let transaction_status;
 
         loop {
             let message = pipe.recv().await?;
@@ -121,9 +119,7 @@ impl PgConnection {
                 }
 
                 BackendMessageFormat::ReadyForQuery => {
-                    // start-up is completed. The frontend can now issue commands
-                    transaction_status = message.decode::<ReadyForQuery>()?.transaction_status;
-
+                    // Transaction status is updated in the bg worker.
                     break;
                 }
 
@@ -136,7 +132,13 @@ impl PgConnection {
             }
         }
 
-        conn.inner.transaction_status = transaction_status;
+        let server_version = conn
+            .inner
+            .shared
+            .remove_parameter_status("server_version")
+            .map(parse_server_version);
+
+        conn.inner.server_version_num = server_version.flatten();
         conn.inner.secret_key = secret_key;
         conn.inner.process_id = process_id;
 

sqlx-postgres/src/connection/executor.rs

@@ -309,8 +309,8 @@ impl PgConnection {
                     }
 
                     BackendMessageFormat::ReadyForQuery => {
-                        // processing of the query string is complete
-                        self.handle_ready_for_query(message)?;
+                        // Processing of the query string is complete, the transaction status is
+                        // updated in the bg worker.
                         break;
                     }
 

sqlx-postgres/src/connection/mod.rs

@@ -8,14 +8,13 @@ use futures_core::future::BoxFuture;
 use futures_util::FutureExt;
 use pipe::Pipe;
 use request::{IoRequest, MessageBuf};
+use worker::Shared;
 
 use crate::common::StatementCache;
 use crate::error::Error;
 use crate::ext::ustr::UStr;
 use crate::io::StatementId;
-use crate::message::{
-    Close, FrontendMessage, Notification, Query, ReadyForQuery, ReceivedMessage, TransactionStatus,
-};
+use crate::message::{Close, FrontendMessage, Notification, Query, TransactionStatus};
 use crate::statement::PgStatementMetadata;
 use crate::transaction::Transaction;
 use crate::types::Oid;
@@ -62,6 +61,8 @@ pub struct PgConnectionInner {
     #[allow(dead_code)]
     secret_key: u32,
 
+    pub(crate) server_version_num: Option<u32>,
+
     // sequence of statement IDs for use in preparing statements
     // in PostgreSQL, the statement is prepared to a user-supplied identifier
     next_statement_id: StatementId,
@@ -74,24 +75,17 @@ pub struct PgConnectionInner {
     cache_type_oid: HashMap<UStr, Oid>,
     cache_elem_type_to_array: HashMap<Oid, Oid>,
 
-    // current transaction status
-    transaction_status: TransactionStatus,
     pub(crate) transaction_depth: usize,
 
     log_settings: LogSettings,
+
+    shared: Shared,
 }
 
 impl PgConnection {
     /// the version number of the server in `libpq` format
     pub fn server_version_num(&self) -> Option<u32> {
-        self.inner.stream.server_version_num
-    }
-
-    #[inline(always)]
-    fn handle_ready_for_query(&mut self, message: ReceivedMessage) -> Result<(), Error> {
-        self.inner.transaction_status = message.decode::<ReadyForQuery>()?.transaction_status;
-
-        Ok(())
+        self.inner.server_version_num
     }
 
     /// Queue a simple query (not prepared) to execute the next time this connection is used.
@@ -103,7 +97,7 @@ impl PgConnection {
     }
 
     pub(crate) fn in_transaction(&self) -> bool {
-        match self.inner.transaction_status {
+        match self.inner.shared.get_transaction_status() {
             TransactionStatus::Transaction => true,
             TransactionStatus::Error | TransactionStatus::Idle => false,
         }
@@ -114,6 +108,7 @@ impl PgConnection {
         options: &PgConnectOptions,
         chan: UnboundedSender<IoRequest>,
         notifications: UnboundedReceiver<Notification>,
+        shared: Shared,
     ) -> Self {
         Self {
             inner: Box::new(PgConnectionInner {
@@ -129,7 +124,8 @@ impl PgConnection {
                 cache_elem_type_to_array: HashMap::new(),
                 transaction_depth: 0,
                 stream,
-                transaction_status: TransactionStatus::Idle,
+                server_version_num: None,
+                shared,
             }),
         }
     }

sqlx-postgres/src/connection/stream.rs

@@ -1,4 +1,3 @@
-use std::collections::BTreeMap;
 use std::ops::{ControlFlow, Deref, DerefMut};
 use std::str::FromStr;
 
@@ -8,8 +7,7 @@ use sqlx_core::bytes::Buf;
 use crate::connection::tls::MaybeUpgradeTls;
 use crate::error::Error;
 use crate::message::{
-    BackendMessage, BackendMessageFormat, EncodeMessage, FrontendMessage, Notice, ParameterStatus,
-    ReceivedMessage,
+    BackendMessage, BackendMessageFormat, EncodeMessage, FrontendMessage, Notice, ReceivedMessage,
 };
 use crate::net::{self, BufferedSocket, Socket};
 use crate::{PgConnectOptions, PgDatabaseError, PgSeverity};
@@ -27,10 +25,6 @@ pub struct PgStream {
     // A trait object is okay here as the buffering amortizes the overhead of both the dynamic
     // function call as well as the syscall.
     inner: BufferedSocket<Box<dyn Socket>>,
-
-    pub(crate) parameter_statuses: BTreeMap<String, String>,
-
-    pub(crate) server_version_num: Option<u32>,
 }
 
 impl PgStream {
@@ -48,8 +42,6 @@ impl PgStream {
 
         Ok(Self {
             inner: BufferedSocket::new(socket),
-            parameter_statuses: BTreeMap::default(),
-            server_version_num: None,
         })
     }
 
@@ -126,25 +118,24 @@ impl PgStream {
                     return Err(message.decode::<PgDatabaseError>()?.into());
                 }
 
-                BackendMessageFormat::ParameterStatus => {
-                    // informs the frontend about the current (initial)
-                    // setting of backend parameters
+                // BackendMessageFormat::ParameterStatus => {
+                //     // informs the frontend about the current (initial)
+                //     // setting of backend parameters
 
-                    let ParameterStatus { name, value } = message.decode()?;
-                    // TODO: handle `client_encoding`, `DateStyle` change
-
-                    match name.as_str() {
-                        "server_version" => {
-                            self.server_version_num = parse_server_version(&value);
-                        }
-                        _ => {
-                            self.parameter_statuses.insert(name, value);
-                        }
-                    }
+                //     let ParameterStatus { name, value } = message.decode()?;
+                //     // TODO: handle `client_encoding`, `DateStyle` change
 
-                    continue;
-                }
+                //     match name.as_str() {
+                //         "server_version" => {
+                //             self.server_version_num = parse_server_version(&value);
+                //         }
+                //         _ => {
+                //             self.parameter_statuses.insert(name, value);
+                //         }
+                //     }
 
+                //     continue;
+                // }
                 BackendMessageFormat::NoticeResponse => {
                     // do we need this to be more configurable?
                     // if you are reading this comment and think so, open an issue
@@ -205,7 +196,7 @@ impl DerefMut for PgStream {
 
 // reference:
 // https://github.com/postgres/postgres/blob/6feebcb6b44631c3dc435e971bd80c2dd218a5ab/src/interfaces/libpq/fe-exec.c#L1030-L1065
-fn parse_server_version(s: &str) -> Option<u32> {
+pub fn parse_server_version(s: String) -> Option<u32> {
     let mut parts = Vec::<u32>::with_capacity(3);
 
     let mut from = 0;

sqlx-postgres/src/connection/worker.rs

@@ -1,13 +1,15 @@
 use std::{
-    collections::VecDeque,
+    collections::{BTreeMap, VecDeque},
     future::Future,
     ops::ControlFlow,
     pin::Pin,
+    sync::{Arc, Mutex, MutexGuard},
     task::{ready, Context, Poll},
 };
 
 use crate::message::{
-    BackendMessageFormat, FrontendMessage, Notification, ReceivedMessage, Terminate,
+    BackendMessageFormat, FrontendMessage, Notification, ParameterStatus, ReadyForQuery,
+    ReceivedMessage, Terminate, TransactionStatus,
 };
 use futures_channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
 use futures_util::{SinkExt, StreamExt};
@@ -38,14 +40,16 @@ pub struct Worker {
     back_log: VecDeque<UnboundedSender<ReceivedMessage>>,
     socket: BufferedSocket<Box<dyn Socket>>,
     notif_chan: UnboundedSender<Notification>,
+    shared: Shared,
 }
 
 impl Worker {
     pub fn spawn(
         socket: BufferedSocket<Box<dyn Socket>>,
         notif_chan: UnboundedSender<Notification>,
-    ) -> UnboundedSender<IoRequest> {
+    ) -> (UnboundedSender<IoRequest>, Shared) {
         let (tx, rx) = unbounded();
+        let shared = Shared::new();
 
         let worker = Worker {
             state: WorkerState::Open,
@@ -54,10 +58,11 @@ impl Worker {
             back_log: VecDeque::new(),
             socket,
             notif_chan,
+            shared: shared.clone(),
         };
 
         spawn(worker);
-        tx
+        (tx, shared)
     }
 
     // Tries to receive the next message from the channel. Also handles termination if needed.
@@ -128,6 +133,9 @@ impl Worker {
         while let Poll::Ready(response) = self.poll_next_message(cx)? {
             match response.format {
                 BackendMessageFormat::ReadyForQuery => {
+                    let rfq: ReadyForQuery = response.clone().decode()?;
+                    self.shared.set_transaction_status(rfq.transaction_status);
+
                     self.send_back(response)?;
                     // Remove from the backlog so we dont send more responses back.
                     let _ = self.back_log.pop_front();
@@ -145,6 +153,9 @@ impl Worker {
                 }
                 BackendMessageFormat::ParameterStatus => {
                     // Asynchronous response - todo
+                    //
+                    let ParameterStatus { name, value } = response.decode()?;
+                    self.shared.insert_parameter_status(name, value);
                 }
                 BackendMessageFormat::NoticeResponse => {
                     // Asynchronous response - todo
@@ -226,3 +237,47 @@ impl Future for Worker {
         self.poll_shutdown(cx)
     }
 }
+
+#[derive(Clone)]
+pub struct Shared(Arc<Mutex<SharedInner>>);
+
+pub struct SharedInner {
+    pub parameter_statuses: BTreeMap<String, String>,
+    pub transaction_status: TransactionStatus,
+}
+
+impl Shared {
+    pub fn new() -> Shared {
+        Shared(Arc::new(Mutex::new(SharedInner {
+            parameter_statuses: BTreeMap::new(),
+            transaction_status: TransactionStatus::Idle,
+        })))
+    }
+
+    fn lock(&self) -> MutexGuard<'_, SharedInner> {
+        self.0
+            .lock()
+            .expect("BUG: failed to get lock on shared state in worker")
+    }
+
+    pub fn get_transaction_status(&self) -> TransactionStatus {
+        self.lock().transaction_status
+    }
+
+    fn set_transaction_status(&self, status: TransactionStatus) {
+        self.lock().transaction_status = status
+    }
+
+    fn insert_parameter_status(&self, name: String, value: String) {
+        self.lock().parameter_statuses.insert(name, value);
+    }
+
+    pub fn remove_parameter_status(&self, name: &str) -> Option<String> {
+        self.lock().parameter_statuses.remove(name)
+    }
+
+    pub fn with_lock<T>(&self, f: impl Fn(&mut SharedInner) -> T) -> T {
+        let mut lock = self.lock();
+        f(&mut lock)
+    }
+}

sqlx-postgres/src/message/ready_for_query.rs

@@ -3,7 +3,7 @@ use sqlx_core::bytes::Bytes;
 use crate::error::Error;
 use crate::message::{BackendMessage, BackendMessageFormat};
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 #[repr(u8)]
 pub enum TransactionStatus {
     /// Not in a transaction block.