Replication support.

Author: jeff-davis

postgres-protocol/src/message/backend.rs

@@ -11,6 +11,7 @@ use std::str;
 
 use crate::Oid;
 
+// top-level message tags
 pub const PARSE_COMPLETE_TAG: u8 = b'1';
 pub const BIND_COMPLETE_TAG: u8 = b'2';
 pub const CLOSE_COMPLETE_TAG: u8 = b'3';
@@ -22,6 +23,7 @@ pub const DATA_ROW_TAG: u8 = b'D';
 pub const ERROR_RESPONSE_TAG: u8 = b'E';
 pub const COPY_IN_RESPONSE_TAG: u8 = b'G';
 pub const COPY_OUT_RESPONSE_TAG: u8 = b'H';
+pub const COPY_BOTH_RESPONSE_TAG: u8 = b'W';
 pub const EMPTY_QUERY_RESPONSE_TAG: u8 = b'I';
 pub const BACKEND_KEY_DATA_TAG: u8 = b'K';
 pub const NO_DATA_TAG: u8 = b'n';
@@ -33,6 +35,10 @@ pub const PARAMETER_DESCRIPTION_TAG: u8 = b't';
 pub const ROW_DESCRIPTION_TAG: u8 = b'T';
 pub const READY_FOR_QUERY_TAG: u8 = b'Z';
 
+// replication message tags
+pub const XLOG_DATA_TAG: u8 = b'w';
+pub const PRIMARY_KEEPALIVE_TAG: u8 = b'k';
+
 #[derive(Debug, Copy, Clone)]
 pub struct Header {
     tag: u8,
@@ -93,6 +99,7 @@ pub enum Message {
     CopyDone,
     CopyInResponse(CopyInResponseBody),
     CopyOutResponse(CopyOutResponseBody),
+    CopyBothResponse(CopyBothResponseBody),
     DataRow(DataRowBody),
     EmptyQueryResponse,
     ErrorResponse(ErrorResponseBody),
@@ -190,6 +197,16 @@ impl Message {
                     storage,
                 })
             }
+            COPY_BOTH_RESPONSE_TAG => {
+                let format = buf.read_u8()?;
+                let len = buf.read_u16::<BigEndian>()?;
+                let storage = buf.read_all();
+                Message::CopyBothResponse(CopyBothResponseBody {
+                    format,
+                    len,
+                    storage,
+                })
+            }
             EMPTY_QUERY_RESPONSE_TAG => Message::EmptyQueryResponse,
             BACKEND_KEY_DATA_TAG => {
                 let process_id = buf.read_i32::<BigEndian>()?;
@@ -278,6 +295,58 @@ impl Message {
     }
 }
 
+/// An enum representing Postgres backend replication messages.
+#[non_exhaustive]
+#[derive(Debug)]
+pub enum ReplicationMessage {
+    XLogData(XLogDataBody),
+    PrimaryKeepAlive(PrimaryKeepAliveBody),
+}
+
+impl ReplicationMessage {
+    pub fn parse(bytes: &Bytes) -> io::Result<ReplicationMessage> {
+        let mut buf = Buffer {
+            bytes: bytes.clone(),
+            idx: 0,
+        };
+
+        let tag = buf.read_u8()?;
+
+        let replication_message = match tag {
+            XLOG_DATA_TAG => {
+                let wal_start = buf.read_u64::<BigEndian>()?;
+                let wal_end = buf.read_u64::<BigEndian>()?;
+                let timestamp = buf.read_i64::<BigEndian>()?;
+                let storage = buf.read_all();
+                ReplicationMessage::XLogData(XLogDataBody {
+                    wal_start,
+                    wal_end,
+                    timestamp,
+                    storage,
+                })
+            }
+            PRIMARY_KEEPALIVE_TAG => {
+                let wal_end = buf.read_u64::<BigEndian>()?;
+                let timestamp = buf.read_i64::<BigEndian>()?;
+                let reply = buf.read_u8()?;
+                ReplicationMessage::PrimaryKeepAlive(PrimaryKeepAliveBody {
+                    wal_end,
+                    timestamp,
+                    reply,
+                })
+            }
+            tag => {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    format!("unknown replication message tag `{}`", tag),
+                ));
+            }
+        };
+
+        Ok(replication_message)
+    }
+}
+
 struct Buffer {
     bytes: Bytes,
     idx: usize,
@@ -524,6 +593,27 @@ impl CopyOutResponseBody {
     }
 }
 
+pub struct CopyBothResponseBody {
+    storage: Bytes,
+    len: u16,
+    format: u8,
+}
+
+impl CopyBothResponseBody {
+    #[inline]
+    pub fn format(&self) -> u8 {
+        self.format
+    }
+
+    #[inline]
+    pub fn column_formats(&self) -> ColumnFormats<'_> {
+        ColumnFormats {
+            remaining: self.len,
+            buf: &self.storage,
+        }
+    }
+}
+
 pub struct DataRowBody {
     storage: Bytes,
     len: u16,
@@ -776,6 +866,65 @@ impl RowDescriptionBody {
     }
 }
 
