feat: add pg_replication module for reactive WAL-based cache invalidation
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
@@ -0,0 +1,387 @@
|
||||
//! # pg_replication example
|
||||
//!
|
||||
//! A live message board backed by PostgreSQL.
|
||||
//!
|
||||
//! Run:
|
||||
//! ```sh
|
||||
//! docker compose up -d # start the PG container (see docker-compose.yml)
|
||||
//! cargo run --example pg_replication --features pg_replication
|
||||
//! ```
|
||||
//! Then open http://localhost:3001 in two browser tabs and watch them stay in sync.
|
||||
//!
|
||||
//! ## How it works
|
||||
//!
|
||||
//! 1. The render function queries the `messages` table directly via `sqlx::PgPool`
|
||||
//! and calls `cx.subscribe(pg_table_key("public", "messages"))` to tell the
|
||||
//! framework "re-render this connection whenever `pg:public.messages` is
|
||||
//! invalidated".
|
||||
//!
|
||||
//! 2. `PgReplicationListener` connects to PostgreSQL over the logical replication
|
||||
//! protocol, buffers changes per transaction, and on each COMMIT publishes
|
||||
//! `CacheKey::Channel { key: "pg:public.messages" }` to the store's broadcast
|
||||
//! channel.
|
||||
//!
|
||||
//! 3. The fanout task picks up the event, finds all connections subscribed to that
|
||||
//! key, and triggers a re-render for each — sending fresh HTML over SSE.
|
||||
//!
|
||||
//! No polling. No websocket boilerplate. Just PostgreSQL WAL + SSE.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::{
|
||||
extract::{Json, State},
|
||||
http::StatusCode,
|
||||
response::{Html, IntoResponse},
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
use maud::{html, Markup, DOCTYPE};
|
||||
use serde::Deserialize;
|
||||
use sqlx::FromRow;
|
||||
use uuid::Uuid;
|
||||
|
||||
use livevue_rs::{
|
||||
pg_replication::{pg_table_key, PgReplicationConfig, PgReplicationListener},
|
||||
server::{ingest_signals, ActionParams, AppState, SharedState},
|
||||
spawn_fanout, RenderContext, Store,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Domain types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone, FromRow)]
|
||||
struct Message {
|
||||
id: i64,
|
||||
body: String,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Database helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn setup_db(pool: &sqlx::PgPool) -> anyhow::Result<()> {
|
||||
// Enable logical replication for this table (requires superuser or
|
||||
// rds_superuser on managed PG).
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS messages (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
body TEXT NOT NULL
|
||||
)",
|
||||
)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
|
||||
// Create the publication if it doesn't exist.
|
||||
// `IF NOT EXISTS` was added in PG 14 — fall back gracefully on older versions.
|
||||
let exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS (SELECT 1 FROM pg_publication WHERE pubname = 'livevue_pub')",
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
|
||||
if !exists {
|
||||
sqlx::query("CREATE PUBLICATION livevue_pub FOR TABLE messages")
|
||||
.execute(pool)
|
||||
.await?;
|
||||
println!("Created publication livevue_pub");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Rendering
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Shared render state — anything the render function needs beyond the store.
|
||||
#[derive(Clone)]
|
||||
struct RenderState {
|
||||
pg: sqlx::PgPool,
|
||||
}
|
||||
|
||||
async fn render_page(cx: &mut RenderContext, rs: &RenderState) -> anyhow::Result<Markup> {
|
||||
// Subscribe to PG-sourced invalidation events for the messages table.
|
||||
// The fanout task will re-render this connection whenever the WAL emitter
|
||||
// fires a `CacheKey::Channel { key: "pg:public.messages" }` event.
|
||||
cx.subscribe(pg_table_key("public", "messages"));
|
||||
|
||||
let messages: Vec<Message> = sqlx::query_as("SELECT id, body FROM messages ORDER BY id DESC")
|
||||
.fetch_all(&rs.pg)
|
||||
.await?;
|
||||
|
||||
let new_msg = cx.signal("newMsg", livevue_rs::global());
|
||||
|
||||
Ok(html! {
|
||||
div #app {
|
||||
h1 { "Live Message Board" }
|
||||
p.subtitle {
|
||||
"Messages are stored in PostgreSQL. "
|
||||
"Add one and watch every open tab update in real time via WAL replication."
|
||||
}
|
||||
|
||||
// ── New message form ─────────────────────────────────────────
|
||||
div.compose {
|
||||
input
|
||||
type="text"
|
||||
placeholder="Type a message…"
|
||||
data-bind=(new_msg.name)
|
||||
data-on:keydown="{if (event.key === 'Enter') $post('/action/post_message')}"
|
||||
;
|
||||
button
|
||||
data-on:click="$post('/action/post_message')"
|
||||
data-attr=(format!("{{disabled: {}.trim() === ''}}", new_msg.val))
|
||||
{ "Send" }
|
||||
}
|
||||
|
||||
// ── Message list ─────────────────────────────────────────────
|
||||
@if messages.is_empty() {
|
||||
p.empty { "No messages yet. Be the first!" }
|
||||
} @else {
|
||||
ul.messages {
|
||||
@for msg in &messages {
|
||||
li {
|
||||
span.id { "#" (msg.id) }
|
||||
span.body { (msg.body) }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Full HTML shell for the initial page load.
|
||||
fn page_shell(conn_id: Uuid, signals_json: &str, inner: &Markup) -> String {
|
||||
html! {
|
||||
(DOCTYPE)
|
||||
html lang="en" {
|
||||
head {
|
||||
meta charset="utf-8";
|
||||
meta name="viewport" content="width=device-width, initial-scale=1";
|
||||
title { "pg_replication example — livevue-rs" }
|
||||
style {
|
||||
(r#"
|
||||
body { font-family: system-ui, sans-serif; max-width: 640px; margin: 2rem auto; padding: 0 1rem; }
|
||||
h1 { margin-bottom: 0.25rem; }
|
||||
.subtitle { color: #666; font-size: 0.9rem; margin-top: 0; }
|
||||
.compose { display: flex; gap: 0.5rem; margin: 1.5rem 0; }
|
||||
.compose input { flex: 1; padding: 0.5rem; border: 1px solid #ccc; border-radius: 4px; font-size: 1rem; }
|
||||
.compose button { padding: 0.5rem 1rem; background: #0070f3; color: white; border: none; border-radius: 4px; cursor: pointer; font-size: 1rem; }
|
||||
.compose button:disabled { background: #aaa; cursor: default; }
|
||||
ul.messages { list-style: none; padding: 0; margin: 0; }
|
||||
ul.messages li { padding: 0.75rem; border-bottom: 1px solid #eee; display: flex; gap: 0.75rem; }
|
||||
.id { color: #999; font-size: 0.85rem; min-width: 2.5rem; }
|
||||
.empty { color: #999; }
|
||||
"#)
|
||||
}
|
||||
// DataStar SDK (loaded from CDN for the example)
|
||||
script
|
||||
type="module"
|
||||
src="https://cdn.jsdelivr.net/gh/starfederation/datastar@v1.0.0-beta.11/bundles/datastar.js"
|
||||
{}
|
||||
}
|
||||
body {
|
||||
div
|
||||
data-signals=(signals_json)
|
||||
data-on:load=(format!("@get('/sse?conn={conn_id}')"))
|
||||
{
|
||||
(inner)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.into_string()
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// HTTP handlers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn index_handler(State(state): State<(SharedState, RenderState)>) -> impl IntoResponse {
|
||||
let (app_state, rs) = &state;
|
||||
let conn_id = Uuid::new_v4();
|
||||
let signals_json = r#"{"newMsg":""}"#;
|
||||
app_state.connections.insert(conn_id, {
|
||||
// We need a sender for the connection, but the SSE handler will replace
|
||||
// it. Use a dummy channel that immediately drops.
|
||||
let (tx, _rx) = tokio::sync::mpsc::channel(1);
|
||||
tx
|
||||
});
|
||||
app_state
|
||||
.connections
|
||||
.update_signals(&conn_id, signals_json.to_string());
|
||||
|
||||
let mut cx =
|
||||
RenderContext::new(conn_id, livevue_rs::SignalStore::from_json(serde_json::json!({"newMsg": ""})), Default::default(), app_state.store.clone());
|
||||
|
||||
let inner = render_page(&mut cx, rs).await.unwrap_or_else(|e| {
|
||||
html! { p { "Render error: " (e) } }
|
||||
});
|
||||
|
||||
Html(page_shell(conn_id, signals_json, &inner))
|
||||
}
|
||||
|
||||
// SSE handler is provided by the framework; it needs SharedState only.
|
||||
// We wrap it to adapt the two-tuple state.
|
||||
async fn sse_adapter(
|
||||
State(state): State<(SharedState, RenderState)>,
|
||||
query: axum::extract::Query<livevue_rs::server::SseParams>,
|
||||
) -> impl IntoResponse {
|
||||
livevue_rs::sse_handler(State(state.0), query).await
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct PostMessageBody {
|
||||
#[serde(rename = "newMsg")]
|
||||
new_msg: Option<String>,
|
||||
}
|
||||
|
||||
async fn post_message_handler(
|
||||
State(state): State<(SharedState, RenderState)>,
|
||||
axum::extract::Query(params): axum::extract::Query<ActionParams>,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> impl IntoResponse {
|
||||
let signals = ingest_signals(&state.0, params.conn, body);
|
||||
let msg: String = signals
|
||||
.get("newMsg")
|
||||
.and_then(|v| v.as_str().map(|s| s.trim().to_string()))
|
||||
.unwrap_or_default();
|
||||
|
||||
if msg.is_empty() {
|
||||
return StatusCode::BAD_REQUEST;
|
||||
}
|
||||
|
||||
if let Err(e) = sqlx::query("INSERT INTO messages (body) VALUES ($1)")
|
||||
.bind(&msg)
|
||||
.execute(&state.1.pg)
|
||||
.await
|
||||
{
|
||||
tracing::error!("failed to insert message: {e}");
|
||||
return StatusCode::INTERNAL_SERVER_ERROR;
|
||||
}
|
||||
|
||||
// Note: we do NOT call store.publish() here — the PgReplicationListener
|
||||
// will detect the INSERT through the WAL stream and fire the invalidation
|
||||
// automatically. This demonstrates the key value of the module: mutations
|
||||
// from *any* source (other services, psql, migrations) are reflected
|
||||
// without any extra wiring.
|
||||
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
fn router(app_state: SharedState, render_state: RenderState) -> Router {
|
||||
let combined = (app_state, render_state);
|
||||
Router::new()
|
||||
.route("/", get(index_handler))
|
||||
.route("/sse", get(sse_adapter))
|
||||
.route("/action/post_message", post(post_message_handler))
|
||||
.with_state(combined)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// main
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::from_default_env()
|
||||
.add_directive("livevue_rs=debug".parse().unwrap())
|
||||
.add_directive("pg_replication=debug".parse().unwrap()),
|
||||
)
|
||||
.init();
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Database setup
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
let database_url = std::env::var("DATABASE_URL")
|
||||
.unwrap_or_else(|_| "postgres://postgres:postgres@localhost:5432/livevue".into());
|
||||
|
||||
let pg = sqlx::PgPool::connect(&database_url).await?;
|
||||
setup_db(&pg).await?;
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Framework store (SQLite in-memory, unused for queries here)
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
// The Store provides the broadcast channel and KV store that the framework
|
||||
// needs. The embedded SQLite pool is not used in this example — queries go
|
||||
// directly to PostgreSQL via the captured `pg` pool in the render closure.
|
||||
let sqlite = sqlx::SqlitePool::connect("sqlite::memory:").await?;
|
||||
let store = Store::new(sqlite);
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Render function
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
let render_state = RenderState { pg: pg.clone() };
|
||||
let rs_clone = render_state.clone();
|
||||
|
||||
let render_fn: livevue_rs::RenderFn = Arc::new(move |mut cx| {
|
||||
let rs = rs_clone.clone();
|
||||
Box::pin(async move {
|
||||
let signals_json = cx.signals.to_json_string();
|
||||
let conn_id = cx.connection_id;
|
||||
|
||||
let inner = render_page(&mut cx, &rs).await?;
|
||||
let keys = cx.take_subscription_keys();
|
||||
|
||||
// Wrap in the full page shell so DataStar can morph the entire document.
|
||||
let html = page_shell(conn_id, &signals_json, &inner);
|
||||
Ok((html, keys))
|
||||
})
|
||||
});
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// Application state + fanout
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
let state = AppState::new(store.clone(), render_fn);
|
||||
spawn_fanout(state.clone());
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// PG replication listener
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
// The replication protocol requires a dedicated connection.
|
||||
|
||||
let replication_store = store.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
// Rebuild the config each iteration — PgReplicationConfig is not
|
||||
// Clone, and start() takes ownership, so we reconstruct from the
|
||||
// env-var values cached in the outer scope.
|
||||
let c = PgReplicationConfig {
|
||||
host: std::env::var("PG_HOST").unwrap_or_else(|_| "localhost".into()),
|
||||
port: std::env::var("PG_PORT").ok()
|
||||
.and_then(|p| p.parse().ok()).unwrap_or(5432),
|
||||
user: std::env::var("PG_USER").unwrap_or_else(|_| "postgres".into()),
|
||||
password: std::env::var("PG_PASSWORD").unwrap_or_else(|_| "postgres".into()),
|
||||
database: std::env::var("PG_DATABASE").unwrap_or_else(|_| "livevue".into()),
|
||||
slot_name: "livevue_slot".into(),
|
||||
publication_name: "livevue_pub".into(),
|
||||
start_lsn: None,
|
||||
};
|
||||
if let Err(e) = PgReplicationListener::new(c, replication_store.clone()).start().await {
|
||||
tracing::error!("pg replication listener error: {e:#}");
|
||||
}
|
||||
tracing::info!("pg replication listener reconnecting in 5 s…");
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
}
|
||||
});
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// HTTP server
|
||||
// -----------------------------------------------------------------------
|
||||
|
||||
let app = router(state, render_state);
|
||||
let listener_tcp = tokio::net::TcpListener::bind("0.0.0.0:3001").await?;
|
||||
println!("pg_replication example listening on http://localhost:3001");
|
||||
axum::serve(listener_tcp, app).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user