Markk116/livevue-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add pg_replication module for reactive WAL-based cache invalidation

Browse Source 
Implements a PostgreSQL logical replication module that streams WAL
changes, buffers them per XID for transaction atomicity, and publishes
CacheKey invalidation events to the existing framework broadcast channel.

Key design choices:
- Transaction atomicity: changes are buffered until COMMIT so the fanout
  always receives a consistent view. A 1000-row bulk-insert emits one
  table-level CacheKey, not 1000 row events.
- Native protocol: uses a self-contained raw TCP + postgres-protocol
  implementation for the replication connection (tokio-postgres 0.7 does
  not expose copy_both_simple publicly).
- Corrected pgoutput v1 parser: original maybe_sql_integration code
  incorrectly read XID from DML messages; DML messages carry no XID in
  proto v1 — only Begin does.
- Clean integration: publishes CacheKey::Channel events to store.events
  so the existing fanout task picks them up with zero changes to the
  fanout loop.
- New cx.subscribe(key) API on RenderContext for manually registering
  subscription keys (needed when queries go through PgPool, not cx.run).

New files:
  src/pg_replication/mod.rs       — PgReplicationListener, key helpers
  src/pg_replication/wal_parser.rs — pgoutput v1 binary protocol parser
  src/pg_replication/emitter.rs   — WalEmitter (tx-buffering + CacheKey emit)
  src/pg_replication/proto.rs     — raw TCP PG wire-protocol connection
  examples/pg_replication/main.rs — live message board example
  docker-compose.yml              — PG 16 container with wal_level=logical

Run the example:
  docker compose up -d
  cargo run --example pg_replication --features pg_replication

