cleanup_status_update and cleanup replication test

Author: ruslantalpa

docker-compose.yml

@@ -6,4 +6,6 @@ services:
       - 5433:5433
     environment:
       - POSTGRES_PASSWORD=pass
+    volumes:
+      - "./docker/sql_setup.sh:/docker-entrypoint-initdb.d/sql_setup.sh"
 

docker/sql_setup.sh

@@ -66,7 +66,6 @@ ssl_cert_file = 'server.crt'
 ssl_key_file = 'server.key'
 wal_level = logical
 log_statement = 'all'
-shared_preload_libraries = 'wal2json'
 EOCONF
 
 cat > "$PGDATA/pg_hba.conf" <<-EOCONF

postgres-protocol/src/message/frontend.rs

@@ -132,6 +132,26 @@ 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,
+    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(0);
+        Ok(())
+    })
+}
+
 pub struct CopyData<T> {
     buf: T,
     len: i32,

tokio-postgres/src/client.rs

@@ -435,16 +435,26 @@ impl Client {
         copy_out::copy_out(self.inner(), statement).await
     }
 
-    /// TODO!
+    /// Executes a 'START_REPLICATION SLOT ...', returning a stream of raw replication events
     pub async fn start_replication(&self, query: &str) -> Result<ReplicationStream, Error> {
         replication::start_replication(self.inner(), query).await
     }
 
-    /// TODO!
+    /// Stoppes the current replication by sending a copy_done message
     pub async fn stop_replication(&self) -> Result<(), Error> {
         replication::stop_replication(self.inner()).await
     }
 
+    /// Notifies PostgreSQL of the last processed WAL
+    pub async fn standby_status_update(
+        &self,
+        write_lsn: i64,
+        flush_lsn: i64,
+        apply_lsn: i64,
+    ) -> Result<(), Error> {
+        replication::standby_status_update(self.inner(), write_lsn, flush_lsn, apply_lsn).await
+    }
+
     /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows.
     ///
     /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that

tokio-postgres/src/replication.rs

@@ -11,7 +11,8 @@ use postgres_protocol::message::frontend;
 use std::marker::PhantomPinned;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-
+use std::time::{SystemTime, UNIX_EPOCH};
+const J2000_EPOCH_GAP: u128 = 946_684_800_000_000;
 pub async fn start_replication(
     client: &InnerClient,
     query: &str,
@@ -34,6 +35,24 @@ pub async fn stop_replication(client: &InnerClient) -> Result<(), Error> {
     Ok(())
 }
 
+pub async fn standby_status_update(
+    client: &InnerClient,
+    write_lsn: i64,
+    flush_lsn: i64,
+    apply_lsn: i64,
+) -> Result<(), Error> {
+    trace!("executing standby_status_update");
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap()
+        .as_micros()
+        - J2000_EPOCH_GAP;
+    let mut buf = BytesMut::new();
+    let _ = frontend::standby_status_update(write_lsn, flush_lsn, apply_lsn, now as i64, &mut buf);
+    let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf.freeze())))?;
+    Ok(())
+}
+
 async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
     let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
     trace!("start in repication");

tokio-postgres/tests/test/main.rs

@@ -5,6 +5,7 @@ use futures::channel::mpsc;
 use futures::{
     future, join, pin_mut, stream, try_join, FutureExt, SinkExt, StreamExt, TryStreamExt,
 };
+use std::convert::TryInto;
 use std::fmt::Write;
 use std::time::Duration;
 use tokio::net::TcpStream;
@@ -21,7 +22,6 @@ mod parse;
 #[cfg(feature = "runtime")]
 mod runtime;
 mod types;
