maarcadetweet/crates/appview/src/handle_sync.rs

//! Background worker that resolves empty `handle` columns in the `posts`
//! table to real `@handle.bsky.social` style strings.
//!
//! ## Why
//!
//! Jetstream events carry no `handle` — only the `did`. The AppView's
//! indexer inserts posts with an empty placeholder (`handle = ''`) and a
//! separate worker is responsible for back-filling it. The UI can render
//! `@<did-prefix>…` as a fallback, but a real handle is far nicer and
//! makes the timeline readable for accounts that post anonymously.
//!
//! ## How
//!
//! `HandleSyncWorker::run_forever()` runs [`Self::run_once`] in a loop,
//! sleeping `interval_secs` between passes. Each pass:
//!
//! 1. Reads up to [`BATCH_SIZE`] distinct DIDs from `posts` where
//!    `handle = ''`.
//! 2. For each DID, dispatches by method:
//!      * `did:plc:` → [`HandleSyncWorker::plc_resolver`]
//!      * `did:web:` → [`HandleSyncWorker::web_resolver`]
//!      * anything else (e.g. `did:garbage:`) → skipped
//!    A resolver returning `Ok(None)` counts as `skipped`, not `failed`.
//!    Before any of these the local PDS is consulted, so a `did:key:`
//!    user on this PDS gets resolved without dialing plc.directory.
//! 3. `UPDATE posts SET handle = $1 WHERE did = $2 AND handle = ''` so
//!    concurrent syncs (or the `/internal/ingest-commit` path, which can
//!    populate handle separately) can't clobber a value written by
//!    someone else in the meantime.
//!
//! ## Testability
//!
//! The two resolvers are type-erased `Arc<dyn DidHandleResolver>`s so
//! tests can swap stubs that map DID → handle without hitting the
//! network. `PlcClient` and `WebResolver` are the production impls.

use anyhow::Result;
use at_identity::DidHandleResolver;
use sqlx::PgPool;
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, info, warn};

/// Max DIDs processed per pass. Keeps individual runs bounded so a
/// back-fill of thousands of empty-handle posts doesn't hammer the PLC.
pub const BATCH_SIZE: i64 = 100;

/// Max concurrent handle-resolve network calls per pass. Each
/// DID in a batch triggers a `POST /xrpc/com.atproto.identity
/// .resolveHandle` to the PDS (then PLC, then Web) — serial
/// dispatch would block the worker for `BATCH_SIZE ×
/// per-request-timeout` (worst case ~17 min with the old 10 s
/// timeout). Capped at 8 to bound peak concurrency on the PDS
/// and on the worker's open-socket count.
const DISPATCH_CONCURRENCY: usize = 8;

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct SyncReport {
    /// Rows whose `handle` column was newly populated this pass.
    pub resolved: usize,
    /// DIDs where the resolver returned `Err(_)` (network / 5xx).
    pub failed: usize,
    /// DIDs where the resolver returned `Ok(None)` (unknown DID,
    /// unsupported method) **or** rows that already had a non-empty
    /// handle when the UPDATE landed.
    pub skipped: usize,
}

pub struct HandleSyncWorker {
    pub db: PgPool,
    /// Local-PDS handle resolver. Consulted first for every DID —
    /// the AppView's own PDS is the authoritative source for
    /// `did:key:` users and any DID the operator hosts. A 404 from
    /// the PDS falls through to the public resolvers below.
    pub pds_resolver: Arc<dyn DidHandleResolver>,
    pub plc_resolver: Arc<dyn DidHandleResolver>,
    pub web_resolver: Arc<dyn DidHandleResolver>,
    pub interval_secs: u64,
}

