use std::collections::HashSet;
use std::str::FromStr;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use anyhow::{Context, Result};
use async_trait::async_trait;
use sqlx::postgres::{PgConnectOptions, PgListener, PgPool, PgPoolOptions};
use tokio::sync::broadcast;
use super::trait_::{
CursorGoneError, EngineEventBus, Event, EventFilter, NewEvent, PruneOpts, Subsystem,
};
const LOCAL_BROADCAST_CAPACITY: usize = 1024;
fn channel_name(ns: &str) -> String {
format!("assay_engine_events_{ns}")
}
pub struct PgEngineEventBus {
pool: PgPool,
local: broadcast::Sender<Arc<Event>>,
listener_pool: PgPool,
}
impl PgEngineEventBus {
pub async fn new(pool: PgPool, db_url: &str) -> Result<Self> {
let listener_opts =
PgConnectOptions::from_str(db_url).context("invalid db_url for listener")?;
let listener_pool = PgPoolOptions::new()
.max_connections(1)
.connect_with(listener_opts)
.await
.context("listener pool connect")?;
let (local, _) = broadcast::channel(LOCAL_BROADCAST_CAPACITY);
Ok(Self {
pool,
local,
listener_pool,
})
}
fn ensure_listener(&self, namespace: &str) {
static REGISTRY: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
let reg = REGISTRY.get_or_init(|| Mutex::new(HashSet::new()));
let ch = channel_name(namespace);
{
let mut g = reg.lock().unwrap();
if !g.insert(ch.clone()) {
return;
}
}
let pool = self.pool.clone();
let listener_pool = self.listener_pool.clone();
let local = self.local.clone();
tokio::spawn(async move {
loop {
match PgListener::connect_with(&listener_pool).await {
Ok(mut listener) => {
if let Err(e) = listener.listen(&ch).await {
tracing::warn!(?e, channel = %ch, "LISTEN failed; retrying in 2s");
tokio::time::sleep(Duration::from_secs(2)).await;
continue;
}
tracing::info!(channel = %ch, "engine_events LISTEN bridge online");
loop {
match listener.recv().await {
Ok(n) => {
let id: i64 = match n.payload().parse() {
Ok(v) => v,
Err(e) => {
tracing::warn!(?e, payload = n.payload(), "bad NOTIFY payload");
continue;
}
};
match sqlx::query_as::<_, PgEventRow>(
"SELECT id, ts, namespace, subsystem, kind, payload
FROM engine.events WHERE id = $1",
)
.bind(id)
.fetch_optional(&pool)
.await
{
Ok(Some(row)) => {
let _ = local.send(Arc::new(row.into_event()));
}
Ok(None) => {
tracing::debug!(id, "NOTIFY id not in table (pruned?)");
}
Err(e) => {
tracing::warn!(?e, id, "engine_events SELECT after NOTIFY failed");
}
}
}
Err(e) => {
tracing::warn!(?e, channel = %ch, "LISTEN recv error; reconnecting");
break;
}
}
}
}
Err(e) => {
tracing::warn!(?e, "PgListener connect failed; retrying in 5s");
tokio::time::sleep(Duration::from_secs(5)).await;
}
}
}
});
}
async fn oldest_id_inner(&self, namespace: &str) -> Result<Option<i64>> {
let row: Option<(i64,)> = sqlx::query_as(
"SELECT id FROM engine.events WHERE namespace = $1 ORDER BY id ASC LIMIT 1",
)
.bind(namespace)
.fetch_optional(&self.pool)
.await
.context("oldest_id query")?;
Ok(row.map(|r| r.0))
}
}
#[derive(sqlx::FromRow)]
struct PgEventRow {
id: i64,
ts: f64,
namespace: String,
subsystem: String,
kind: String,
payload: serde_json::Value,
}
impl PgEventRow {
fn into_event(self) -> Event {
let ss = Subsystem::from_str(&self.subsystem);
Event {
id: self.id,
ts: self.ts,
namespace: self.namespace,
subsystem: ss,
kind: self.kind,
payload: self.payload,
}
}
}
#[async_trait]
impl EngineEventBus for PgEngineEventBus {
async fn publish_committed(&self, ev: NewEvent<'_>) -> Result<i64> {
let mut tx = self.pool.begin().await?;
let row: (i64, f64) = sqlx::query_as(
"INSERT INTO engine.events (namespace, subsystem, kind, payload)
VALUES ($1, $2, $3, $4)
RETURNING id, ts",
)
.bind(ev.namespace)
.bind(ev.subsystem.as_str())
.bind(ev.kind)
.bind(&ev.payload)
.fetch_one(&mut *tx)
.await?;
let ch = channel_name(ev.namespace);
sqlx::query("SELECT pg_notify($1, $2)")
.bind(&ch)
.bind(row.0.to_string())
.execute(&mut *tx)
.await?;
tx.commit().await?;
let event = Arc::new(Event {
id: row.0,
ts: row.1,
namespace: ev.namespace.to_string(),
subsystem: ev.subsystem,
kind: ev.kind.to_string(),
payload: ev.payload,
});
let _ = self.local.send(event);
Ok(row.0)
}
async fn read_since(
&self,
namespace: &str,
after: Option<i64>,
filter: &EventFilter,
limit: u32,
) -> std::result::Result<Vec<Event>, CursorGoneError> {
if let Some(a) = after
&& let Ok(Some(oldest)) = self.oldest_id_inner(namespace).await
&& a < oldest - 1
{
return Err(CursorGoneError { after: a, oldest });
}
let rows: Vec<PgEventRow> = sqlx::query_as(
"SELECT id, ts, namespace, subsystem, kind, payload
FROM engine.events
WHERE namespace = $1 AND ($2::bigint IS NULL OR id > $2)
ORDER BY id ASC
LIMIT $3",
)
.bind(namespace)
.bind(after)
.bind(limit as i64)
.fetch_all(&self.pool)
.await
.map_err(|e| {
tracing::error!(?e, "read_since DB error");
CursorGoneError {
after: after.unwrap_or(0),
oldest: 0,
}
})?;
Ok(rows
.into_iter()
.map(PgEventRow::into_event)
.filter(|e| filter.matches(e))
.collect())
}
fn subscribe(&self, namespace: &str) -> broadcast::Receiver<Arc<Event>> {
self.ensure_listener(namespace);
self.local.subscribe()
}
async fn prune(&self, before_ts: f64) -> Result<u64> {
let n = sqlx::query("DELETE FROM engine.events WHERE ts < $1")
.bind(before_ts)
.execute(&self.pool)
.await
.context("engine_events prune (global)")?
.rows_affected();
Ok(n)
}
async fn prune_with(&self, opts: PruneOpts) -> Result<u64> {
let n = match &opts.namespace {
Some(ns) => sqlx::query(
"DELETE FROM engine.events WHERE namespace = $1 AND ts < $2",
)
.bind(ns)
.bind(opts.before_ts)
.execute(&self.pool)
.await
.context("engine_events prune_with (scoped)")?
.rows_affected(),
None => sqlx::query("DELETE FROM engine.events WHERE ts < $1")
.bind(opts.before_ts)
.execute(&self.pool)
.await
.context("engine_events prune_with (global)")?
.rows_affected(),
};
Ok(n)
}
async fn oldest_id(&self, namespace: &str) -> Result<Option<i64>> {
self.oldest_id_inner(namespace).await
}
}
#[cfg(test)]
#[path = "pg_test.rs"]
mod pg_test;