feat: add support for quick active health checks on connections

tokio-postgres/src/client.rs

@@ -531,6 +531,12 @@ impl Client {
         simple_query::batch_execute(self.inner(), query).await
     }
 
+    /// Check the connection is alive and wait for the confirmation.
+    pub async fn check_connection(&self) -> Result<(), Error> {
+        // sync is a very quick message to test the connection health.
+        query::sync(self.inner()).await
+    }
+
     /// Begins a new database transaction.
     ///
     /// The transaction will roll back by default - use the `commit` method to commit it.

tokio-postgres/src/query.rs

@@ -323,3 +323,13 @@ impl RowStream {
         self.rows_affected
     }
 }
+
+pub async fn sync(client: &InnerClient) -> Result<(), Error> {
+    let buf = Bytes::from_static(b"S\0\0\0\x04");
+    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+    match responses.next().await? {
+        Message::ReadyForQuery(_) => Ok(()),
+        _ => Err(Error::unexpected_message()),
+    }
+}

tokio-postgres/tests/test/main.rs

@@ -147,6 +147,12 @@ async fn scram_password_ok() {
     connect("user=scram_user password=password dbname=postgres").await;
 }
 
+#[tokio::test]
+async fn sync() {
+    let client = connect("user=postgres").await;
+    client.check_connection().await.unwrap();
+}
+
 #[tokio::test]
 async fn pipelined_prepare() {
     let client = connect("user=postgres").await;