fix: possible page_service hang on cancel (#5696)

Fixes #5341, one more suspected case, see:
https://github.com/neondatabase/neon/issues/5341#issuecomment-1783052379

- races `MaybeWriteOnly::shutdown` with cancellation
- switches to using `AsyncWriteExt::write_buf`
- notes cancellation safety for shutdown

libs/postgres_backend/src/lib.rs

@@ -242,6 +242,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> MaybeWriteOnly<IO> {
         }
     }
 
+    /// Cancellation safe as long as the underlying IO is cancellation safe.
     async fn shutdown(&mut self) -> io::Result<()> {
         match self {
             MaybeWriteOnly::Full(framed) => framed.shutdown().await,
@@ -393,13 +394,23 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
         shutdown_watcher: F,
     ) -> Result<(), QueryError>
     where
-        F: Fn() -> S,
+        F: Fn() -> S + Clone,
         S: Future,
     {
-        let ret = self.run_message_loop(handler, shutdown_watcher).await;
-        // socket might be already closed, e.g. if previously received error,
-        // so ignore result.
-        self.framed.shutdown().await.ok();
+        let ret = self
+            .run_message_loop(handler, shutdown_watcher.clone())
+            .await;
+
+        tokio::select! {
+            _ = shutdown_watcher() => {
+                // do nothing; we most likely got already stopped by shutdown and will log it next.
+            }
+            _ = self.framed.shutdown() => {
+                // socket might be already closed, e.g. if previously received error,
+                // so ignore result.
+            },
+        }
+
         match ret {
             Ok(()) => Ok(()),
             Err(QueryError::Shutdown) => {

libs/pq_proto/src/framed.rs

@@ -214,27 +214,24 @@ where
     }
 }
 
+/// Cancellation safe as long as the AsyncWrite is cancellation safe.
 async fn flush<S: AsyncWrite + Unpin>(
     stream: &mut S,
     write_buf: &mut BytesMut,
 ) -> Result<(), io::Error> {
     while write_buf.has_remaining() {
-        let bytes_written = stream.write(write_buf.chunk()).await?;
+        let bytes_written = stream.write_buf(write_buf).await?;
         if bytes_written == 0 {
             return Err(io::Error::new(
                 ErrorKind::WriteZero,
                 "failed to write message",
             ));
         }
-        // The advanced part will be garbage collected, likely during shifting
-        // data left on next attempt to write to buffer when free space is not
-        // enough.
-        write_buf.advance(bytes_written);
     }
-    write_buf.clear();
     stream.flush().await
 }
 
+/// Cancellation safe as long as the AsyncWrite is cancellation safe.
 async fn shutdown<S: AsyncWrite + Unpin>(
     stream: &mut S,
     write_buf: &mut BytesMut,