use futures::{FutureExt, TryFutureExt};
use rand::Rng;
use super::*;
use std::collections::HashSet;
use crate::ingest::indexer::message::{BulkMessage, EraseRequest, RepoOp};
use crate::ingest::indexer::repos::{on_owners, repo_ops};
impl ReposControl {
/// repositories hydrant still has work queued for, a backfill or resync
/// entry, in the same did order as [`ReposControl::iter`]
pub fn iter_queued(&self, cursor: Option<&Did>) -> impl Iterator<Item = Result<RepoInfo>> {
use std::ops::Bound;
let (pending_start, resync_start) = match cursor {
Some(cursor) => {
let repo_key = keys::repo_key(cursor);
// past every entry the cursor's own did has
let last = keys::pending_did_key(&repo_key, &keys::pending_key(u64::MAX));
(Bound::Excluded(last), Bound::Excluded(repo_key))
}
None => (Bound::Unbounded, Bound::Unbounded),
};
let indexer = &self.0.db.indexer;
let pending = indexer
.pending_did_range((pending_start, Bound::Unbounded))
.map(|guard| {
let key = guard.key().into_diagnostic()?;
keys::split_pending_did_key(&key)
.map(|(repo_key, _)| repo_key.to_vec())
.ok_or_else(|| miette::miette!("invalid pending_dids key {key:?}"))
});
let resync = indexer
.resync
.range((resync_start, Bound::Unbounded))
.map(|guard| Ok(guard.key().into_diagnostic()?.to_vec()));
let state = self.0.clone();
merge_sorted(pending, resync)
.filter_map(move |repo_key| queued_row(&state, repo_key).transpose())
}
/// request one or more repositories to be resynced.
///
/// note that they may not immediately start backfilling if:
/// - other repos already filled the backfill concurrency limit,
/// - or there are many repos pending already.
///
/// this will also clear any error state the repo may have been in,
/// allowing it to resync again.
pub async fn resync(&self, dids: impl IntoIterator<Item = Did>) -> Result<Vec<Did>> {
let queued = self.on_owner_shards(RepoOp::Resync, dids).await;
self.0.notify_backfill();
queued
}
/// explicitly track one or more repositories, enqueuing them for backfill if needed.
///
/// - if a repo is new, a fresh [`RepoState`] is created and backfill is queued.
/// - if a repo is already known but untracked, it is marked tracked and re-enqueued.
/// - if a repo has a row but was never queued, the row stays and backfill is queued.
/// - if a repo is already tracked, this is a no-op.
pub async fn track(&self, dids: impl IntoIterator<Item = Did>) -> Result<Vec<Did>> {
let queued = self.on_owner_shards(RepoOp::Track, dids).await;
self.0.notify_backfill();
queued
}
/// stop tracking one or more repositories. hydrant will stop processing new events
/// for them and remove them from the pending/resync queues, but existing indexed
/// records are **not** deleted. a repo whose backfill is writing its records right now is
/// left alone and missing from the returned list.
pub async fn untrack(&self, dids: impl IntoIterator<Item = Did>) -> Result<Vec<Did>> {
self.on_owner_shards(RepoOp::Untrack, dids).await
}
/// runs `op` for every one of `dids` on its owner shard, see [`on_owners`], and returns the
/// ones it changed in the order they came. every shard is answered before the one persist,
/// so it runs with nothing held. a shard that fails or is gone fails the call, but what the
/// others did stays
async fn on_owner_shards(
&self,
op: RepoOp,
dids: impl IntoIterator<Item = Did>,
) -> Result<Vec<Did>> {
let dids: Vec<Did> = dids.into_iter().collect();
let answers = on_owners(
&self.0,
dids.clone(),
|did| did,
move |state, shard, batch| repo_ops(state, shard.as_deref(), op, &batch),
)
.await;
let mut changed = HashSet::new();
let mut failed = None;
for answer in answers {
match answer.and_then(|ran| ran) {
Ok(dids) => changed.extend(dids),
Err(e) => failed = Some(e),
}
}
self.0.db.run(|db| db.persist()).await?;
if let Some(e) = failed {
return Err(e);
}
Ok(dids.into_iter().filter(|did| changed.remove(did)).collect())
}
/// stop ingesting a repository and forget everything stored for it, so
/// tracking it again later backfills it from nothing. `exclude` also adds
/// it to the filter excludes, which keeps it from being indexed again
/// until it's taken back out. its indexer shard does the erasing, a step at a
/// time between its live events, and this answers once the last step is in.
pub async fn erase_repo(
&self,
did: &Did,
exclude: bool,
) -> Result<crate::types::DeleteBodiesReport> {
self.0.ensure_operator_body_deletion_supported()?;
let indexer =
self.0.indexer.lock().clone().ok_or_else(|| {
miette::miette!("the indexer isn't running, so nothing can erase")
})?;
let (reply, report) = tokio::sync::oneshot::channel();
let erase = EraseRequest {
did: did.clone(),
exclude,
reply,
};
indexer
.bulk
.send(BulkMessage::Erase(erase))
.await
.map_err(|_| miette::miette!("the indexer stopped before the erase started"))?;
let report = report
.await
.map_err(|_| miette::miette!("the indexer failed to erase {did}, try again"))?;
// not under anything, so the shard doesn't wait on the fsync
self.0.db.run(|db| db.persist()).await?;
Ok(report)
}
}
impl RepoHandle {
/// erase selected body copies for one record while preserving its CID and
/// the repository's authoritative topology.
pub async fn delete_record_bodies(
&self,
collection: &str,
rkey: &str,
target: crate::types::DeleteBodyTarget,
) -> Result<crate::types::DeleteBodiesReport> {
self.state.ensure_operator_body_deletion_supported()?;
Nsid::new(collection).into_diagnostic()?;
Rkey::new(rkey).into_diagnostic()?;
let did = self.did.clone();
let collection = collection.to_owned();
let rkey = DbRkey::new(rkey);
let report = super::write_repo(&self.state, &self.did, move |db| {
crate::db::redact_record_bodies(db, &did, &collection, &rkey, target)
})
.await?;
// after the shard answers, so neither it nor the repo's writers wait on the fsync
self.state.db.run(|db| db.persist()).await?;
Ok(report)
}
/// gets a record from this repository.
pub async fn get_record(&self, collection: &str, rkey: &str) -> Result<Option<Record>> {
if !self.state.stores_record_bodies() {
miette::bail!("block storage is not available in this mode");
}
let did = self.did.clone();
let db_key = keys::record_key(&did, collection, &DbRkey::new(rkey));
self.state
.db
.run(move |db| {
use miette::WrapErr;
let Some(body) = db.indexer.head(db_key).into_diagnostic()? else {
return Ok(None);
};
if crate::db::is_cid_record_value(&body) {
return Ok(None);
}
let value = serde_ipld_dagcbor::from_slice::<Data>(&body)
.into_diagnostic()
.wrap_err("cant parse record body")?;
// the cid is no longer stored; recompute it from the body
let cid = crate::car::compute_cid(body.as_ref());
let cid = Cid::ipld(cid);
Ok(Some(Record { did, cid, value }))
})
.await
}
/// lists records from this repository.
pub async fn list_records(
&self,
collection: &str,
limit: usize,
reverse: bool,
cursor: Option<&str>,
) -> Result<RecordList> {
if !self.state.stores_record_bodies() {
miette::bail!("block storage is not available in this mode");
}
let did = self.did.clone().into_static();
let prefix = keys::record_prefix_collection(&did, collection);
let cursor = cursor.map(|c| c.to_smolstr());
self.state
.db
.run(move |db| {
let mut results = Vec::new();
let mut next_cursor = None;
let mut last_returned = None;
let iter: Box<dyn Iterator<Item = _>> = if !reverse {
// newest-first: walk backwards from the cursor, or from the
// end of the collection. the exclusive upper bound is the
// cursor key itself, or `prefix ff`: rkeys cannot contain
// 0xff, so every key in the collection sorts below it
let end_key = if let Some(cursor) = &cursor {
let mut k = prefix.clone();
k.extend_from_slice(DbRkey::new(cursor).to_smolstr().as_bytes());
k
} else {
let mut k = prefix.clone();
k.push(0xff);
k
};
Box::new(
db.indexer
.head_range(prefix.as_slice()..end_key.as_slice())
.rev(),
)
} else {
let start_key = if let Some(cursor) = &cursor {
// strictly after the cursor rkey: its key plus the next
// possible byte
let mut k = prefix.clone();
k.extend_from_slice(DbRkey::new(cursor).to_smolstr().as_bytes());
k.push(0);
k
} else {
prefix.clone()
};
Box::new(db.indexer.head_range(start_key.as_slice()..))
};
for item in iter {
let (key, body) = item.into_inner().into_diagnostic()?;
if !key.starts_with(prefix.as_slice()) {
break;
}
if crate::db::is_cid_record_value(&body) {
continue;
}
let rkey = keys::parse_rkey_text(&key[prefix.len()..])?;
if results.len() >= limit {
next_cursor = last_returned;
break;
}
let value: Data = serde_ipld_dagcbor::from_slice(&body).unwrap_or(Data::Null);
// the cid is no longer stored; recompute it from the body
let cid = crate::car::compute_cid(body.as_ref());
let cid = Cid::ipld(cid);
results.push(ListedRecord {
rkey: Rkey::new(rkey.to_smolstr()).expect("that rkey is validated"),
cid,
value,
});
last_returned = Some(rkey);
}
Ok((results, next_cursor))
})
.await
.map(|(records, next_cursor)| RecordList {
records,
cursor: next_cursor
.map(|rkey| Rkey::new(rkey.to_smolstr()).expect("that rkey is validated")),
})
}
/// generates a streaming CAR v1 response body for this repository.
///
/// returns `None` if the repo has no commit yet (still backfilling) or is an
/// unmigrated repo that does not have the necessary data to reconstruct the
/// root commit from.
///
/// ## notes
/// - calling this if you are using collection allowlist will always result
/// in an error since the commit root won't match the reconstructed CID.
/// - calling this for big repositories will incur more resource cost due to
/// hydrant's structure, the whole MST is always reconstructed.
pub async fn generate_car(
&self,
) -> Result<Option<impl futures::Stream<Item = std::io::Result<bytes::Bytes>> + Send + 'static>>
{
if self.state.ephemeral {
miette::bail!("CAR generation is not supported in ephemeral mode");
}
if !self.state.stores_record_bodies() {
miette::bail!("block storage is not available in this mode");
}
use crate::car::{CarHeader, CarWriter};
// a backfill writes its records in steps, so until it's done they can be a mix of old
// and new that no root matches
let info = self.info().await?;
if info.as_ref().is_some_and(|info| info.pending) {
return Ok(None);
}
let commit = match self.state().await? {
Some(state) => match state.root {
Some(c) => c,
None => return Ok(None),
},
None => return Ok(None),
};
let atp_commit = match commit.into_atp_commit(self.did.clone().into_static()) {
Some(c) => c,
None => return Ok(None),
};
let commit_cid = atp_commit.to_cid()?;
let commit_cbor = atp_commit.to_cbor()?;
let did = self.did.clone().into_static();
let app_state = self.state.clone();
// build mst and collect leaf blocks in a single blocking pass
let (computed_root, mst_nodes, leaf_blocks) =
tokio::task::spawn_blocking(move || -> Result<_> {
let prefix = keys::record_prefix_did(&did);
let mut leaves = Vec::new();
let mut leaf_blocks = Vec::new();
for guard in app_state.db.indexer.head_prefix(&prefix) {
let (key, body) = guard.into_inner().into_diagnostic()?;
let (collection, rkey) = keys::split_record_suffix(&key[prefix.len()..])?;
let mst_key = format!("{collection}/{rkey}");
if crate::db::is_cid_record_value(&body) {
miette::bail!("record body was deleted for {}/{mst_key}", did);
}
// the cid is no longer stored; recompute it from the body
let ipld_cid = crate::car::compute_cid(body.as_ref());
leaves.push((mst_key, ipld_cid));
leaf_blocks.push((ipld_cid, bytes::Bytes::copy_from_slice(body.as_ref())));
}
let (root, nodes) = crate::mst::build_tree(leaves)?;
Ok((root, nodes, leaf_blocks))
})
.await
.into_diagnostic()??;
// sanity check: rebuilt root should match stored commit data in full-index mode
if computed_root != atp_commit.data {
// a backfill that failed partway left its steps, so until the retry the records
// can be a mix that no root matches
if info.is_some_and(|info| info.resync.is_some()) {
return Ok(None);
}
miette::bail!(
"mst root mismatch for {}: computed {computed_root} != stored {}",
self.did,
atp_commit.data,
);
}
// stream the car directly to the response
let (reader, writer) = tokio::io::duplex(64 * 1024);
tokio::spawn(
async move {
let header = CarHeader::new_v1(vec![commit_cid]);
let mut car_writer = CarWriter::new(header, writer);
// write commit first, then mst nodes + leaf blocks
car_writer
.write(commit_cid, &commit_cbor)
.await
.into_diagnostic()?;
for (cid, bytes) in mst_nodes {
car_writer.write(cid, &bytes).await.into_diagnostic()?;
}
for (cid, bytes) in leaf_blocks {
car_writer.write(cid, &bytes).await.into_diagnostic()?;
}
car_writer.finish().await.into_diagnostic()?;
Result::<_, miette::Report>::Ok(())
}
.inspect_err(|e| tracing::error!("can't generate car: {e}")),
);
Ok(Some(tokio_util::io::ReaderStream::new(reader)))
}
/// gets how many records of a collection this repository has.
pub async fn count_records(&self, collection: &str) -> Result<u64> {
let did = self.did.clone().into_static();
let collection = collection.to_string();
self.state
.db
.run(move |db| db::get_record_count(db, &did, &collection))
.await
}
}
/// the union of two ascending key streams, each key once
fn merge_sorted(
a: impl Iterator<Item = Result<Vec<u8>>>,
b: impl Iterator<Item = Result<Vec<u8>>>,
) -> impl Iterator<Item = Result<Vec<u8>>> {
let (mut a, mut b) = (a.peekable(), b.peekable());
let mut last: Option<Vec<u8>> = None;
std::iter::from_fn(move || {
loop {
let take_a = match (a.peek(), b.peek()) {
(None, None) => return None,
(Some(Err(_)), _) | (Some(_), None) => true,
(_, Some(Err(_))) | (None, Some(_)) => false,
(Some(Ok(x)), Some(Ok(y))) => x <= y,
};
let key = match if take_a { a.next() } else { b.next() }? {
Ok(key) => key,
Err(e) => return Some(Err(e)),
};
if last.as_ref() != Some(&key) {
last = Some(key.clone());
return Some(Ok(key));
}
}
})
}
/// the row for a repo a queue names, or `None` if the repo is gone or its
/// did-order entry outlived the real one
fn queued_row(state: &AppState, repo_key: Result<Vec<u8>>) -> Result<Option<RepoInfo>> {
let repo_key = repo_key?;
let Some(bytes) = state.db.repos.get(&repo_key).into_diagnostic()? else {
return Ok(None);
};
let did = TrimmedDid::try_from(repo_key.as_slice())?.to_did();
let repo_state = crate::db::deser_repo_state(&bytes)?.into_static();
let info = repo_row(state, did, repo_state)?;
Ok((info.pending || info.resync.is_some()).then_some(info))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ingest::indexer::message::shard_of;
fn record_body(text: &str) -> Vec<u8> {
serde_ipld_dagcbor::to_vec(&serde_json::json!({
"$type": "app.bsky.feed.post",
"text": text,
}))
.unwrap()
}
#[tokio::test]
async fn get_and_list_hide_redacted_heads() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let mut config = crate::config::Config::default();
config.database_path = tmp.path().to_path_buf();
let state = Arc::new(AppState::new(&config)?);
let did = Did::new_static("did:plc:testredactedreads").into_diagnostic()?;
let visible = record_body("visible");
let hidden_cid = crate::car::compute_cid(&record_body("hidden"));
let mut batch = state.db.inner.batch();
for rkey in ["a", "c"] {
state.db.indexer.stage_head(
&mut batch,
keys::record_key(&did, "app.bsky.feed.post", &DbRkey::new(rkey)),
&visible,
);
}
state.db.indexer.stage_head(
&mut batch,
keys::record_key(&did, "app.bsky.feed.post", &DbRkey::new("b")),
hidden_cid.to_bytes(),
);
batch.commit().into_diagnostic()?;
let handle = ReposControl(state).get(&did);
let visible_record = handle
.get_record("app.bsky.feed.post", "a")
.await?
.expect("visible record");
let visible_json = serde_json::to_value(&visible_record.value).into_diagnostic()?;
assert_eq!(visible_json["$type"], "app.bsky.feed.post");
assert_eq!(visible_json["text"], "visible");
assert!(
handle
.get_record("app.bsky.feed.post", "b")
.await?
.is_none()
);
let listed = handle
.list_records("app.bsky.feed.post", 10, true, None)
.await?;
assert_eq!(
listed
.records
.iter()
.map(|record| record.rkey.as_str())
.collect::<Vec<_>>(),
["a", "c"]
);
let first = handle
.list_records("app.bsky.feed.post", 1, true, None)
.await?;
assert_eq!(first.cursor.as_ref().map(|rkey| rkey.as_str()), Some("a"));
let second = handle
.list_records(
"app.bsky.feed.post",
1,
true,
first.cursor.as_ref().map(|rkey| rkey.as_str()),
)
.await?;
assert_eq!(second.records[0].rkey.as_str(), "c");
Ok(())
}
#[tokio::test]
async fn erase_repo_forgets_everything_stored_for_the_repo() -> miette::Result<()> {
// one entry per step, and as many as an erase step takes
for chunk_records in [1, 0] {
erases_everything_stored_for_the_repo(chunk_records).await?;
}
Ok(())
}
async fn erases_everything_stored_for_the_repo(chunk_records: usize) -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let mut config = crate::config::Config::default();
config.database_path = tmp.path().to_path_buf();
config.backfill_chunk_records = chunk_records;
let state = Arc::new(AppState::new(&config)?);
let _shard = crate::ingest::indexer::FirehoseWorker::serve_one_shard(state.clone());
let did = Did::new_static("did:plc:testoperatorrepo").into_diagnostic()?;
let metadata = crate::types::RepoMetadata::backfilling(7);
let head = record_body("head");
let old = record_body("old");
#[cfg(feature = "indexer_stream")]
let archived = record_body("archived");
#[cfg(feature = "indexer_stream")]
let archived_cid = crate::car::compute_cid(&archived);
let record_key =
keys::record_key(&did, "app.bsky.feed.post", &DbRkey::new("3kzbif5moe22m"));
let mut batch = state.db.inner.batch();
batch.insert(
&state.db.repos,
keys::repo_key(&did),
crate::db::ser_repo_state(&crate::types::RepoState::synced())?,
);
batch.insert(
&state.db.repo_metadata,
keys::repo_metadata_key(&did),
crate::db::ser_repo_meta(&metadata)?,
);
state
.db
.indexer
.stage_pending_insert(&mut batch, &did, &keys::pending_key(7));
batch.insert(&state.db.crawler, keys::crawler_retry_key(&did), []);
state.db.indexer.stage_head(&mut batch, &record_key, &head);
state.db.indexer.stage_history(
&mut batch,
keys::history_key(
&record_key,
&crate::db::types::DbTid::new_from_bytes(1_u64.to_be_bytes()),
),
&old,
);
#[cfg(feature = "indexer_stream")]
batch.insert(
&state.db.stream.event_bodies,
keys::event_body_key(&record_key, &archived_cid),
&archived,
);
batch.insert(
&state.db.indexer.resync_buffer,
keys::resync_buffer_key(
&did,
crate::db::types::DbTid::new_from_bytes(2_u64.to_be_bytes()),
),
b"buffered body bytes",
);
crate::db::set_record_count(&mut batch, &state.db, &did, "app.bsky.feed.post", 1);
batch.commit().into_diagnostic()?;
let repos = ReposControl(state.clone());
// the answer comes once the last step is in, not before
let answered = Arc::new(std::sync::atomic::AtomicBool::new(false));
let steps = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let (seen_answered, counted) = (answered.clone(), steps.clone());
*state.after_job_step.lock() = Some(Arc::new(move |step| {
counted.store(step, std::sync::atomic::Ordering::SeqCst);
assert!(!seen_answered.load(std::sync::atomic::Ordering::SeqCst));
}));
let report = repos.erase_repo(&did, false).await?;
answered.store(true, std::sync::atomic::Ordering::SeqCst);
*state.after_job_step.lock() = None;
let steps = steps.load(std::sync::atomic::Ordering::SeqCst);
// one per kind of entry and the repo's own, since every kind here fits in one step
assert!(steps >= 5, "{steps} steps");
assert_eq!(report.heads_redacted, 1);
assert_eq!(report.history_redacted, 1);
assert_eq!(report.buffered_commits_deleted, 1);
let gone = |ks: &fjall::Keyspace, prefix: Vec<u8>| ks.prefix(prefix).next().is_none();
let db = &state.db;
assert!(gone(db.repos.raw(), keys::repo_key(&did)));
assert!(gone(db.repo_metadata.raw(), keys::repo_metadata_key(&did)));
assert!(
!db.indexer
.is_pending(&keys::pending_key(7))
.into_diagnostic()?
);
db.indexer.assert_pending_dids_match();
assert!(gone(db.crawler.raw(), keys::crawler_retry_key(&did)));
assert!(gone(db.crawler.raw(), keys::bulk_job_key(&did)));
assert!(gone(
db.counts.raw(),
crate::db::keys::did_collection_prefix(&did)
));
assert!(db.indexer.head(&record_key).into_diagnostic()?.is_none());
assert!(gone(
db.indexer.resync_buffer.raw(),
keys::resync_buffer_prefix(&did)
));
#[cfg(feature = "indexer_stream")]
assert!(gone(
db.stream.event_bodies.raw(),
keys::record_prefix_did(&did)
));
// with nothing left, tracking it again is a fresh backfill
assert_eq!(repos.track([did.clone()]).await?, vec![did.clone()]);
repos.erase_repo(&did, true).await?;
assert!(
db.filter
.contains_key(crate::db::filter::exclude_key(did.as_str())?)
.into_diagnostic()?
);
assert!(repos.track([did.clone()]).await?.is_empty());
Ok(())
}
async fn fresh_state() -> miette::Result<(tempfile::TempDir, Arc<AppState>)> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
Ok((tmp, state))
}
fn many_dids(n: usize) -> Vec<Did> {
(0..n)
.map(|i| Did::new_owned(format!("did:web:repo{i}.example")).unwrap())
.collect()
}
#[tokio::test]
async fn a_long_track_with_a_shard_gone_applies_the_rest_and_a_resend_finishes_it()
-> miette::Result<()> {
let (_tmp, state) = fresh_state().await?;
let shards =
crate::ingest::indexer::FirehoseWorker::serve_shards(state.clone(), 4, Some(1));
let persisted = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let counted = persisted.clone();
*state.db.before_persist.lock() = Some(Arc::new(move || {
counted.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(())
}));
let dids = many_dids(600);
let on_gone_shard = |did: &&Did| shard_of(did, 4) == 1;
let err = ReposControl(state.clone())
.track(dids.clone())
.await
.unwrap_err();
assert!(err.to_string().contains("gone"), "{err}");
assert_eq!(persisted.load(std::sync::atomic::Ordering::SeqCst), 1);
let queued = |did: &Did| {
state
.db
.repo_metadata
.contains_key(keys::repo_metadata_key(did))
.unwrap()
};
for did in &dids {
assert_eq!(queued(did), !on_gone_shard(&did), "{did}");
}
let missing: Vec<_> = dids.iter().filter(on_gone_shard).cloned().collect();
assert!(!missing.is_empty());
// sending it all again, like a caller retrying after a crash between the sends, only
// changes what the first call didn't get to
drop(shards);
let _shards = crate::ingest::indexer::FirehoseWorker::serve_shards(state.clone(), 4, None);
assert_eq!(
ReposControl(state.clone()).track(dids.clone()).await?,
missing
);
assert!(ReposControl(state.clone()).track(dids).await?.is_empty());
assert_eq!(state.db.indexer.pending_iter().count(), 600);
state.db.indexer.assert_pending_dids_match();
Ok(())
}
#[tokio::test]
async fn a_record_redaction_runs_on_the_repos_indexer_shard() -> miette::Result<()> {
let (_tmp, state) = fresh_state().await?;
let did = Did::new_static("did:plc:testoperatorrepo").into_diagnostic()?;
let record_key =
keys::record_key(&did, "app.bsky.feed.post", &DbRkey::new("3kzbif5moe22m"));
let mut batch = state.db.inner.batch();
state
.db
.indexer
.stage_head(&mut batch, &record_key, record_body("head"));
batch.commit().into_diagnostic()?;
// the repo's shard never runs, so a redaction that goes through it never happens
let _shards =
crate::ingest::indexer::FirehoseWorker::serve_shards(state.clone(), 1, Some(0));
let redacted = ReposControl(state.clone())
.get(&did)
.delete_record_bodies(
"app.bsky.feed.post",
"3kzbif5moe22m",
crate::types::DeleteBodyTarget::All,
)
.await;
assert!(redacted.is_err());
assert_eq!(
state
.db
.indexer
.head(&record_key)
.into_diagnostic()?
.as_deref(),
Some(record_body("head").as_slice())
);
Ok(())
}
#[tokio::test]
async fn queue_moves_keep_the_did_order_copy_in_step() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let repos = ReposControl(state.clone());
let did = Did::new_static("did:plc:testoperatorrepo").into_diagnostic()?;
let other = Did::new_static("did:web:example.com").into_diagnostic()?;
let in_step = || state.db.indexer.assert_pending_dids_match();
repos.track([did.clone(), other.clone()]).await?;
in_step();
// already queued, so nothing moves
repos.track([did.clone()]).await?;
in_step();
repos.untrack([did.clone(), other.clone()]).await?;
in_step();
assert_eq!(state.db.indexer.pending_iter().count(), 0);
repos.track([did.clone()]).await?;
in_step();
// an untracked repo comes back under a fresh pending key
assert_eq!(repos.resync([other.clone()]).await?, vec![other.clone()]);
in_step();
assert_eq!(state.db.indexer.pending_iter().count(), 2);
Ok(())
}
#[tokio::test]
#[cfg(all(feature = "indexer_stream", not(feature = "backlinks")))]
async fn test_generate_car_ephemeral_rejected() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let mut config = crate::config::Config::default();
config.database_path = tmp.path().to_path_buf();
config.ephemeral = true;
let state = Arc::new(AppState::new(&config)?);
let repos = ReposControl(state);
let did = Did::new_static("did:plc:testephemeral").into_diagnostic()?;
let handle = repos.get(&did);
let res = handle.generate_car().await;
let Err(err) = res else {
panic!("generate_car in ephemeral mode should fail");
};
assert!(err.to_string().contains("ephemeral mode"));
Ok(())
}
#[tokio::test]
async fn test_generate_car_root_mismatch_fails() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let mut config = crate::config::Config::default();
config.database_path = tmp.path().to_path_buf();
config.ephemeral = false;
let state = Arc::new(AppState::new(&config)?);
let did = Did::new_static("did:plc:testmismatch").into_diagnostic()?;
let dummy_cid = crate::car::compute_cid(b"dummy_root");
let repo_state = crate::types::RepoState {
root: Some(crate::types::v2::Commit {
version: 3,
rev: crate::db::types::DbTid::from(&jacquard_common::types::string::Tid::now_0()),
data: dummy_cid,
prev: None,
sig: bytes::Bytes::new(),
}),
..crate::types::RepoState::synced()
};
let did_key = keys::repo_key(&did);
state
.db
.repos
.insert(&did_key, crate::db::ser_repo_state(&repo_state)?)
.into_diagnostic()?;
let repos = ReposControl(state);
let handle = repos.get(&did);
let res = handle.generate_car().await;
let Err(err) = res else {
panic!("generate_car with mismatched root should fail");
};
assert!(err.to_string().contains("mst root mismatch"));
Ok(())
}
#[tokio::test]
async fn get_repo_of_a_repo_a_failed_backfill_left_half_written_answers_not_found()
-> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let did = Did::new_static("did:plc:testmismatch").into_diagnostic()?;
// a root the stored records don't add up to, and the retry the failure queued
let repo_state = crate::types::RepoState {
root: Some(crate::types::v2::Commit {
version: 3,
rev: crate::db::types::DbTid::from(&jacquard_common::types::string::Tid::now_0()),
data: crate::car::compute_cid(b"new root"),
prev: None,
sig: bytes::Bytes::new(),
}),
..crate::types::RepoState::synced()
};
let resync = crate::types::ResyncState::Error {
kind: crate::types::ResyncErrorKind::Generic,
retry_count: 1,
next_retry: 0,
};
let mut batch = state.db.inner.batch();
batch.insert(
&state.db.repos,
keys::repo_key(&did),
crate::db::ser_repo_state(&repo_state)?,
);
batch.insert(
&state.db.indexer.resync,
keys::repo_key(&did),
rmp_serde::to_vec(&resync).into_diagnostic()?,
);
batch.commit().into_diagnostic()?;
let handle = ReposControl(state).get(&did);
assert!(handle.generate_car().await?.is_none());
Ok(())
}
#[tokio::test]
async fn generate_car_reports_a_deleted_head_body() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let mut config = crate::config::Config::default();
config.database_path = tmp.path().to_path_buf();
let state = Arc::new(AppState::new(&config)?);
let did = Did::new_static("did:plc:testredactedcar").into_diagnostic()?;
let body = record_body("deleted");
let cid = crate::car::compute_cid(&body);
let repo_state = crate::types::RepoState {
root: Some(crate::types::v2::Commit {
version: 3,
rev: crate::db::types::DbTid::from(&jacquard_common::types::string::Tid::now_0()),
data: cid,
prev: None,
sig: bytes::Bytes::new(),
}),
..crate::types::RepoState::synced()
};
let mut batch = state.db.inner.batch();
batch.insert(
&state.db.repos,
keys::repo_key(&did),
crate::db::ser_repo_state(&repo_state)?,
);
state.db.indexer.stage_head(
&mut batch,
keys::record_key(&did, "app.bsky.feed.post", &DbRkey::new("self")),
cid.to_bytes(),
);
batch.commit().into_diagnostic()?;
let handle = ReposControl(state).get(&did);
let err = match handle.generate_car().await {
Ok(_) => panic!("redacted head must make getRepo fail"),
Err(err) => err,
};
assert!(err.to_string().contains("record body was deleted"));
Ok(())
}
/// a repo per queue state, one of them stranded by a remove that named the
/// wrong did, and one with nothing queued
#[tokio::test]
async fn queued_listing_is_the_full_listing_minus_whats_not_queued() -> miette::Result<()> {
use crate::types::{ResyncErrorKind, ResyncState};
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let did = |s: &'static str| Did::new_static(s).into_diagnostic();
let pending = did("did:web:a.com")?;
let prefixed = did("did:web:a.com.au")?;
let both = did("did:plc:yk4q3id7id6p5z3bypvshc64")?;
let synced = did("did:web:synced.example")?;
let stale = did("did:web:stale.example")?;
let gone = did("did:web:gone.example")?;
let error = rmp_serde::to_vec(&ResyncState::Error {
kind: ResyncErrorKind::Generic,
retry_count: 1,
next_retry: 0,
})
.into_diagnostic()?;
let deactivated = rmp_serde::to_vec(&ResyncState::Gone {
status: RepoStatus::Deactivated,
})
.into_diagnostic()?;
let db = &state.db;
let mut batch = db.inner.batch();
for (index_id, did) in [&pending, &prefixed, &both, &synced, &stale, &gone]
.into_iter()
.enumerate()
{
batch.insert(
&db.repos,
keys::repo_key(did),
crate::db::ser_repo_state(&RepoState::backfilling())?,
);
batch.insert(
&db.repo_metadata,
keys::repo_metadata_key(did),
crate::db::ser_repo_meta(&RepoMetadata::backfilling(index_id as u64))?,
);
}
db.indexer
.stage_pending_insert(&mut batch, &pending, &keys::pending_key(0));
db.indexer
.stage_pending_insert(&mut batch, &both, &keys::pending_key(2));
// a superseded key still queued for the same repo
db.indexer
.stage_pending_insert(&mut batch, &both, &keys::pending_key(20));
db.indexer
.stage_pending_insert(&mut batch, &stale, &keys::pending_key(4));
batch.insert(&db.indexer.resync, keys::repo_key(&prefixed), &error);
batch.insert(&db.indexer.resync, keys::repo_key(&both), &error);
batch.insert(&db.indexer.resync, keys::repo_key(&gone), &deactivated);
batch.commit().into_diagnostic()?;
let mut batch = db.inner.batch();
db.indexer
.stage_pending_remove(&mut batch, &synced, &keys::pending_key(4));
batch.commit().into_diagnostic()?;
let repos = ReposControl(state.clone());
let expected = repos
.iter(None)
.filter(|row| {
row.as_ref()
.map_or(true, |info| info.pending || info.resync.is_some())
})
.map(|row| row.map(|info| info.did))
.collect::<miette::Result<Vec<_>>>()?;
assert_eq!(expected.len(), 4, "{expected:?}");
let queued = repos
.iter_queued(None)
.map(|row| row.map(|info| info.did))
.collect::<miette::Result<Vec<_>>>()?;
assert_eq!(queued, expected);
let mut paged = Vec::new();
let mut cursor = None;
// one more page than there are rows, so a cursor that stops moving fails here
for _ in 0..=expected.len() {
let Some(row) = repos.iter_queued(cursor.as_ref()).next().transpose()? else {
break;
};
cursor = Some(row.did.clone());
paged.push(row.did);
}
assert_eq!(paged, expected);
Ok(())
}
#[tokio::test]
async fn a_repo_that_changes_queue_between_pages_is_listed_once() -> miette::Result<()> {
use crate::types::{ResyncErrorKind, ResyncState};
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let repos = ReposControl(state.clone());
let mut dids = [
"did:web:one.example",
"did:web:two.example",
"did:web:three.example",
]
.map(|s| Did::new_static(s).unwrap());
dids.sort_by_key(keys::repo_key);
repos.track(dids.clone()).await?;
let error = rmp_serde::to_vec(&ResyncState::Error {
kind: ResyncErrorKind::Transport,
retry_count: 0,
next_retry: 0,
})
.into_diagnostic()?;
let db = &state.db;
// the way a failed backfill moves a repo, both writes in one batch
let fail = |did: &Did| -> miette::Result<()> {
let index_id = crate::db::deser_repo_meta(
&db.repo_metadata
.get(keys::repo_metadata_key(did))
.into_diagnostic()?
.expect("tracked repo has metadata"),
)?
.index_id;
let mut batch = db.inner.batch();
db.indexer
.stage_pending_remove(&mut batch, did, &keys::pending_key(index_id));
batch.insert(&db.indexer.resync, keys::repo_key(did), &error);
batch.commit().into_diagnostic()
};
let page = |cursor: Option<&Did>| -> miette::Result<Option<Did>> {
Ok(repos
.iter_queued(cursor)
.next()
.transpose()?
.map(|info| info.did))
};
let first = page(None)?.expect("three queued");
fail(&dids[1])?;
let second = page(Some(&first))?.expect("two left");
fail(&dids[2])?;
let third = page(Some(&second))?.expect("one left");
assert_eq!([first, second, third.clone()], dids);
assert!(page(Some(&third))?.is_none());
Ok(())
}
/// a row like the ones the relay and the crawler make, with the repo counted but no metadata
fn row_without_metadata(state: &AppState, did: &Did) -> miette::Result<fjall::Slice> {
let mut row = RepoState::backfilling();
row.active = false;
row.status = RepoStatus::Deactivated;
row.handle = Some(
jacquard_common::types::string::Handle::new_static("alice.test").into_diagnostic()?,
);
row.pds = Some("https://pds.example/".into());
row.last_message_time = Some(1_783_785_600_123);
row.last_account_time = Some(1_783_785_600_456);
let mut txn = crate::db::Txn::new(&state.db);
txn.batch.insert(
&state.db.indexer.resync,
keys::repo_key(did),
rmp_serde::to_vec(&crate::types::ResyncState::Gone {
status: RepoStatus::Deactivated,
})
.into_diagnostic()?,
);
txn.transition_lifecycle(did, GaugeState::Resync(None))?;
txn.put_repo(did, &row)?;
txn.counts.add_repos(1);
txn.commit()?;
Ok(state
.db
.repos
.get(keys::repo_key(did))
.into_diagnostic()?
.expect("repo row"))
}
/// queued once under the key its new metadata names, with the row left as it was
fn assert_queued_without_touching_the_row(
state: &AppState,
did: &Did,
row: &fjall::Slice,
) -> miette::Result<()> {
let db = &state.db;
assert_eq!(
db.repos
.get(keys::repo_key(did))
.into_diagnostic()?
.as_ref(),
Some(row),
"the row was written over"
);
let metadata = crate::db::deser_repo_meta(
&db.repo_metadata
.get(keys::repo_metadata_key(did))
.into_diagnostic()?
.expect("queued repo has metadata"),
)?;
assert!(metadata.tracked);
assert!(
db.indexer
.is_pending(&keys::pending_key(metadata.index_id))
.into_diagnostic()?
);
assert_eq!(db.indexer.pending_iter().count(), 1);
db.indexer.assert_pending_dids_match();
assert!(
!db.indexer
.resync
.contains_key(keys::repo_key(did))
.into_diagnostic()?
);
assert_eq!(db.get_count_sync("pending"), 1);
assert_eq!(db.get_count_sync("repos"), 1);
Ok(())
}
#[tokio::test]
async fn resync_queues_a_row_without_metadata() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let did = Did::new_static("did:web:listed.example").into_diagnostic()?;
let row = row_without_metadata(&state, &did)?;
let repos = ReposControl(state.clone());
assert_eq!(repos.resync([did.clone()]).await?, vec![did.clone()]);
assert_queued_without_touching_the_row(&state, &did, &row)
}
#[tokio::test]
async fn track_keeps_a_row_without_metadata_and_counts_it_once() -> miette::Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let did = Did::new_static("did:web:listed.example").into_diagnostic()?;
let row = row_without_metadata(&state, &did)?;
let repos = ReposControl(state.clone());
assert_eq!(repos.track([did.clone()]).await?, vec![did.clone()]);
assert_queued_without_touching_the_row(&state, &did, &row)?;
// queued already, so a second track leaves it alone
assert!(repos.track([did.clone()]).await?.is_empty());
assert_queued_without_touching_the_row(&state, &did, &row)
}
#[tokio::test]
async fn a_bulk_resync_queues_valid_repos_and_a_row_without_metadata() -> miette::Result<()> {
use crate::types::{ResyncErrorKind, ResyncState};
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&config)?);
let db = &state.db;
let repos = ReposControl(state.clone());
let valid =
["did:web:one.example", "did:web:two.example"].map(|s| Did::new_static(s).unwrap());
repos.track(valid.clone()).await?;
let error = rmp_serde::to_vec(&ResyncState::Error {
kind: ResyncErrorKind::Transport,
retry_count: 0,
next_retry: 0,
})
.into_diagnostic()?;
for did in &valid {
let index_id = crate::db::deser_repo_meta(
&db.repo_metadata
.get(keys::repo_metadata_key(did))
.into_diagnostic()?
.expect("tracked repo has metadata"),
)?
.index_id;
let mut txn = crate::db::Txn::new(db);
db.indexer
.stage_pending_remove(&mut txn.batch, did, &keys::pending_key(index_id));
txn.batch
.insert(&db.indexer.resync, keys::repo_key(did), &error);
txn.transition_lifecycle(did, GaugeState::Resync(Some(ResyncErrorKind::Transport)))?;
txn.commit()?;
}
let orphan = Did::new_static("did:web:listed.example").into_diagnostic()?;
let row = row_without_metadata(&state, &orphan)?;
assert_eq!(db.get_count_sync("resync"), 3);
assert_eq!(db.get_count_sync("error_transport"), 2);
let batch = vec![valid[0].clone(), orphan.clone(), valid[1].clone()];
assert_eq!(repos.resync(batch.clone()).await?, batch);
assert_eq!(
db.repos
.get(keys::repo_key(&orphan))
.into_diagnostic()?
.as_ref(),
Some(&row),
"the orphan's row was written over"
);
for did in &batch {
let metadata = crate::db::deser_repo_meta(
&db.repo_metadata
.get(keys::repo_metadata_key(did))
.into_diagnostic()?
.expect("queued repo has metadata"),
)?;
assert!(metadata.tracked);
assert!(
db.indexer
.is_pending(&keys::pending_key(metadata.index_id))
.into_diagnostic()?
);
assert!(
!db.indexer
.resync
.contains_key(keys::repo_key(did))
.into_diagnostic()?
);
}
assert_eq!(db.indexer.pending_iter().count(), 3);
db.indexer.assert_pending_dids_match();
assert_eq!(db.get_count_sync("repos"), 3);
assert_eq!(db.get_count_sync("pending"), 3);
assert_eq!(db.get_count_sync("resync"), 0);
assert_eq!(db.get_count_sync("error_transport"), 0);
Ok(())
}
}