impl HandleSyncWorker {
    /// Pick the right resolver based on the DID's method prefix and
    /// return its result. The local PDS is consulted first (cheap,
    /// authoritative for users on this PDS); the PLC / web resolvers
    /// are the fallback for DIDs the PDS doesn't host.
    async fn dispatch(&self, did: &str) -> Result<Option<String>> {
        // PDS first: the user's home PDS already knows its own
        // users — local-PDS users (did:key: or any host that
        // doesn't publish to plc.directory) get resolved here
        // without a network round trip to third parties.
        if let Some(h) = self.pds_resolver.resolve_handle(did).await? {
            if !h.is_empty() {
                return Ok(Some(h));
            }
        }
        if did.starts_with("did:plc:") {
            self.plc_resolver.resolve_handle(did).await
        } else if did.starts_with("did:web:") {
            self.web_resolver.resolve_handle(did).await
        } else {
            Ok(None)
        }
    }

    /// Drive [`Self::run_once`] on a fixed-interval loop until the
    /// process exits. Intended for `tokio::spawn`.
    pub async fn run_forever(self) {
        info!(
            interval_secs = self.interval_secs,
            "handle-sync worker started"
        );
        loop {
            match self.run_once().await {
                Ok(report) => {
                    if report.resolved > 0
                        || report.failed > 0
                        || report.skipped > 0
                    {
                        info!(
                            resolved = report.resolved,
                            failed = report.failed,
                            skipped = report.skipped,
                            "handle-sync pass complete"
                        );
                    } else {
                        debug!("handle-sync pass: nothing to do");
                    }
                }
                Err(e) => {
                    warn!(error = %e, "handle-sync pass aborted; will retry");
                }
            }
            tokio::time::sleep(Duration::from_secs(self.interval_secs)).await;
        }
    }

    /// One bounded scan: find up to [`BATCH_SIZE`] distinct DIDs whose
    /// posts have an empty handle, resolve them, and update the rows
    /// where the handle is still empty (race-safe).
    ///
    /// Unresolvable DIDs (e.g. `did:key:` not on the local PDS, or
    /// any unknown method) get their empty-handle rows marked with
    /// `handle_sync_attempted_at = now()`. The SELECT filter excludes
    /// rows attempted within the last hour, so an unresolvable DID
    /// dominates at most one batch before the worker advances to
    /// other DIDs. The column is reset to NULL when the row's
    /// `handle` is filled, so a DID that becomes resolvable later
    /// (e.g. the user joins the local PDS) gets re-attempted.
    pub async fn run_once(&self) -> Result<SyncReport> {
        let dids: Vec<(String,)> = sqlx::query_as(
            r#"SELECT DISTINCT did
               FROM posts
               WHERE handle = ''
                 AND (handle_sync_attempted_at IS NULL
                      OR handle_sync_attempted_at < now() - interval '1 hour')
               ORDER BY did
               LIMIT $1"#,
        )
        .bind(BATCH_SIZE)
        .fetch_all(&self.db)
        .await?;

        let mut report = SyncReport::default();
        if dids.is_empty() {
            return Ok(report);
        }

        // Dispatch in parallel — the PDS / PLC / Web resolvers are
        // independent network calls. Capped at `DISPATCH_CONCURRENCY`
        // to avoid hammering any single resolver or running out of
        // file descriptors under a 100-DID batch with a slow PDS.
        // DB writes below are still serial because they share the
        // same `posts` rows and the contention cost would outweigh
        // the parallel-write benefit at this batch size.
        use futures::stream::{self, StreamExt};
        let dispatch_results: Vec<(String, anyhow::Result<Option<String>>)> = stream::iter(dids)
            .map(|(did,)| async move {
                let r = self.dispatch(&did).await;
                (did, r)
            })
            .buffer_unordered(DISPATCH_CONCURRENCY)
            .collect()
            .await;

        for (did, result) in dispatch_results {
            match result {
                Ok(Some(handle)) => {
                    if handle.is_empty() {
                        report.skipped += 1;
                        mark_attempted(&self.db, &did).await?;
                        continue;
                    }
                    let res = sqlx::query(
                        "UPDATE posts SET handle = $1, \
                                handle_sync_attempted_at = NULL \
                         WHERE did = $2 AND handle = ''",
                    )
                    .bind(&handle)
                    .bind(&did)
                    .execute(&self.db)
                    .await?;
                    if res.rows_affected() > 0 {
                        report.resolved += res.rows_affected() as usize;
                    } else {
                        // Another worker / ingest path already filled it
                        // between our SELECT and UPDATE.
                        report.skipped += 1;
                    }
                }
                Ok(None) => {
                    report.skipped += 1;
                    mark_attempted(&self.db, &did).await?;
                }
                Err(e) => {
                    warn!(did = %did, error = %e, "handle resolve failed");
                    report.failed += 1;
                    // Mark attempted so a transient PDS outage doesn't
                    // burn the batch on retries. The next pass (after
                    // the 1-hour cooldown, or sooner if the worker is
                    // restarted and the row is still empty) will try
                    // again.
                    if let Err(e) = mark_attempted(&self.db, &did).await {
                        warn!(did = %did, error = %e, "handle_sync mark_attempted failed");
                    }
                }
            }
        }
        Ok(report)
    }
}