+#[derive(Debug)]
+pub struct XLogDataBody {
+    wal_start: u64,
+    wal_end: u64,
+    timestamp: i64,
+    storage: Bytes,
+}
+
+impl XLogDataBody {
+    #[inline]
+    pub fn wal_start(&self) -> u64 {
+        self.wal_start
+    }
+
+    #[inline]
+    pub fn wal_end(&self) -> u64 {
+        self.wal_end
+    }
+
+    #[inline]
+    pub fn timestamp(&self) -> i64 {
+        self.timestamp
+    }
+
+    #[inline]
+    pub fn data(&self) -> &[u8] {
+        &self.storage
+    }
+
+    #[inline]
+    pub fn into_bytes(self) -> Bytes {
+        self.storage
+    }
+}
+
+#[derive(Debug)]
+pub struct PrimaryKeepAliveBody {
+    wal_end: u64,
+    timestamp: i64,
+    reply: u8,
+}
+
+impl PrimaryKeepAliveBody {
+    #[inline]
+    pub fn wal_end(&self) -> u64 {
+        self.wal_end
+    }
+
+    #[inline]
+    pub fn timestamp(&self) -> i64 {
+        self.timestamp
+    }
+
+    #[inline]
+    pub fn reply(&self) -> u8 {
+        self.reply
+    }
+}
+
 pub struct Fields<'a> {
     buf: &'a [u8],
     remaining: u16,

postgres-protocol/src/message/frontend.rs

@@ -132,6 +132,48 @@ pub fn close(variant: u8, name: &str, buf: &mut BytesMut) -> io::Result<()> {
     })
 }
 
+#[inline]
+pub fn standby_status_update(
+    write_lsn: i64,
+    flush_lsn: i64,
+    apply_lsn: i64,
+    timestamp: i64,
+    reply: u8,
+    buf: &mut BytesMut,
+) -> io::Result<()> {
+    buf.put_u8(b'd');
+    write_body(buf, |buf| {
+        buf.put_u8(b'r');
+        buf.put_i64(write_lsn + 1);
+        buf.put_i64(flush_lsn + 1);
+        buf.put_i64(apply_lsn + 1);
+        buf.put_i64(timestamp);
+        buf.put_u8(reply);
+        Ok(())
+    })
+}
+
+#[inline]
+pub fn hot_standby_feedback(
+    timestamp: i64,
+    global_xmin: u32,
+    global_xmin_epoch: u32,
+    catalog_xmin: u32,
+    catalog_xmin_epoch: u32,
+    buf: &mut BytesMut,
+) -> io::Result<()> {
+    buf.put_u8(b'd');
+    write_body(buf, |buf| {
+        buf.put_u8(b'h');
+        buf.put_i64(timestamp);
+        buf.put_u32(global_xmin);
+        buf.put_u32(global_xmin_epoch);
+        buf.put_u32(catalog_xmin);
+        buf.put_u32(catalog_xmin_epoch);
+        Ok(())
+    })
+}
+
 pub struct CopyData<T> {
     buf: T,
     len: i32,

tokio-postgres/src/config.rs

@@ -56,6 +56,16 @@ pub enum ChannelBinding {
     Require,
 }
 
+/// Replication mode configuration.
+#[derive(Debug, Copy, Clone, PartialEq)]
+#[non_exhaustive]
+pub enum ReplicationMode {
+    /// Physical replication.
+    Physical,
+    /// Logical replication.
+    Logical,
+}
+
 /// A host specification.
 #[derive(Debug, Clone, PartialEq)]
 pub enum Host {
@@ -159,6 +169,7 @@ pub struct Config {
     pub(crate) keepalives_idle: Duration,
     pub(crate) target_session_attrs: TargetSessionAttrs,
     pub(crate) channel_binding: ChannelBinding,
+    pub(crate) replication_mode: Option<ReplicationMode>,
 }
 
 impl Default for Config {
@@ -184,6 +195,7 @@ impl Config {
             keepalives_idle: Duration::from_secs(2 * 60 * 60),
             target_session_attrs: TargetSessionAttrs::Any,
             channel_binding: ChannelBinding::Prefer,
+            replication_mode: None,
         }
     }
 
@@ -387,6 +399,17 @@ impl Config {
         self.channel_binding
     }
 
+    /// Set replication mode.
+    pub fn replication_mode(&mut self, replication_mode: ReplicationMode) -> &mut Config {
+        self.replication_mode = Some(replication_mode);
+        self
+    }
+
+    /// Get replication mode.
+    pub fn get_replication_mode(&self) -> Option<ReplicationMode> {
+        self.replication_mode
+    }
+
     fn param(&mut self, key: &str, value: &str) -> Result<(), Error> {
         match key {
             "user" => {
@@ -534,6 +557,12 @@ impl fmt::Debug for Config {
             }
         }
 
+        let replication_mode_str = match self.replication_mode {
+            None => "false",
+            Some(ReplicationMode::Physical) => "true",
+            Some(ReplicationMode::Logical) => "database",
+        };
+
         f.debug_struct("Config")
             .field("user", &self.user)
             .field("password", &self.password.as_ref().map(|_| Redaction {}))
@@ -548,6 +577,7 @@ impl fmt::Debug for Config {
             .field("keepalives_idle", &self.keepalives_idle)
             .field("target_session_attrs", &self.target_session_attrs)
             .field("channel_binding", &self.channel_binding)
+            .field("replication", &replication_mode_str.to_string())
             .finish()
     }
 }

tokio-postgres/src/connect_raw.rs

@@ -1,5 +1,5 @@
 use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
-use crate::config::{self, Config};
+use crate::config::{self, Config, ReplicationMode};
 use crate::connect_tls::connect_tls;
 use crate::maybe_tls_stream::MaybeTlsStream;
 use crate::tls::{TlsConnect, TlsStream};
@@ -124,6 +124,12 @@ where
     if let Some(application_name) = &config.application_name {
         params.push(("application_name", &**application_name));
     }
+    if let Some(replication_mode) = &config.replication_mode {
+        match replication_mode {
+            ReplicationMode::Physical => params.push(("replication", "true")),
+            ReplicationMode::Logical => params.push(("replication", "database")),
+        }
+    }
 
     let mut buf = BytesMut::new();
     frontend::startup_message(params, &mut buf).map_err(Error::encode)?;