Remove sync postgres_backend, tidy up its split usage.

- Add support for splitting async postgres_backend into read and write halfes.
  Safekeeper needs this for bidirectional streams. To this end, encapsulate
  reading-writing postgres messages to framed.rs with split support without any
  additional changes (relying on BufRead for reading and BytesMut out buffer for
  writing).
- Use async postgres_backend throughout safekeeper (and in proxy auth link
  part).
- 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/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::QueryError,
-    postgres_backend::{self, AuthType, PostgresBackend},
-};
 
 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

@@ -51,7 +51,7 @@ 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)
+        let msg = FeStartupPacket::read(&mut self.stream)
             .await
             .map_err(ConnectionError::into_io_error)?
             .ok_or_else(err_connection)?;
@@ -73,7 +73,7 @@ impl<S: AsyncRead + Unpin> PqStream<S> {
     }
 
     async fn read_message(&mut self) -> io::Result<FeMessage> {
-        FeMessage::read_fut(&mut self.stream)
+        FeMessage::read(&mut self.stream)
             .await
             .map_err(ConnectionError::into_io_error)?
             .ok_or_else(err_connection)