Remove sync postgres_backend_async, tidy up its split usage.

- Use postgres_backend_async throughout safekeeper.
- Use framed.rs in postgres_backend_async, similar to tokio_util::codec::Framed
  but with slightly more efficient split. Now IO functions are also cancellation
  safe. Also, there is no allocation in each message read anymore, and data in
  read messages still directly point to input buffer without copies.
- In both safekeeper COPY streams, do read-write from the same thread/task with
  select! for easier error handling.
- Tidy up finishing CopyBoth streams in safekeeper sending and receiving WAL
  -- join split parts back catching errors from them before returning.

Initially I hoped to do that read-write without split at all, through polling
IO:
https://github.com/neondatabase/neon/pull/3522
However that turned out to be more complicated than I initially expected
due to 1) borrow checking and 2) anon Future types. 1) required Rc<Refcell<...>>
which is Send construct just to satisfy the checker; 2) can be workaround with
transmute. But this is so messy that I decided to leave split.

proxy stream.rs is adapted with minimal changes. It also benefits from framed.rs
improvements described above.

proxy/Cargo.toml

@@ -31,6 +31,7 @@ once_cell.workspace = true
 opentelemetry.workspace = true
 parking_lot.workspace = true
 pin-project-lite.workspace = true
+postgres_backend.workspace = true
 pq_proto.workspace = true
 prometheus.workspace = true
 rand.workspace = true

proxy/src/console/mgmt.rs

@@ -4,13 +4,11 @@ use crate::{
 };
 use anyhow::Context;
 use once_cell::sync::Lazy;
+use postgres_backend::{self, AuthType, PostgresBackend, QueryError};
 use pq_proto::{BeMessage, SINGLE_COL_ROWDESC};
-use std::{net::TcpStream, thread};
+use std::future;
+use tokio::net::{TcpListener, TcpStream};
 use tracing::{error, info, info_span};
-use utils::{
-    postgres_backend::{self, AuthType, PostgresBackend},
-    postgres_backend_async::QueryError,
-};
 
 static CPLANE_WAITERS: Lazy<Waiters<ComputeReady>> = Lazy::new(Default::default);
 
@@ -33,7 +31,7 @@ pub fn notify(psql_session_id: &str, msg: ComputeReady) -> Result<(), waiters::N
 
 /// Console management API listener task.
 /// It spawns console response handlers needed for the link auth.
-pub async fn task_main(listener: tokio::net::TcpListener) -> anyhow::Result<()> {
+pub async fn task_main(listener: TcpListener) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("mgmt has shut down");
     }
@@ -42,18 +40,12 @@ pub async fn task_main(listener: tokio::net::TcpListener) -> anyhow::Result<()>
         let (socket, peer_addr) = listener.accept().await?;
         info!("accepted connection from {peer_addr}");
 
-        let socket = socket.into_std()?;
         socket
             .set_nodelay(true)
             .context("failed to set client socket option")?;
-        socket
-            .set_nonblocking(false)
-            .context("failed to set client socket option")?;
 
-        // TODO: replace with async tasks.
-        thread::spawn(move || {
-            let tid = std::thread::current().id();
-            let span = info_span!("mgmt", thread = format_args!("{tid:?}"));
+        tokio::task::spawn(async move {
+            let span = info_span!("mgmt", peer = %peer_addr);
             let _enter = span.enter();
 
             info!("started a new console management API thread");
@@ -61,16 +53,16 @@ pub async fn task_main(listener: tokio::net::TcpListener) -> anyhow::Result<()>
                 info!("console management API thread is about to finish");
             }
 
-            if let Err(e) = handle_connection(socket) {
+            if let Err(e) = handle_connection(socket).await {
                 error!("thread failed with an error: {e}");
             }
         });
     }
 }
 
