use std::{
sync::{
Arc,
atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering},
},
time::Duration,
};
use sqlx::{PgPool, postgres::PgListener};
use tokio::sync::broadcast;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
pub const PG_CHANNEL: &str = "bus_events";
#[derive(Clone, Debug)]
pub struct BusEvent(pub serde_json::Value);
impl BusEvent {
pub fn kind(&self) -> &str {
self.0.get("kind").and_then(|v| v.as_str()).unwrap_or("")
}
fn uuid_field(&self, key: &str) -> Option<Uuid> {
self.0
.get(key)
.and_then(|v| v.as_str())
.and_then(|s| Uuid::parse_str(s).ok())
}
pub fn team_id(&self) -> Option<Uuid> {
self.uuid_field("team_id")
}
pub fn recipient_agent_id(&self) -> Option<Uuid> {
self.uuid_field("recipient_agent_id")
}
pub fn sender_agent_id(&self) -> Option<Uuid> {
self.uuid_field("sender_agent_id")
}
pub fn channel_id(&self) -> Option<Uuid> {
self.uuid_field("channel_id")
}
pub fn message_id(&self) -> Option<i64> {
self.0.get("id").and_then(|v| v.as_i64())
}
pub fn recipient_session(&self) -> Option<&str> {
self.0.get("recipient_session").and_then(|v| v.as_str())
}
pub fn sender_session(&self) -> Option<&str> {
self.0.get("sender_session").and_then(|v| v.as_str())
}
pub fn is_announcement(&self) -> bool {
self.0
.get("announce")
.and_then(|v| v.as_bool())
.unwrap_or(false)
}
pub fn is_direct_message(&self) -> bool {
self.kind() == "message" && self.recipient_agent_id().is_some()
}
pub fn visible_to(&self, team_id: Uuid, agent_id: Uuid, session: &str) -> bool {
if self.team_id() != Some(team_id) {
return false;
}
if self.is_direct_message() {
if self.sender_agent_id() == Some(agent_id) {
return true;
}
if self.recipient_agent_id() != Some(agent_id) {
return false;
}
return match self.recipient_session() {
Some(addressed) => addressed == session,
None => true,
};
}
true
}
}
pub const DEFAULT_PING_SECS: u64 = 30;
const MISSED_ECHOES: u32 = 3;
pub struct ListenerHealth {
ping_every: Duration,
attached: AtomicBool,
last_echo: AtomicI64,
attachments: AtomicU64,
echoes: AtomicU64,
}
fn unix_now() -> i64 {
chrono::Utc::now().timestamp()
}
impl ListenerHealth {
fn new(ping_every: Duration) -> Self {
Self {
ping_every,
attached: AtomicBool::new(false),
last_echo: AtomicI64::new(0),
attachments: AtomicU64::new(0),
echoes: AtomicU64::new(0),
}
}
pub fn ping_every(&self) -> Duration {
self.ping_every
}
pub fn stale_after(&self) -> Duration {
self.ping_every * MISSED_ECHOES
}
fn attached_now(&self) {
self.last_echo.store(0, Ordering::SeqCst);
self.attached.store(true, Ordering::SeqCst);
self.attachments.fetch_add(1, Ordering::SeqCst);
}
fn echoed(&self) {
self.last_echo.store(unix_now(), Ordering::SeqCst);
self.echoes.fetch_add(1, Ordering::SeqCst);
}
fn detached(&self) {
self.attached.store(false, Ordering::SeqCst);
}
pub fn report(&self) -> serde_json::Value {
let attached = self.attached.load(Ordering::SeqCst);
let last_echo = self.last_echo.load(Ordering::SeqCst);
let age = (last_echo > 0).then(|| unix_now() - last_echo);
let listener = if !attached {
"detached"
} else if age.is_some_and(|a| a <= self.stale_after().as_secs() as i64) {
"live"
} else {
"silent"
};
serde_json::json!({
"listener": listener,
"last_echo_seconds": age.filter(|_| attached),
"attachments": self.attachments.load(Ordering::SeqCst),
"echoes": self.echoes.load(Ordering::SeqCst),
"ping_seconds": self.ping_every.as_secs(),
})
}
pub fn is_live(&self) -> bool {
self.report()["listener"] == "live"
}
}
#[derive(Clone)]
pub struct EventHub {
tx: broadcast::Sender<BusEvent>,
health: Arc<ListenerHealth>,
}
impl Default for EventHub {
fn default() -> Self {
Self::new()
}
}
impl EventHub {
pub fn new() -> Self {
Self::with_ping(Duration::from_secs(DEFAULT_PING_SECS))
}
pub fn with_ping(ping_every: Duration) -> Self {
let (tx, _) = broadcast::channel(256);
Self {
tx,
health: Arc::new(ListenerHealth::new(ping_every.max(Duration::from_secs(1)))),
}
}
pub fn listener(&self) -> &ListenerHealth {
&self.health
}
pub fn subscribe(&self) -> broadcast::Receiver<BusEvent> {
self.tx.subscribe()
}
pub fn publish(&self, event: BusEvent) {
let _ = self.tx.send(event);
}
}
pub async fn run_pg_listener(pool: PgPool, hub: EventHub, ct: CancellationToken) {
let health = hub.listener();
let replica = Uuid::new_v4().to_string();
loop {
if ct.is_cancelled() {
return;
}
match PgListener::connect_with(&pool).await {
Ok(mut listener) => {
if let Err(e) = listener.listen(PG_CHANNEL).await {
tracing::warn!(error = %e, "LISTEN failed; retrying");
} else {
tracing::info!("event listener attached to '{PG_CHANNEL}'");
health.attached_now();
let mut ticker = tokio::time::interval(health.ping_every());
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut unanswered: u32 = 0;
loop {
tokio::select! {
_ = ct.cancelled() => return,
_ = ticker.tick() => {
if unanswered >= MISSED_ECHOES {
tracing::warn!(
unanswered,
"event listener has not heard its own ping; the \
connection is dead however open it looks. Reattaching"
);
break;
}
unanswered += 1;
let ping = serde_json::json!({ "kind": "ping", "replica": replica })
.to_string();
if let Err(e) = sqlx::query("SELECT pg_notify($1, $2)")
.bind(PG_CHANNEL)
.bind(&ping)
.execute(&pool)
.await
{
tracing::warn!(error = %e, "could not ping the event channel");
}
}
recv = listener.try_recv() => match recv {
Ok(Some(notification)) => {
match serde_json::from_str(notification.payload()) {
Ok(value) => {
let event = BusEvent(value);
if event.kind() == "ping" {
if event.0.get("replica").and_then(|v| v.as_str())
== Some(replica.as_str())
{
unanswered = 0;
health.echoed();
}
continue;
}
hub.publish(event)
}
Err(e) => tracing::warn!(
error = %e,
payload = notification.payload(),
"unparseable bus event"
),
}
}
Ok(None) => tracing::debug!("event listener reconnected"),
Err(e) => {
tracing::warn!(error = %e, "event listener error");
break;
}
}
}
}
health.detached();
}
}
Err(e) => {
tracing::warn!(error = %e, "could not attach event listener; retrying");
}
}
tokio::select! {
_ = ct.cancelled() => return,
_ = tokio::time::sleep(std::time::Duration::from_secs(2)) => {}
}
}
}