maarcadetweet/crates/appview/src/firehose.rs

//! Connects `at-firehose::JetstreamConsumer` into the AppView indexer.
//!
//! The handler is intentionally tiny: it routes events by `kind`, delegates
//! all DB work to [`crate::indexer`], and feeds a small `mpsc` channel that
//! a background task drains to flush the cursor.
//!
//! Stats (events processed, last cursor seen, last-event timestamp for
//! connection health) live in an `Arc<Stats>` so both the consumer closure
//! and the HTTP `/healthz` handler can read them without locking.

use anyhow::Result;
use at_firehose::JetstreamEvent;
use serde_json::Value;
use sqlx::PgPool;
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use tracing::{debug, info, warn};

use crate::indexer;

/// Shared counters consumed by `/healthz`.
pub struct Stats {
    pub events_processed: AtomicU64,
    /// Microseconds since epoch of the most recent event seen.
    pub last_event_time_us: AtomicI64,
    /// Microseconds since epoch of the last persisted cursor (best-effort).
    pub last_cursor_persisted_us: AtomicI64,
    /// Whether the Jetstream WebSocket is currently up. The consumer
    /// writes this; `/healthz` reads it. Wrapped in `Arc` so the consumer
    /// can hold its own clone without borrowing from us.
    pub jetstream_connected: Arc<AtomicBool>,
    /// Whether the **local PDS** firehose WebSocket is currently up
    /// ([`crate::pds_firehose`]). Separate from `jetstream_connected`
    /// because the two streams fail independently and for different
    /// reasons: a dead Jetstream means no view of the wider network, a
    /// dead PDS firehose means the AppView has lost the guaranteed
    /// delivery path for its *own* users' records and is running on the
    /// best-effort push alone. `/healthz` has to be able to say which.
    pub pds_connected: AtomicBool,
    /// Number of `#commit` frames applied from the PDS firehose.
    pub pds_frames_processed: AtomicU64,
    /// Highest `seq` applied from the PDS firehose in this process.
    pub pds_last_seq: AtomicI64,
}

impl Default for Stats {
    fn default() -> Self {
        Self {
            events_processed: AtomicU64::new(0),
            last_event_time_us: AtomicI64::new(0),
            last_cursor_persisted_us: AtomicI64::new(0),
            jetstream_connected: Arc::new(AtomicBool::new(false)),
            pds_connected: AtomicBool::new(false),
            pds_frames_processed: AtomicU64::new(0),
            pds_last_seq: AtomicI64::new(0),
        }
    }
}

impl std::fmt::Debug for Stats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Stats")
            .field("events_processed", &self.events_processed())
            .field("last_event_time_us", &self.last_event_time_us.load(Ordering::Relaxed))
            .field("jetstream_connected", &self.jetstream_connected())
            .field("pds_connected", &self.pds_connected())
            .field("pds_frames_processed", &self.pds_frames_processed())
            .field("pds_last_seq", &self.pds_last_seq())
            .finish()
    }
}

impl Stats {
    pub fn new() -> Arc<Self> {
        Arc::new(Self::default())
    }

    /// Cheap clone of the `connected` flag (what the consumer stores into).
    pub fn jetstream_connected_arc(&self) -> Arc<AtomicBool> {
        self.jetstream_connected.clone()
    }

    /// Lag in milliseconds (last event time − local wall clock). When the
    /// local clock is ahead (very common — Jetstream events usually look
    /// "in the past" by a few seconds because the spec is `time_us` from
    /// the producer), this returns 0.
    pub fn lag_ms(&self) -> i64 {
        let last = self.last_event_time_us.load(Ordering::Relaxed);
        if last == 0 {
            return 0;
        }
        let now_us = chrono::Utc::now().timestamp_micros();
        let lag_us = (now_us - last).max(0);
        lag_us / 1000
    }

    pub fn events_processed(&self) -> u64 {
        self.events_processed.load(Ordering::Relaxed)
    }

    pub fn jetstream_connected(&self) -> bool {
        self.jetstream_connected.load(Ordering::Relaxed)
    }

    /// Is the local PDS firehose connected right now?
    pub fn pds_connected(&self) -> bool {
        self.pds_connected.load(Ordering::Relaxed)
    }

    pub fn pds_frames_processed(&self) -> u64 {
        self.pds_frames_processed.load(Ordering::Relaxed)
    }

    /// Highest PDS-firehose `seq` this process has applied. 0 before the
    /// first frame — note this is the *in-process* high-water mark, not
    /// the persisted cursor, which lives in `pds_firehose_cursor` and
    /// survives restarts.
    pub fn pds_last_seq(&self) -> i64 {
        self.pds_last_seq.load(Ordering::Relaxed)
    }
}

/// The thing the Jetstream consumer calls once per event.
#[derive(Clone)]
pub struct IndexHandler {
    pub db: PgPool,
    pub stats: Arc<Stats>,
    /// Sender side of the cursor-flush channel. The consumer closure pushes
    /// `event.time_us` every 100 events; a background task drains it and
    /// writes the running maximum to `jetstream_cursor` at most once every
    /// 500ms. The bound of 32 is plenty — the background task keeps up.
    cursor_tx: mpsc::Sender<i64>,
}

impl IndexHandler {
    pub fn new(db: PgPool, stats: Arc<Stats>, cursor_tx: mpsc::Sender<i64>) -> Self {
        Self { db, stats, cursor_tx }
    }

