diff --git a/src/ingest/relay/context.rs b/src/ingest/relay/context.rs index 073e081..a5644db 100644 --- a/src/ingest/relay/context.rs +++ b/src/ingest/relay/context.rs @@ -45,6 +45,10 @@ pub(crate) struct WorkerContext<'a> { pub(crate) http: PublicHttpClient, pub(crate) error_counts: HashMap>, pub(crate) wrong_host_authority: HashMap>, + /// the repo row as [`Self::load_repo_state`] read it, `None` when it made the row up + pub(crate) loaded: Option>, + /// the repo row this message writes, put on the batch by [`Self::stage_repo_row`] + pub(crate) staged: Option>, } impl WorkerContext<'_> { @@ -56,6 +60,38 @@ impl WorkerContext<'_> { is_suppressed } + pub(crate) fn stage_repo_state(&mut self, repo_state: &RepoState<'_>) { + self.staged = Some(repo_state.clone().into_static()); + } + + /// puts the row this message staged on `txn`. the relay holds no lock while it resolves a + /// doc or probes a pds, so a backfill or the indexer may have written the row since it was + /// read, and only what this message changed goes onto the row as it is now + pub(crate) fn stage_repo_row(&mut self, txn: &mut db::Txn<'_>, did: &Did) -> Result<()> { + let loaded = self.loaded.take(); + let Some(staged) = self.staged.take() else { + return Ok(()); + }; + let key = keys::repo_key(did); + let row = match loaded { + #[cfg(feature = "indexer")] + Some(before) => { + txn.hold_repo_write_lock(did); + let Some(bytes) = self.state.db.repos.get(&key).into_diagnostic()? else { + // erased while this message was in flight + return Ok(()); + }; + staged.rebased(&before, db::deser_repo_state(&bytes)?.into_static()) + } + // a row load_repo_state made up has nothing to rebase onto, and without the + // indexer there are no record locks to take + _ => staged, + }; + txn.batch + .insert(&self.state.db.repos, key, db::ser_repo_state(&row)?); + Ok(()) + } + pub(crate) fn reset_error(&mut self, host: &str) { if let Some(count) = self.error_counts.get_mut(&util::hash(&host)) { *count = 0; @@ -135,7 +171,6 @@ impl WorkerContext<'_> { pub(crate) fn refresh_doc(&mut self, did: &Did, repo_state: &mut RepoState) -> Result<()> { let refresh_started = StatsInstant::now(); let result = (|| { - let db = &self.state.db; self.state.resolver.invalidate_sync(did); let resolve_started = StatsInstant::now(); let doc = Handle::current() @@ -145,12 +180,7 @@ impl WorkerContext<'_> { let doc = doc?; repo_state.update_from_doc(doc); repo_state.touch(); - - self.batch.insert( - &db.repos, - keys::repo_key(did), - db::ser_repo_state(repo_state)?, - ); + self.stage_repo_state(repo_state); Ok(()) })(); self.stats.record_refresh_doc(refresh_started.elapsed()); @@ -361,6 +391,7 @@ impl WorkerContext<'_> { &mut self, msg: &WorkerMessage, ) -> Result>> { + self.loaded = None; let db = &self.state.db; let did = msg.msg.did().expect("we checked if valid"); let repo_key = keys::repo_key(did); @@ -392,6 +423,7 @@ impl WorkerContext<'_> { self.stats.record_repo_state_outcome( crate::ingest::firehose_stats::RepoStateLoadOutcome::Hit, ); + self.loaded = Some(repo_state.clone()); return Ok(Some(repo_state)); } @@ -495,11 +527,7 @@ impl WorkerContext<'_> { RelayWorker::pds_host(repo_state.pds.as_deref()).as_deref(), ); - self.batch.insert( - &db.repos, - &repo_key, - crate::db::ser_repo_state(&repo_state)?, - ); + self.stage_repo_state(&repo_state); self.sink.new_repo(did.clone()); diff --git a/src/ingest/relay/handlers.rs b/src/ingest/relay/handlers.rs index 818579d..87a8955 100644 --- a/src/ingest/relay/handlers.rs +++ b/src/ingest/relay/handlers.rs @@ -4,7 +4,6 @@ use tokio::runtime::Handle; use tracing::{debug, warn}; use url::Url; -use crate::db::{self, keys}; use crate::types::{RepoState, RepoStatus}; use crate::ingest::stream::{Account, AccountStatus, Commit, Identity, Sync}; @@ -40,8 +39,6 @@ impl RelayWorker { debug!(broken = ?chain_break, "chain break, forwarding anyway"); } - let repo_key = keys::repo_key(&commit.repo); - ctx.sink.commit( ctx.state, &mut ctx.batch, @@ -53,11 +50,7 @@ impl RelayWorker { repo_state.root = Some(commit_obj.into()); repo_state.touch(); - ctx.batch.insert( - &ctx.state.db.repos, - repo_key, - db::ser_repo_state(repo_state)?, - ); + ctx.stage_repo_state(repo_state); Ok(()) } @@ -78,17 +71,11 @@ impl RelayWorker { return Ok(()); }; - let repo_key = keys::repo_key(&sync.did); - ctx.sink.sync(ctx.state, &mut ctx.batch, firehose, sync)?; repo_state.root = Some(validated.commit_obj.into()); repo_state.touch(); - ctx.batch.insert( - &ctx.state.db.repos, - repo_key, - db::ser_repo_state(repo_state)?, - ); + ctx.stage_repo_state(repo_state); Ok(()) } @@ -146,16 +133,10 @@ impl RelayWorker { Self::pds_host(repo_state.pds.as_deref()).as_deref(), ); - let repo_key = keys::repo_key(&identity.did); - ctx.sink .identity(ctx.state, &mut ctx.batch, firehose, identity, changed, true)?; - ctx.batch.insert( - &ctx.state.db.repos, - repo_key, - db::ser_repo_state(repo_state)?, - ); + ctx.stage_repo_state(repo_state); Ok(()) } @@ -212,8 +193,6 @@ impl RelayWorker { Self::pds_host(repo_state.pds.as_deref()).as_deref(), ); - let repo_key = keys::repo_key(&account.did); - ctx.sink.account( ctx.state, &mut ctx.batch, @@ -225,11 +204,7 @@ impl RelayWorker { )?; repo_state.touch(); - ctx.batch.insert( - &ctx.state.db.repos, - repo_key, - db::ser_repo_state(repo_state)?, - ); + ctx.stage_repo_state(repo_state); Ok(()) } diff --git a/src/ingest/relay/worker.rs b/src/ingest/relay/worker.rs index dba0d3c..2153348 100644 --- a/src/ingest/relay/worker.rs +++ b/src/ingest/relay/worker.rs @@ -119,6 +119,8 @@ impl RelayWorker { http, error_counts: Default::default(), wrong_host_authority: Default::default(), + loaded: None, + staged: None, }; while let Some(msg) = rx.blocking_recv() { @@ -171,7 +173,11 @@ impl RelayWorker { let batch = std::mem::replace(&mut ctx.batch, ctx.state.db.inner.batch()); let count_deltas = std::mem::take(&mut ctx.count_deltas); - let txn = crate::db::Txn::from_parts(&ctx.state.db, batch, count_deltas); + let mut txn = crate::db::Txn::from_parts(&state.db, batch, count_deltas); + if let Err(e) = ctx.stage_repo_row(&mut txn, &did) { + shard_stats.record_process_error(); + error!(did = %did, err = %e, "relay shard: failed to stage repo state"); + } let (staged, commit_timings) = match ctx.sink.commit_txn(&state, txn) { Ok(staged) => staged, Err(e) => { @@ -323,3 +329,228 @@ impl RelayWorker { } } } + +#[cfg(all(test, feature = "indexer"))] +mod tests { + use std::time::Duration; + + use axum::Router; + use bytes::Bytes; + use jacquard_common::IntoStatic; + use jacquard_common::types::did::Did; + use jacquard_common::types::tid::Tid; + + use super::*; + use crate::db::{self, keys}; + use crate::ingest::indexer::{IndexerMessage, IndexerTx}; + use crate::ingest::mailbox::{ShardedReceiver, ShardedSender}; + use crate::ingest::stream::{Datetime, Identity}; + use crate::types::RepoState; + + const DID: &str = "did:plc:aaaaaaaaaaaaaaaaaaaaaaaa"; + const DOC_PDS: &str = "https://pds.example/"; + + /// a plc stub answering every path with `DID`'s doc, but only once `gate` is free, + /// so a test can hold the relay inside `resolve_doc` + async fn held_doc_stub( + gate: Arc>, + arrived: Arc, + ) -> u16 { + let doc = serde_json::json!({ + "@context": ["https://www.w3.org/ns/did/v1"], + "id": DID, + "service": [{ + "id": "#atproto_pds", + "type": "AtprotoPersonalDataServer", + "serviceEndpoint": DOC_PDS, + }], + }) + .to_string(); + let app = Router::new().fallback(move || { + let doc = doc.clone(); + let gate = gate.clone(); + let arrived = arrived.clone(); + async move { + arrived.notify_one(); + drop(gate.lock().await); + ( + [(axum::http::header::CONTENT_TYPE, "application/json")], + doc, + ) + } + }); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + port + } + + /// a relay worker for one repo whose doc lookups wait for [`HeldRelay::release`] + struct HeldRelay { + state: Arc, + did: Did, + buffer_tx: BufferTx, + forwarded: ShardedReceiver, + held: Option>, + arrived: Arc, + _tmp: tempfile::TempDir, + } + + impl HeldRelay { + async fn start(row: &RepoState<'_>) -> miette::Result { + let tmp = tempfile::tempdir().into_diagnostic()?; + let plc = Url::parse("http://public.example/").into_diagnostic()?; + let mut state = AppState::new(&crate::config::Config { + database_path: tmp.path().to_path_buf(), + plc_urls: vec![plc.clone()], + identity_cache_size: 64, + ..Default::default() + })?; + let gate = Arc::new(tokio::sync::Mutex::new(())); + let held = gate.clone().lock_owned().await; + let arrived = Arc::new(tokio::sync::Notify::new()); + let port = held_doc_stub(gate, arrived.clone()).await; + let client = reqwest::Client::builder() + .no_proxy() + .resolve( + "public.example", + std::net::SocketAddr::from((std::net::Ipv4Addr::LOCALHOST, port)), + ) + .build() + .into_diagnostic()?; + state.resolver = crate::resolver::Resolver::with_plc_urls_and_client( + vec![plc], + 64, + PublicHttpClient::test(client), + ); + let state = Arc::new(state); + + let (indexer_tx, mut forwarded) = ShardedSender::channel(1); + let (buffer_tx, worker) = RelayWorker::new( + state.clone(), + SinkSeed::new(IndexerTx { inner: indexer_tx }), + false, + 1, + ValidationOptions { + rev_clock_skew_secs: 300, + verify_mst: false, + verify_cids: false, + }, + ); + let handle = Handle::current(); + std::thread::spawn(move || worker.run(handle)); + + let relay = Self { + state, + did: Did::new_static(DID).into_diagnostic()?, + buffer_tx, + forwarded: forwarded.pop().expect("one indexer shard"), + held: Some(held), + arrived, + _tmp: tmp, + }; + relay.put_repo(row)?; + Ok(relay) + } + + /// sends an identity event with no handle, so the relay goes to the doc for one, + /// and comes back once it's waiting there + async fn identity_waiting_on_the_doc(&self) -> miette::Result<()> { + self.buffer_tx + .send(IngestMessage::Firehose { + url: Url::parse("wss://relay.example").into_diagnostic()?, + is_pds: false, + msg: SubscribeReposMessage::Identity(Box::new(Identity { + did: self.did.clone(), + handle: None, + seq: 1, + time: Datetime(chrono::Utc::now().into()), + })), + }) + .await + .into_diagnostic()?; + self.arrived.notified().await; + Ok(()) + } + + /// lets the doc through, then waits until the relay has committed and forwarded the event + async fn release(&mut self) -> miette::Result<()> { + drop(self.held.take()); + tokio::time::timeout(Duration::from_secs(10), self.forwarded.recv()) + .await + .into_diagnostic()? + .expect("the relay forwards the identity event"); + Ok(()) + } + + fn put_repo(&self, row: &RepoState<'_>) -> miette::Result<()> { + let mut batch = self.state.db.inner.batch(); + batch.insert( + &self.state.db.repos, + keys::repo_key(&self.did), + db::ser_repo_state(row)?, + ); + batch.commit().into_diagnostic() + } + + fn repo(&self) -> miette::Result>> { + self.state + .db + .repos + .get(keys::repo_key(&self.did)) + .into_diagnostic()? + .map(|bytes| Ok(db::deser_repo_state(&bytes)?.into_static())) + .transpose() + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn an_identity_event_keeps_a_root_written_while_it_resolved() -> miette::Result<()> { + let mut relay = HeldRelay::start(&RepoState::backfilling()).await?; + relay.identity_waiting_on_the_doc().await?; + + // a backfill's persist lands, under the repo lock, while the relay waits on the doc + let root = crate::types::Commit { + version: 3, + rev: crate::db::types::DbTid::from(&Tid::new("3jzfcijpj2z2a").into_diagnostic()?), + data: jacquard_repo::mst::util::compute_cid(b"backfilled-root").into_diagnostic()?, + prev: None, + sig: Bytes::from_static(b"signature"), + }; + let mut backfilled = RepoState::backfilling(); + backfilled.root = Some(root.clone()); + let mut txn = crate::db::Txn::new(&relay.state.db); + txn.hold_repo_write_lock(&relay.did); + txn.batch.insert( + &relay.state.db.repos, + keys::repo_key(&relay.did), + db::ser_repo_state(&backfilled)?, + ); + txn.commit()?; + + relay.release().await?; + let after = relay.repo()?.expect("repo row"); + assert_eq!(after.pds.as_deref(), Some(DOC_PDS)); + assert_eq!( + after.root, + Some(root), + "the identity event rolled back the backfill's root" + ); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn an_identity_event_does_not_bring_back_a_repo_erased_while_it_resolved() + -> miette::Result<()> { + let mut relay = HeldRelay::start(&RepoState::backfilling()).await?; + relay.identity_waiting_on_the_doc().await?; + + let mut batch = relay.state.db.inner.batch(); + batch.remove(&relay.state.db.repos, keys::repo_key(&relay.did)); + batch.commit().into_diagnostic()?; + + relay.release().await?; + assert!(relay.repo()?.is_none()); + Ok(()) + } +}