-use std::sync::{Arc, Mutex};
 use std::thread;
 async fn connect_raw(s: &str) -> Result<(Client, Connection<TcpStream, NoTlsStream>), Error> {
     let socket = TcpStream::connect("127.0.0.1:5433").await.unwrap();
@@ -808,70 +808,77 @@ async fn replication_start() {
 
     client
         .batch_execute(
-            "CREATE TABLE IF NOT EXISTS foo (
+            "CREATE TABLE IF NOT EXISTS replication (
             id SERIAL,
             name TEXT
         );
-        DROP PUBLICATION IF EXISTS rust;
-        CREATE PUBLICATION rust FOR ALL TABLES;
         ",
         )
         .await
         .unwrap();
     let _ = r_client
-        .simple_query("DROP_REPLICATION_SLOT rust_slot")
-        .await;
-    // let _ = r_client.simple_query("CREATE_REPLICATION_SLOT rust_slot LOGICAL pgoutput NOEXPORT_SNAPSHOT").await;
-    let _ = r_client
-        .simple_query("CREATE_REPLICATION_SLOT rust_slot LOGICAL wal2json NOEXPORT_SNAPSHOT")
+        .simple_query(
+            "CREATE_REPLICATION_SLOT rust_slot TEMPORARY LOGICAL test_decoding NOEXPORT_SNAPSHOT",
+        )
         .await;
 
     let stream = r_client
-        //.start_replication("START_REPLICATION SLOT rust_slot LOGICAL 0/0 (proto_version '1',  publication_names 'rust')")
-        .start_replication("START_REPLICATION SLOT rust_slot LOGICAL 0/0 (\"pretty-print\" '1', \"format-version\" '2')")
+        .start_replication(
+            "START_REPLICATION SLOT rust_slot  LOGICAL 0/0 (\"skip-empty-xacts\" '1')",
+        )
         .await
         .unwrap();
 
-    let events = Arc::new(Mutex::new(vec![]));
-    let total_events = Arc::new(Mutex::new(0));
-    let total_events_1 = Arc::clone(&total_events);
-    let events_1 = Arc::clone(&events);
-    let t = tokio::spawn(async move {
-        let events = stream.try_collect::<Vec<_>>().await.unwrap();
-        let mut num = total_events_1.lock().unwrap();
-        *num = events.len();
-        let mut ev = events_1.lock().unwrap();
-        *ev = events;
-    });
-
     client
-        .query("INSERT INTO foo (name) VALUES ('ann')", &[])
+        .query("INSERT INTO replication (name) VALUES ('ann')", &[])
         .await
         .unwrap();
-    client
-        .query("INSERT INTO foo (name) VALUES ('jim'), ('joe')", &[])
-        .await
-        .unwrap();
-    client.query("DROP TABLE foo", &[]).await.unwrap();
 
     thread::sleep(time::Duration::from_secs(1)); //give a chance to pg to send the events
     let _ = r_client.stop_replication().await;
-    let _ = t.await;
-    println!("events {:?}", *events.lock().unwrap());
-    //assert_eq!(*total_events.lock().unwrap(), 2);
-    for e in &*events.lock().unwrap() {
+    let events = stream.try_collect::<Vec<_>>().await.unwrap();
+    client.query("DROP TABLE replication", &[]).await.unwrap();
+
+    let mut total_events = 0;
+    for e in &events {
         match e[0].into() {
             'k' => {
-                // let message_type = "keepalive";
-                // println!("type: ({}), {}", message_type, e.len());
+                //keepalive message
+                let current_wal_end =
+                    i64::from_be_bytes(e.slice(1..9).as_ref().try_into().unwrap());
+                let timestamp = i64::from_be_bytes(e.slice(9..17).as_ref().try_into().unwrap());
+                let reply: char = e[17].into();
+                println!(
+                    "keepalive tiemstamp: {} current_wal_end: {:X}/{:X} reply: {}",
+                    timestamp,
+                    (current_wal_end >> 32) as i32,
+                    current_wal_end as i32,
+                    reply as i8
+                );
             }
             'w' => {
-                let message_type = "dataframe";
-                let slice = e.slice(25..);
-                let data = std::str::from_utf8(slice.as_ref()).unwrap();
-                println!("type: ({}), {}, {}", message_type, e.len(), data);
+                // WAL message
+                let current_wal = i64::from_be_bytes(e.slice(1..9).as_ref().try_into().unwrap());
+                let current_wal_end =
+                    i64::from_be_bytes(e.slice(9..17).as_ref().try_into().unwrap());
+                let timestamp = i64::from_be_bytes(e.slice(17..25).as_ref().try_into().unwrap());
+                let _data = e.slice(25..); //the format of these bytes depends on the logical decoder
+                println!(
+                    "WAL timestamp: {} current_wal: {:X}/{:X} current_wal_end: {:X}/{:X} {:?}",
+                    timestamp,
+                    (current_wal >> 32) as i32,
+                    current_wal as i32,
+                    (current_wal_end >> 32) as i32,
+                    current_wal_end as i32,
+                    _data
+                );
+                total_events += 1;
+                //while in replication state, one needs to send updates from time to time like this
+                //let _ = r_pg_client.standby_status_update(current_wal, current_wal, current_wal).await.unwrap();
             }
             _ => {}
         };
     }
+    // in this case we receive 3 events (BEGIN,INSERT,COMMIT)
+    assert_eq!(total_events, 3);
 }