/// Stamp `handle_sync_attempted_at = now()` on every empty-handle
/// row for `did`. Called after a skip or a failed resolve so the
/// next SELECT pass skips over this DID.
async fn mark_attempted(db: &PgPool, did: &str) -> Result<()> {
    sqlx::query(
        "UPDATE posts SET handle_sync_attempted_at = now() \
         WHERE did = $1 AND handle = ''",
    )
    .bind(did)
    .execute(db)
    .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use std::collections::HashMap;
    use std::sync::Mutex;
    use std::time::Duration;
    use tokio::time::timeout;

    /// Stub resolver driven by a fixed DID → handle map. Counts how
    /// many times each DID was queried so the limit test can assert
    /// the worker capped the batch correctly.
    struct StubResolver {
        mapping: Mutex<HashMap<String, Option<String>>>,
        queried: Mutex<Vec<String>>,
    }

    impl StubResolver {
        fn new(mapping: HashMap<String, Option<String>>) -> Self {
            Self {
                mapping: Mutex::new(mapping),
                queried: Mutex::new(Vec::new()),
            }
        }
        fn into_arc(self) -> Arc<dyn DidHandleResolver> {
            Arc::new(self)
        }
    }

    #[async_trait]
    impl DidHandleResolver for StubResolver {
        async fn resolve_handle(&self, did: &str) -> Result<Option<String>> {
            self.queried.lock().unwrap().push(did.to_string());
            Ok(self.mapping.lock().unwrap().get(did).cloned().flatten())
        }
    }

    /// Convenience: a worker whose `plc_resolver` and `web_resolver`
    /// both point at the same stub. Existing tests don't care which
    /// method the DIDs use because the stub is method-agnostic.
    fn worker_with(db: PgPool, stub: Arc<dyn DidHandleResolver>) -> HandleSyncWorker {
        HandleSyncWorker {
            db,
            pds_resolver: Arc::clone(&stub),
            plc_resolver: Arc::clone(&stub),
            web_resolver: Arc::clone(&stub),
            interval_secs: 999,
        }
    }

    async fn try_test_db() -> Option<PgPool> {
        let url = std::env::var("DATABASE_URL_APPVIEW").ok()?;
        match timeout(Duration::from_secs(2), sqlx::PgPool::connect(&url)).await {
            Ok(Ok(pool)) => match sqlx::migrate!("../../migrations/appview")
                .run(&pool)
                .await
            {
                Ok(()) => Some(pool),
                Err(_) => None,
            },
            _ => None,
        }
    }

    fn unique_did(suffix: &str) -> String {
        format!("did:plc:stubsync_{}_{}", suffix, uuid::Uuid::new_v4().simple())
    }

    async fn seed_post(
        db: &PgPool,
        did: &str,
        rkey: &str,
        handle: &str,
    ) -> Result<()> {
        let uri = format!("at://{did}/app.twi.post/{rkey}");
        sqlx::query(
            r#"INSERT INTO posts
                 (uri, did, handle, rkey, collection, text, cid,
                  parent_uri, root_uri, langs, created_at)
               VALUES ($1,$2,$3,$4,'app.twi.post','stub','bafy',NULL,NULL,NULL, now())
               ON CONFLICT (uri) DO NOTHING"#,
        )
        .bind(&uri)
        .bind(did)
        .bind(handle)
        .bind(rkey)
        .execute(db)
        .await?;
        Ok(())
    }

    async fn get_handle(db: &PgPool, did: &str) -> Option<String> {
        sqlx::query_scalar::<_, String>(
            "SELECT handle FROM posts WHERE did = $1 ORDER BY indexed_at DESC LIMIT 1",
        )
        .bind(did)
        .fetch_optional(db)
        .await
        .ok()
        .flatten()
        .filter(|s| !s.is_empty())
    }

    /// Smoke test: a stub resolver maps one DID → handle. After
    /// `run_once()` the worker should populate the `handle` column on
    /// every empty post for that DID and report it as `resolved`.
    #[tokio::test]
    async fn run_once_returns_report() {
        let Some(db) = try_test_db().await else {
            eprintln!("appview DB unavailable; skipping");
            return;
        };
        let did = unique_did("report");
        // Wipe any previous stub rows for this slot.
        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await
            .unwrap();

        // Seed two posts for the same DID, both with empty handle.
        for rk in ["rka", "rkb"] {
            seed_post(&db, &did, rk, "").await.unwrap();
        }

        let expected_handle = format!("handle.{}", uuid::Uuid::new_v4().simple());
        let resolver = StubResolver::new(HashMap::from([(
            did.clone(),
            Some(expected_handle.clone()),
        )]))
        .into_arc();

        let worker = worker_with(db, resolver);
        let report = worker.run_once().await.unwrap();
        assert!(report.resolved >= 2, "expected ≥2 resolved, got {report:?}");
        assert_eq!(report.failed, 0);
        assert_eq!(report.skipped, 0);

        let h = get_handle(&worker.db, &did).await;
        assert_eq!(h.as_deref(), Some(expected_handle.as_str()));
    }

    /// DIDs that already have a non-empty handle on **every** post must
    /// not be re-queried — the worker scans `WHERE handle = ''`.
    #[tokio::test]
    async fn skip_dids_with_handle() {
        let Some(db) = try_test_db().await else {
            eprintln!("appview DB unavailable; skipping");
            return;
        };
        let did = unique_did("skip");
        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await
            .unwrap();

        // Insert one post with a pre-filled handle.
        seed_post(&db, &did, "rkA", "pre-existing.handle").await.unwrap();

        // The stub would return a different handle if asked.
        let resolver = StubResolver::new(HashMap::from([(
            did.clone(),
            Some("different.handle".into()),
        )]))
        .into_arc();

        let worker = worker_with(db, resolver);
        let report = worker.run_once().await.unwrap();
        assert_eq!(report.resolved, 0);
        assert_eq!(report.failed, 0);
        assert_eq!(report.skipped, 0); // DID was already filtered out by the SELECT

        let h = get_handle(&worker.db, &did).await;
        assert_eq!(h.as_deref(), Some("pre-existing.handle"));
    }

    /// When `run_once()` finds more than [`BATCH_SIZE`] empty-handle DIDs,
    /// only the first batch is processed this pass; the rest stay empty
    /// for the next pass.
    #[tokio::test]
    async fn respects_batch_limit() {
        let Some(db) = try_test_db().await else {
            eprintln!("appview DB unavailable; skipping");
            return;
        };

        // Seed BATCH_SIZE + 5 distinct DIDs, all empty handles. We
        // can't seed 100+000 for real so we use the BATCH_SIZE bound
        // directly. The same code path runs with any row count.
        let prefix = unique_did("batch");
        let mut all_dids = Vec::new();
        // Insert a separate DIDs table-style marker so we can clean
        // them all up afterwards without touching other test data.
        for i in 0..(BATCH_SIZE as usize + 5) {
            let did = format!("{prefix}_{i}");
            all_dids.push(did.clone());
            seed_post(&db, &did, "rk", "").await.unwrap();
        }

        let mut mapping = HashMap::new();
        for did in &all_dids {
            mapping.insert(did.clone(), Some(format!("h.{}", &did[did.len()-6..])));
        }
        let resolver = StubResolver::new(mapping).into_arc();

        let worker = HandleSyncWorker {
            db: db.clone(),
            pds_resolver: Arc::clone(&resolver),
            plc_resolver: Arc::clone(&resolver),
            web_resolver: Arc::clone(&resolver),
            interval_secs: 999,
        };
        let report = worker.run_once().await.unwrap();
        // Exactly BATCH_SIZE rows updated (one post per DID).
        assert_eq!(
            report.resolved as i64,
            BATCH_SIZE,
            "resolved should equal batch size: {report:?}"
        );
        assert_eq!(report.failed, 0);

        // The remaining 5 DIDs must still have empty handles.
        let remaining: i64 =
            sqlx::query_scalar("SELECT COUNT(*) FROM posts WHERE handle = '' AND did LIKE $1")
                .bind(format!("{prefix}_%"))
                .fetch_one(&db)
                .await
                .unwrap();
        assert_eq!(
            remaining, 5,
            "expected 5 unresolved DIDs left, got {remaining}"
        );

        // Cleanup so repeated test runs stay hygienic.
        let _ = sqlx::query("DELETE FROM posts WHERE did LIKE $1")
            .bind(format!("{prefix}_%"))
            .execute(&db)
            .await;
    }

    /// The `UPDATE … WHERE handle = ''` clause must guard against races
    /// with another writer: if /internal/ingest-commit fills the
    /// handle between our SELECT and UPDATE, our UPDATE is a no-op.
    #[tokio::test]
    async fn update_does_not_overwrite_concurrent_write() {
        let Some(db) = try_test_db().await else {
            eprintln!("appview DB unavailable; skipping");
            return;
        };
        let did = unique_did("race");
        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await
            .unwrap();

        seed_post(&db, &did, "rk", "").await.unwrap();
        let resolver = StubResolver::new(HashMap::from([(
            did.clone(),
            Some("from-resolver".into()),
        )]))
        .into_arc();

        let worker = HandleSyncWorker {
            db: db.clone(),
            pds_resolver: Arc::clone(&resolver),
            plc_resolver: Arc::clone(&resolver),
            web_resolver: Arc::clone(&resolver),
            interval_secs: 999,
        };

        // Concurrent writer: race with the worker's UPDATE by setting
        // handle directly while run_once is reading it.
        // In practice the SELECT happens first, so the UPDATE WHERE
        // clause is what protects us. Simulate the "other writer won"
        // outcome directly here: write a handle, then call run_once —
        // since the SELECT excludes non-empty rows, run_once sees an
        // empty batch.
        sqlx::query("UPDATE posts SET handle = 'from-ingest' WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await
            .unwrap();

        let report = worker.run_once().await.unwrap();
        assert_eq!(report.resolved, 0, "must not touch already-handled rows");
        let h = get_handle(&worker.db, &did).await;
        assert_eq!(h.as_deref(), Some("from-ingest"));
    }

    /// Dispatch test: a `did:web:` DID must be routed to the
    /// `web_resolver` (not the PLC one). Without this routing, every
    /// `did:web:` post would stay `@<did-prefix>…` forever.
    #[tokio::test]
    async fn dispatches_did_web_to_web_resolver() {
        let Some(db) = try_test_db().await else {
            eprintln!("appview DB unavailable; skipping");
            return;
        };
        let did = format!("did:web:example.com:user:{}", uuid::Uuid::new_v4().simple());
        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await
            .unwrap();
        seed_post(&db, &did, "rk", "").await.unwrap();

        // Two stubs that disagree on the answer. The dispatcher's
        // job is to pick the right one based on the DID method.
        let plc = StubResolver::new(HashMap::from([(
            did.clone(),
            Some("WRONG-PLC-HANDLE".into()),
        )]))
        .into_arc();
        let web = StubResolver::new(HashMap::from([(
            did.clone(),
            Some("web-handle.example.com".into()),
        )]))
        .into_arc();

        let worker = HandleSyncWorker {
            db: db.clone(),
            pds_resolver: Arc::clone(&plc),
            plc_resolver: plc,
            web_resolver: web,
            interval_secs: 999,
        };
        let report = worker.run_once().await.unwrap();
        assert_eq!(
            report.resolved, 1,
            "did:web must resolve through the web resolver, got {report:?}"
        );
        assert_eq!(
            report.failed, 0,
            "did:web must not route to the PLC resolver"
        );
        let h = get_handle(&worker.db, &did).await;
        assert_eq!(h.as_deref(), Some("web-handle.example.com"));

        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await;
    }

    /// DIDs whose method isn't `did:plc:` or `did:web:` (e.g. `did:key:`)
    /// are silently skipped — neither resolver is consulted.
    #[tokio::test]
    async fn unknown_methods_are_skipped() {
        let Some(db) = try_test_db().await else {
            eprintln!("appview DB unavailable; skipping");
            return;
        };
        let did = format!("did:key:z{}", uuid::Uuid::new_v4().simple());
        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await
            .unwrap();
        seed_post(&db, &did, "rk", "").await.unwrap();

        // Both stubs would happily return a handle if asked. The
        // dispatcher's prefix check must prevent that — neither
        // resolver should ever see a `did:key:` DID. We share the
        // query logs via `Arc<Mutex<...>>` so the test can read them
        // back after the worker has run.
        let plc_log: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let web_log: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let plc = TrackingResolver::new(
            HashMap::from([(did.clone(), Some("plc-handle".into()))]),
            Arc::clone(&plc_log),
        );
        let web = TrackingResolver::new(
            HashMap::from([(did.clone(), Some("web-handle".into()))]),
            Arc::clone(&web_log),
        );
        let plc_arc: Arc<dyn DidHandleResolver> = Arc::new(plc);
        let web_arc: Arc<dyn DidHandleResolver> = Arc::new(web);

        let worker = HandleSyncWorker {
            db: db.clone(),
            pds_resolver: Arc::clone(&plc_arc),
            plc_resolver: plc_arc,
            web_resolver: web_arc,
            interval_secs: 999,
        };
        let report = worker.run_once().await.unwrap();
        assert_eq!(report.resolved, 0, "did:key must not resolve, got {report:?}");
        assert_eq!(
            report.failed, 0,
            "did:key must not be treated as a failure"
        );
        assert_eq!(report.skipped, 1, "did:key must be skipped");

        // No resolver was consulted.
        assert!(
            plc_log.lock().unwrap().is_empty(),
            "PLC resolver must not be called for did:key"
        );
        assert!(
            web_log.lock().unwrap().is_empty(),
            "web resolver must not be called for did:key"
        );

        let _ = sqlx::query("DELETE FROM posts WHERE did = $1")
            .bind(&did)
            .execute(&db)
            .await;
    }

    /// Stub resolver that records every DID it's queried into a
    /// shared log so tests can verify dispatch routing. Distinct
    /// from `StubResolver`, which owns its log and would need
    /// downcasting to read it back through an `Arc<dyn _>`.
    struct TrackingResolver {
        mapping: HashMap<String, Option<String>>,
        queried: Arc<Mutex<Vec<String>>>,
    }

    impl TrackingResolver {
        fn new(
            mapping: HashMap<String, Option<String>>,
            queried: Arc<Mutex<Vec<String>>>,
        ) -> Self {
            Self { mapping, queried }
        }
    }

    #[async_trait]
    impl DidHandleResolver for TrackingResolver {
        async fn resolve_handle(&self, did: &str) -> Result<Option<String>> {
            self.queried.lock().unwrap().push(did.to_string());
            Ok(self.mapping.get(did).cloned().flatten())
        }
    }
}