https://claude.ai/code/session_01SLGgXeqKV2o7KTmCaQZZPg
This commit is contained in:
Claude
2026-03-08 22:12:01 +00:00
parent 2358ca29c7
commit b57dbd0fa0
10 changed files with 2212 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pg_replication/proto.rs
+341
View File
@@ -0,0 +1,341 @@
//! Raw PostgreSQL wire-protocol connection for logical replication.
//!
//! `tokio-postgres` 0.7 does not expose a `copy_both_simple` API suitable for
//! the replication stream. This module implements the minimum subset of the
//! PostgreSQL wire protocol needed to:
//!
//! 1. Connect with `replication=database` startup parameter.
//! 2. Perform MD5 or cleartext password authentication.
//! 3. Send simple queries (for `CREATE_REPLICATION_SLOT` and `START_REPLICATION`).
//! 4. Enter `COPY BOTH` mode and stream `CopyData` frames.
//! 5. Send `StandbyStatusUpdate` CopyData frames back to the server.
//!
//! `postgres-protocol` is used for building most frontend messages and for
//! MD5 hash computation. The backend parser is bypassed for `CopyBothResponse`
//! (tag `'W'` / 0x57) which is not present in that crate's public API.
use super::standby_status_payload;
use anyhow::{bail, Context as _};
use bytes::{Buf, Bytes, BytesMut};
use postgres_protocol::authentication::md5_hash;
use postgres_protocol::message::frontend;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
// ---------------------------------------------------------------------------
// RawPgConn
// ---------------------------------------------------------------------------
/// A raw PostgreSQL connection suitable for replication streaming.
pub struct RawPgConn {
stream: TcpStream,
/// Accumulates bytes read from the server until a complete frame arrives.
read_buf: BytesMut,
}
impl RawPgConn {
// -----------------------------------------------------------------------
// Constructor
// -----------------------------------------------------------------------
/// Connect to a PostgreSQL server in **logical replication** mode.
///
/// Sends a startup message with `replication=database` so the backend
/// enters walsender mode. Handles MD5 and cleartext password auth.
pub async fn connect(
host: &str,
port: u16,
user: &str,
password: &str,
database: &str,
) -> anyhow::Result<Self> {
let addr = format!("{host}:{port}");
let stream = TcpStream::connect(&addr)
.await
.with_context(|| format!("TCP connect to {addr}"))?;
let mut conn = RawPgConn {
stream,
read_buf: BytesMut::with_capacity(4096),
};
// --- Startup message -----------------------------------------------
// Parameters include `replication=database` which instructs PostgreSQL
// to accept replication commands on this connection.
let mut startup_buf = BytesMut::new();
frontend::startup_message(
[
("user", user),
("database", database),
("replication", "database"),
("application_name", "livevue-rs"),
]
.iter()
.copied(),
&mut startup_buf,
)?;
conn.stream.write_all(&startup_buf).await?;
conn.stream.flush().await?;
// --- Authentication loop -------------------------------------------
loop {
let (tag, body) = conn.read_frame().await?;
match tag {
b'R' => {
// Authentication request
conn.handle_auth_request(&body, user, password).await?;
}
b'K' => {
// BackendKeyData — ignore
}
b'S' => {
// ParameterStatus — ignore
}
b'Z' => {
// ReadyForQuery — handshake complete
break;
}
b'E' => {
let msg = parse_error_message(&body);
bail!("PostgreSQL startup error: {msg}");
}
b'N' => {
// NoticeResponse — ignore
}
other => {
tracing::debug!("startup: unexpected tag 0x{:02x}", other);
}
}
}
Ok(conn)
}
// -----------------------------------------------------------------------
// Simple query (used for CREATE_REPLICATION_SLOT etc.)
// -----------------------------------------------------------------------
/// Send a simple query and drain responses until `ReadyForQuery`.
///
/// Returns `Err` if the server sends an `ErrorResponse`.
pub async fn simple_query(&mut self, query: &str) -> anyhow::Result<()> {
let mut buf = BytesMut::new();
frontend::query(query, &mut buf)?;
self.stream.write_all(&buf).await?;
self.stream.flush().await?;
loop {
let (tag, body) = self.read_frame().await?;
match tag {
b'Z' => break, // ReadyForQuery
b'E' => {
let msg = parse_error_message(&body);
bail!("query error: {msg}");
}
_ => {} // CommandComplete, RowDescription, DataRow, NoticeResponse …
}
}
Ok(())
}
// -----------------------------------------------------------------------
// Enter COPY BOTH mode for replication streaming
// -----------------------------------------------------------------------
/// Send `START_REPLICATION` and absorb the `CopyBothResponse`.
///
/// After this call the connection is in COPY BOTH mode. Use
/// [`read_copy_data`] / [`send_standby_status`] to exchange messages.
pub async fn start_replication(&mut self, query: &str) -> anyhow::Result<()> {
let mut buf = BytesMut::new();
frontend::query(query, &mut buf)?;
self.stream.write_all(&buf).await?;
self.stream.flush().await?;
// Expect CopyBothResponse ('W') followed by the replication stream.
// postgres-protocol doesn't parse 'W', so we handle it manually.
loop {
let (tag, _body) = self.read_frame().await?;
match tag {
b'W' => {
// CopyBothResponse — body contains column-format info we
// don't need for the replication stream. Just break out.
break;
}
b'E' => {
let msg = parse_error_message(&_body);
bail!("START_REPLICATION error: {msg}");
}
b'N' => {} // NoticeResponse — skip
other => {
tracing::debug!("start_replication: unexpected tag 0x{:02x}", other);
}
}
}
Ok(())
}
// -----------------------------------------------------------------------
// Replication stream I/O
// -----------------------------------------------------------------------
/// Read one `CopyData` payload from the replication stream.
///
/// Returns:
/// - `Ok(Some(bytes))` — a CopyData payload (XLogData or PrimaryKeepalive).
/// - `Ok(None)` — the server sent `CopyDone` or `ErrorResponse` (stream end).
/// - `Err(_)` — I/O or protocol error.
pub async fn read_copy_data(&mut self) -> anyhow::Result<Option<Bytes>> {
let (tag, body) = self.read_frame().await?;
match tag {
b'd' => Ok(Some(body)),
b'c' => Ok(None), // CopyDone
b'E' => {
let msg = parse_error_message(&body);
bail!("replication error: {msg}");
}
other => {
tracing::debug!("copy stream: unexpected tag 0x{:02x}", other);
Ok(Some(Bytes::new()))
}
}
}
/// Send a `StandbyStatusUpdate` CopyData frame to acknowledge `lsn`.
pub async fn send_standby_status(&mut self, lsn: u64) -> anyhow::Result<()> {
let payload = standby_status_payload(lsn);
let mut buf = BytesMut::new();
frontend::CopyData::new(&payload[..])?.write(&mut buf);
self.stream.write_all(&buf).await?;
self.stream.flush().await?;
Ok(())
}
// -----------------------------------------------------------------------
// Low-level frame I/O
// -----------------------------------------------------------------------
/// Read one complete PostgreSQL message frame.
///
/// Frame format: `tag (u8) | length (i32, BE, includes itself) | body`.
/// Returns `(tag, body_bytes)`.
async fn read_frame(&mut self) -> anyhow::Result<(u8, Bytes)> {
// Ensure we have at least 5 bytes (1 tag + 4 length).
while self.read_buf.len() < 5 {
let n = self.stream.read_buf(&mut self.read_buf).await?;
if n == 0 {
bail!("server closed connection unexpectedly");
}
}
let tag = self.read_buf[0];
// Length field is i32 big-endian, includes the 4-length bytes but NOT
// the tag byte. Total frame length = 1 (tag) + length.
let body_len = (i32::from_be_bytes(
self.read_buf[1..5].try_into().expect("4 byte slice"),
) as usize)
.saturating_sub(4); // subtract the 4 length bytes themselves
let total = 5 + body_len; // tag(1) + length(4) + body
while self.read_buf.len() < total {
let n = self.stream.read_buf(&mut self.read_buf).await?;
if n == 0 {
bail!("server closed connection while reading message body");
}
}
self.read_buf.advance(5); // consume tag + length
let body = self.read_buf.split_to(body_len).freeze();
Ok((tag, body))
}
// -----------------------------------------------------------------------
// Authentication
// -----------------------------------------------------------------------
async fn handle_auth_request(
&mut self,
body: &[u8],
user: &str,
password: &str,
) -> anyhow::Result<()> {
if body.len() < 4 {
bail!("AuthenticationRequest body too short");
}
let auth_type = i32::from_be_bytes(body[0..4].try_into().expect("4 bytes"));
match auth_type {
0 => {
// AuthenticationOk — nothing to do
}
3 => {
// AuthenticationCleartextPassword
let mut buf = BytesMut::new();
frontend::password_message(password.as_bytes(), &mut buf)?;
self.stream.write_all(&buf).await?;
self.stream.flush().await?;
}
5 => {
// AuthenticationMD5Password — salt is bytes 4..8
if body.len() < 8 {
bail!("MD5 auth: body too short for salt");
}
let salt: [u8; 4] = body[4..8].try_into().expect("4 byte salt");
let hashed = md5_hash(user.as_bytes(), password.as_bytes(), salt);
let mut buf = BytesMut::new();
frontend::password_message(hashed.as_bytes(), &mut buf)?;
self.stream.write_all(&buf).await?;
self.stream.flush().await?;
}
10 => {
// AuthenticationSASL (SCRAM-SHA-256) — not implemented here.
// For Docker dev setups, use `md5` or `trust` in pg_hba.conf.
bail!(
"SCRAM-SHA-256 authentication is not supported by the \
replication module. Configure PostgreSQL to use `md5` \
or `trust` for this user in pg_hba.conf."
);
}
other => {
bail!("unsupported authentication type: {other}");
}
}
Ok(())
}
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// Extract a human-readable message from an `ErrorResponse` body.
///
/// Fields are null-terminated key-value pairs. We extract the 'M' (message)
/// and 'D' (detail) fields.
fn parse_error_message(body: &[u8]) -> String {
let mut parts = Vec::new();
let mut i = 0;
while i < body.len() {
let field_type = body[i];
i += 1;
if field_type == 0 {
break;
}
// Find the null terminator of the field value.
let start = i;
while i < body.len() && body[i] != 0 {
i += 1;
}
let value = String::from_utf8_lossy(&body[start..i]).into_owned();
i += 1; // skip null terminator
if matches!(field_type, b'M' | b'D' | b'C') {
parts.push(value);
}
}
if parts.is_empty() {
"unknown PostgreSQL error".into()
} else {
parts.join(" — ")
}
}