-fn handle_connection(socket: TcpStream) -> Result<(), QueryError> {
-    let pgbackend = PostgresBackend::new(socket, AuthType::Trust, None, true)?;
-    pgbackend.run(&mut MgmtHandler)
+async fn handle_connection(socket: TcpStream) -> Result<(), QueryError> {
+    let pgbackend = PostgresBackend::new(socket, AuthType::Trust, None)?;
+    pgbackend.run(&mut MgmtHandler, future::pending::<()>).await
 }
 
 /// A message received by `mgmt` when a compute node is ready.
@@ -78,16 +70,21 @@ pub type ComputeReady = Result<DatabaseInfo, String>;
 
 // TODO: replace with an http-based protocol.
 struct MgmtHandler;
+#[async_trait::async_trait]
 impl postgres_backend::Handler for MgmtHandler {
-    fn process_query(&mut self, pgb: &mut PostgresBackend, query: &str) -> Result<(), QueryError> {
-        try_process_query(pgb, query).map_err(|e| {
+    async fn process_query(
+        &mut self,
+        pgb: &mut PostgresBackend,
+        query: &str,
+    ) -> Result<(), QueryError> {
+        try_process_query(pgb, query).await.map_err(|e| {
             error!("failed to process response: {e:?}");
             e
         })
     }
 }
 
-fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> Result<(), QueryError> {
+async fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> Result<(), QueryError> {
     let resp: KickSession = serde_json::from_str(query).context("Failed to parse query as json")?;
 
     let span = info_span!("event", session_id = resp.session_id);
@@ -98,11 +95,11 @@ fn try_process_query(pgb: &mut PostgresBackend, query: &str) -> Result<(), Query
         Ok(()) => {
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::DataRow(&[Some(b"ok")]))?
-                .write_message(&BeMessage::CommandComplete(b"SELECT 1"))?;
+                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         }
         Err(e) => {
             error!("failed to deliver response to per-client task");
-            pgb.write_message(&BeMessage::ErrorResponse(&e.to_string(), None))?;
+            pgb.write_message_noflush(&BeMessage::ErrorResponse(&e.to_string(), None))?;
         }
     }
 

proxy/src/stream.rs

@@ -1,45 +1,40 @@
 use crate::error::UserFacingError;
 use anyhow::bail;
-use bytes::BytesMut;
 use pin_project_lite::pin_project;
-use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket};
+use pq_proto::framed::{ConnectionError, Framed};
+use pq_proto::{BeMessage, FeMessage, FeStartupPacket, ProtocolError};
 use rustls::ServerConfig;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::{io, task};
 use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio_rustls::server::TlsStream;
 
-pin_project! {
-    /// Stream wrapper which implements libpq's protocol.
-    /// NOTE: This object deliberately doesn't implement [`AsyncRead`]
-    /// or [`AsyncWrite`] to prevent subtle errors (e.g. trying
-    /// to pass random malformed bytes through the connection).
-    pub struct PqStream<S> {
-        #[pin]
-        stream: S,
-        buffer: BytesMut,
-    }
+/// Stream wrapper which implements libpq's protocol.
+/// NOTE: This object deliberately doesn't implement [`AsyncRead`]
+/// or [`AsyncWrite`] to prevent subtle errors (e.g. trying
+/// to pass random malformed bytes through the connection).
+pub struct PqStream<S> {
+    framed: Framed<S>,
 }
 
 impl<S> PqStream<S> {
     /// Construct a new libpq protocol wrapper.
     pub fn new(stream: S) -> Self {
         Self {
-            stream,
-            buffer: Default::default(),
+            framed: Framed::new(stream),
         }
     }
 
     /// Extract the underlying stream.
     pub fn into_inner(self) -> S {
-        self.stream
+        self.framed.into_inner()
     }
 
     /// Get a shared reference to the underlying stream.
     pub fn get_ref(&self) -> &S {
-        &self.stream
+        self.framed.get_ref()
     }
 }
 
@@ -50,16 +45,19 @@ fn err_connection() -> io::Error {
 impl<S: AsyncRead + Unpin> PqStream<S> {
     /// Receive [`FeStartupPacket`], which is a first packet sent by a client.
     pub async fn read_startup_packet(&mut self) -> io::Result<FeStartupPacket> {
-        // TODO: `FeStartupPacket::read_fut` should return `FeStartupPacket`
-        let msg = FeStartupPacket::read_fut(&mut self.stream)
+        self.framed
+            .read_startup_message()
             .await
             .map_err(ConnectionError::into_io_error)?
-            .ok_or_else(err_connection)?;
+            .ok_or_else(err_connection)
+    }
 
-        match msg {
-            FeMessage::StartupPacket(packet) => Ok(packet),
-            _ => panic!("unreachable state"),
-        }
+    async fn read_message(&mut self) -> io::Result<FeMessage> {
+        self.framed
+            .read_message()
+            .await
+            .map_err(ConnectionError::into_io_error)?
+            .ok_or_else(err_connection)
     }
 
     pub async fn read_password_message(&mut self) -> io::Result<bytes::Bytes> {
@@ -71,19 +69,14 @@ impl<S: AsyncRead + Unpin> PqStream<S> {
             )),
         }
     }
-
-    async fn read_message(&mut self) -> io::Result<FeMessage> {
-        FeMessage::read_fut(&mut self.stream)
-            .await
-            .map_err(ConnectionError::into_io_error)?
-            .ok_or_else(err_connection)
-    }
 }
 
 impl<S: AsyncWrite + Unpin> PqStream<S> {
     /// Write the message into an internal buffer, but don't flush the underlying stream.
     pub fn write_message_noflush(&mut self, message: &BeMessage<'_>) -> io::Result<&mut Self> {
-        BeMessage::write(&mut self.buffer, message)?;
+        self.framed
+            .write_message(message)
+            .map_err(ProtocolError::into_io_error)?;
         Ok(self)
     }
 
@@ -96,9 +89,7 @@ impl<S: AsyncWrite + Unpin> PqStream<S> {
 
     /// Flush the output buffer into the underlying stream.
     pub async fn flush(&mut self) -> io::Result<&mut Self> {
-        self.stream.write_all(&self.buffer).await?;
-        self.buffer.clear();
-        self.stream.flush().await?;
+        self.framed.flush().await?;
         Ok(self)
     }