helper ReplicationStream type for replication protocol

This can be optionally used with a CopyBoth stream to decode the
replication protocol

Signed-off-by: Petros Angelatos <petrosagg@gmail.com>

Author: petrosagg

tokio-postgres/src/lib.rs

@@ -164,6 +164,7 @@ mod maybe_tls_stream;
 mod portal;
 mod prepare;
 mod query;
+pub mod replication;
 pub mod row;
 mod simple_query;
 #[cfg(feature = "runtime")]

tokio-postgres/src/replication.rs

@@ -0,0 +1,95 @@
+//! Utilities for working with the PostgreSQL replication copy both format.
+
+use crate::copy_both::CopyBothDuplex;
+use crate::Error;
+use bytes::{Bytes, BytesMut};
+use futures::{ready, SinkExt, Stream};
+use pin_project_lite::pin_project;
+use postgres_protocol::message::backend::ReplicationMessage;
+use postgres_protocol::message::frontend;
+use postgres_types::PgLsn;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+pin_project! {
+    /// A type which deserializes the postgres replication protocol. This type can be used with
+    /// both physical and logical replication to get access to the byte content of each replication
+    /// message.
+    ///
+    /// The replication *must* be explicitly completed via the `finish` method.
+    pub struct ReplicationStream {
+        #[pin]
+        stream: CopyBothDuplex<Bytes>,
+    }
+}
+
+impl ReplicationStream {
+    /// Creates a new ReplicationStream that will wrap the underlying CopyBoth stream
+    pub fn new(stream: CopyBothDuplex<Bytes>) -> Self {
+        Self { stream }
+    }
+
+    /// Send standby update to server.
+    pub async fn standby_status_update(
+        self: Pin<&mut Self>,
+        write_lsn: PgLsn,
+        flush_lsn: PgLsn,
+        apply_lsn: PgLsn,
+        ts: i64,
+        reply: u8,
+    ) -> Result<(), Error> {
+        let mut this = self.project();
+
+        let mut buf = BytesMut::new();
+        frontend::standby_status_update(
+            write_lsn.into(),
+            flush_lsn.into(),
+            apply_lsn.into(),
+            ts as i64,
+            reply,
+            &mut buf,
+        )
+        .expect("writes to memory never fail");
+        this.stream.send(buf.freeze()).await
+    }
+
+    /// Send hot standby feedback message to server.
+    pub async fn hot_standby_feedback(
+        self: Pin<&mut Self>,
+        timestamp: i64,
+        global_xmin: u32,
+        global_xmin_epoch: u32,
+        catalog_xmin: u32,
+        catalog_xmin_epoch: u32,
+    ) -> Result<(), Error> {
+        let mut this = self.project();
+
+        let mut buf = BytesMut::new();
+        frontend::hot_standby_feedback(
+            timestamp,
+            global_xmin,
+            global_xmin_epoch,
+            catalog_xmin,
+            catalog_xmin_epoch,
+            &mut buf,
+        )
+        .expect("writes to memory never fail");
+        this.stream.send(buf.freeze()).await
+    }
+}
+
+impl Stream for ReplicationStream {
+    type Item = Result<ReplicationMessage<Bytes>, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+
+        match ready!(this.stream.poll_next(cx)) {
+            Some(Ok(buf)) => {
+                Poll::Ready(Some(ReplicationMessage::parse(&buf).map_err(Error::parse)))
+            }
+            Some(Err(err)) => Poll::Ready(Some(Err(err))),
+            None => Poll::Ready(None),
+        }
+    }
+}