    pub async fn handle(&self, ev: JetstreamEvent) -> Result<()> {
        // Update per-event stats first so /healthz reflects liveness even
        // when DB writes are slow.
        self.stats.events_processed.fetch_add(1, Ordering::Relaxed);
        let prev = self
            .stats
            .last_event_time_us
            .fetch_max(ev.time_us, Ordering::Relaxed);
        if ev.time_us < prev {
            // Out-of-order event: keep the larger value but still try to
            // process. (Jetstream delivers near-monotonically but it's
            // not guaranteed.)
            self.stats
                .last_event_time_us
                .store(prev, Ordering::Relaxed);
        }

        let ok: bool = match ev.kind.as_str() {
            "commit" => match indexer::apply_commit(&self.db, &ev).await {
                Ok(true) => true,
                Ok(false) => {
                    debug!(kind = "commit", "apply_commit skipped event (unrecognized sub-collection)");
                    true // cursor can still advance — we "saw" the event
                }
                Err(e) => {
                    warn!(error = %e, kind = "commit", "apply_commit failed; NOT advancing cursor — expect reconnect replay");
                    false
                }
            },
            "identity" => {
                if let Err(e) = handle_identity(&self.db, &ev).await {
                    warn!(error = %e, did = %ev.did, "handle_identity failed");
                    return Ok(()); // don't advance cursor; let next replay retry
                }
                true
            }
            "account" => {
                if let Err(e) = handle_account(&self.db, &ev).await {
                    warn!(error = %e, did = %ev.did, "handle_account failed");
                    return Ok(()); // don't advance cursor; let next replay retry
                }
                true
            }
            other => {
                debug!(kind = %other, "ignoring event of unknown kind");
                true
            }
        };

        // ONLY advance the cursor when the event was processed or skipped
        // legitimately. Failures leave the cursor pinned so reconnect replays
        // the event.
        if !ok {
            return Ok(());
        }

        // Forward the cursor only every 100 events to avoid hammering the DB.
        if self.stats.events_processed() % 100 == 0 {
            let _ = self.cursor_tx.try_send(ev.time_us);
        }

        Ok(())
    }
}

/// `identity` event — Jetstream tells us a DID's handle changed.
///
/// The Jetstream payload includes `identity.handle` (the *current*
/// handle, since the event fires after every handle change) and
/// optionally `identity.did` (the DID — redundant with the outer
/// `ev.did` but we accept both). We pull the handle out and run it
/// through `indexer::backfill_handle` so every existing post row for
/// that DID gets the new value. The COALESCE guard inside
/// `indexer::PostRow::from_record` keeps empty strings from
/// clobbering this backfilled value when a later `commit` event
/// arrives.
async fn handle_identity(db: &PgPool, ev: &JetstreamEvent) -> Result<()> {
    let handle = extract_handle(&ev.identity).or_else(|| extract_handle(&ev.account));
    let Some(handle) = handle else {
        // Some identity events carry only a DID-doc rotation signal
        // with no handle payload — those are uninteresting for our
        // purpose. Advance the cursor anyway.
        debug!(did = %ev.did, "identity event without a usable handle payload");
        return Ok(());
    };
    let rows = indexer::backfill_handle(db, &ev.did, &handle).await?;
    info!(
        did = %ev.did,
        handle = %handle,
        rows_updated = rows,
        "backfilled handle on posts"
    );
    Ok(())
}

/// `account` event — Jetstream tells us an account's active/deactive
/// status changed. We mirror the handle-backfill behaviour in case
/// the `account` payload carries the verified handle alongside
/// `active`; many real-world identities show the handle there even
/// when no `identity` event was emitted.
async fn handle_account(db: &PgPool, ev: &JetstreamEvent) -> Result<()> {
    let Some(handle) = extract_handle(&ev.account) else {
        return Ok(());
    };
    let rows = indexer::backfill_handle(db, &ev.did, &handle).await?;
    info!(
        did = %ev.did,
        handle = %handle,
        rows_updated = rows,
        "backfilled handle on posts (account event)"
    );
    Ok(())
}

/// Pull a handle string out of a Jetstream event fragment. Returns
/// `None` if the fragment is absent or doesn't carry a usable
/// `handle` string field.
fn extract_handle(fragment: &Option<Value>) -> Option<String> {
    fragment
        .as_ref()
        .and_then(|v| v.get("handle"))
        .and_then(|v| v.as_str())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
}

/// Spawn the background task that drains the cursor-flush channel and
/// writes the running maximum to the DB. Returns when the receiver is
/// dropped (i.e. the main process is shutting down).
pub fn spawn_cursor_flush(
    db: PgPool,
    mut rx: mpsc::Receiver<i64>,
    stats: Arc<Stats>,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        let mut buf: Vec<i64> = Vec::with_capacity(128);
        let mut interval = tokio::time::interval(Duration::from_millis(500));
        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);

        loop {
            tokio::select! {
                biased;
                maybe = rx.recv() => {
                    match maybe {
                        Some(ts) => buf.push(ts),
                        None => {
                            // Channel closed — flush whatever's left and
                            // exit cleanly.
                            if !buf.is_empty() {
                                let max = buf.iter().copied().max().unwrap_or(0);
                                if let Err(e) = indexer::cursor_advance(&db, max).await {
                                    warn!(error = %e, "final cursor flush failed");
                                } else {
                                    stats.last_cursor_persisted_us.store(max, Ordering::Relaxed);
                                }
                                buf.clear();
                            }
                            return;
                        }
                    }
                }
                _ = interval.tick() => {
                    if buf.is_empty() { continue; }
                    let max = buf.iter().copied().max().unwrap_or(0);
                    if let Err(e) = indexer::cursor_advance(&db, max).await {
                        warn!(error = %e, "cursor flush failed");
                    } else {
                        stats.last_cursor_persisted_us.store(max, Ordering::Relaxed);
                    }
                    buf.clear();
                }
            }
        }
    })
}