//! the short writes to a repo that aren't a firehose event, like track, untrack and resync,
//! which go to the repo's owner shard so they're in order with everything else written for it
use std::sync::Arc;
use std::sync::atomic::Ordering;
use jacquard_common::IntoStatic;
use jacquard_common::types::did::Did;
use miette::{IntoDiagnostic, Result};
use super::bulk::{Busy, ShardState, busy_of};
use super::message::{IndexerTx, RepoOp, shard_of};
use super::refetch::{self, Charged};
use crate::db::{self, Db, Txn, keys};
use crate::state::AppState;
use crate::types::{GaugeState, RepoMetadata, RepoState};
/// most repos one message of a multi-repo write carries to their shard, so a long list doesn't
/// hold the shard for all of it before it gets back to its live messages
pub(crate) const REPOS_PER_MESSAGE: usize = 256;
/// where a write ran, see [`place`]
enum Placed<T, I> {
Here(T),
/// the shards run, so the input comes back for them
Shards(IndexerTx, I),
}
/// runs `here` on `input` right here when no shard takes work, holding the shards' sender off
/// meanwhile so they can't start writing under it. otherwise hands `input` back with the sender
fn place<I, T>(
state: &AppState,
input: I,
here: impl FnOnce(&AppState, I) -> T,
) -> Result<Placed<T, I>> {
let indexer = state.indexer.lock();
if let Some(tx) = indexer.clone() {
return Ok(Placed::Shards(tx, input));
}
no_shard_running(state)?;
Ok(Placed::Here(here(state, input)))
}
/// the shards' sender is gone once hydrant stops too, but a shard can still be in its last
/// step then, and a write in place would race it
pub(crate) fn no_shard_running(state: &AppState) -> Result<()> {
if state.stop.is_set() && state.shards_running.load(Ordering::SeqCst) > 0 {
miette::bail!("hydrant is stopping and its indexer shards are still running");
}
Ok(())
}
/// runs `work` on `did`'s owner shard, see [`super::message::OnShard`], and waits for it, on a thread that may
/// block. with no shards taking work it runs right here instead, see [`place`]
pub(crate) fn on_owner_blocking<T: Send + 'static>(
state: &AppState,
did: &Did,
work: impl FnOnce(&AppState, Option<&mut ShardState>) -> T + Send + Sync + 'static,
) -> Result<T> {
match place(state, work, |state, work| work(state, None))? {
Placed::Here(done) => Ok(done),
Placed::Shards(tx, work) => {
tx.on_owner_blocking(did, move |state, shard| work(state, Some(shard)))
}
}
}
/// [`on_owner_blocking`] from async code
pub(crate) async fn on_owner<T: Send + 'static>(
state: &Arc<AppState>,
did: &Did,
work: impl FnOnce(&AppState, Option<&mut ShardState>) -> T + Send + Sync + 'static,
) -> Result<T> {
let here = state.clone();
let placed = state
.db
.run(move |_| place(&here, work, |state, work| work(state, None)))
.await?;
match placed {
Placed::Here(done) => Ok(done),
Placed::Shards(tx, work) => {
tx.on_owner(did, move |state, shard| work(state, Some(shard)))
.await
}
}
}
/// runs `work` on every one of `items`, on each one's owner shard at most [`REPOS_PER_MESSAGE`]
/// to a message, or all at once right here when no shard takes work, see [`place`]. an answer
/// per message, so a shard that failed or is gone fails its own and the others' writes stay
pub(crate) async fn on_owners<I, T>(
state: &Arc<AppState>,
items: Vec<I>,
did_of: fn(&I) -> &Did,
work: impl Fn(&AppState, Option<&mut ShardState>, Vec<I>) -> T + Clone + Send + Sync + 'static,
) -> Vec<Result<T>>
where
I: Clone + Send + Sync + 'static,
T: Send + 'static,
{
let (here, in_place) = (state.clone(), work.clone());
let placed = state
.db
.run(move |_| place(&here, items, |state, items| in_place(state, None, items)))
.await;
let (tx, items) = match placed {
Ok(Placed::Here(done)) => return vec![Ok(done)],
Ok(Placed::Shards(tx, items)) => (tx, items),
Err(e) => return vec![Err(e)],
};
let num_shards = tx.inner.num_shards();
let mut by_shard = vec![Vec::new(); num_shards];
for item in items {
by_shard[shard_of(did_of(&item), num_shards)].push(item);
}
let messages = by_shard
.iter()
.flat_map(|items| items.chunks(REPOS_PER_MESSAGE))
.map(|batch| {
let (owner, batch) = (did_of(&batch[0]).clone(), batch.to_vec());
let (tx, work) = (tx.clone(), work.clone());
async move {
tx.on_owner(&owner, move |state, shard| work(state, Some(shard), batch))
.await
}
});
futures::future::join_all(messages).await
}
/// runs on the repo's owner shard, so it stays in order with the other writes to that repo
pub(crate) async fn write_repo<T: Send + 'static>(
state: &Arc<AppState>,
did: &Did,
work: impl FnOnce(&Db) -> Result<T> + Send + Sync + 'static,
) -> Result<T> {
on_owner(state, did, move |state, _| work(&state.db)).await?
}
/// [`repo_op`] for each of `dids`, and the ones it changed. a failure keeps what the dids
/// before it did
pub(crate) fn repo_ops(
state: &AppState,
shard: Option<&ShardState>,
op: RepoOp,
dids: &[Did],
) -> Result<Vec<Did>> {
let mut changed = Vec::new();
for did in dids {
if repo_op(state, shard, op, did)? {
changed.push(did.clone());
}
}
Ok(changed)
}
/// runs `op` for `did` in its own txn and says whether it changed anything. `shard` is the
/// repo's owner, or `None` when no shard is running
pub(crate) fn repo_op(
state: &AppState,
shard: Option<&ShardState>,
op: RepoOp,
did: &Did,
) -> Result<bool> {
let busy = busy_of(state, shard, did)?;
let mut txn = Txn::new(&state.db);
let mut parked = false;
let changed = match op {
// an erase or a wipe would forget whatever the backfill wrote
RepoOp::Resync | RepoOp::ResyncCrawled if busy.as_ref().is_some_and(Busy::drops_live) => {
false
}
RepoOp::Resync | RepoOp::ResyncCrawled => {
let cause = (op == RepoOp::ResyncCrawled).then_some("crawl");
let queued = queue_backfill(state, did, &mut txn, cause)?;
parked = queued == Queued::Deferred;
queued == Queued::Queued
}
RepoOp::Track if busy.as_ref().is_some_and(Busy::drops_live) => false,
RepoOp::Track => track(state, did, &mut txn)?,
// a backfill partway through writing its steps would be left half written, so the
// caller has to try again once it's done
RepoOp::Untrack if busy.map(|busy| busy.holds(state)).transpose()? == Some(true) => false,
RepoOp::Untrack => untrack(&state.db, did, &mut txn)?,
};
if changed || parked {
txn.commit()?;
}
Ok(changed)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Queued {
Queued,
/// out of refetches. the park is in the txn, so the caller has to commit it
Deferred,
Skipped,
}
/// queues a backfill of a repo hydrant knows, under a fresh queue key, unless it's excluded or
/// already queued. passing a `cause` costs a refetch, `None` is free
pub(crate) fn queue_backfill(
state: &AppState,
did: &Did,
txn: &mut Txn<'_>,
cause: Option<&str>,
) -> Result<Queued> {
let db = &state.db;
if txn.is_excluded(did)? {
return Ok(Queued::Skipped);
}
let did_key = keys::repo_key(did);
let Some(row) = db.repos.get(&did_key).into_diagnostic()? else {
return Ok(Queued::Skipped);
};
let mut repo = db::deser_repo_state(&row)?.into_static();
let metadata_key = keys::repo_metadata_key(did);
// rows the relay or the crawler made have no metadata until something queues them
let mut metadata = match db.repo_metadata.get(&metadata_key).into_diagnostic()? {
Some(bytes) => {
let mut metadata = db::deser_repo_meta(&bytes)?;
let old_pending = keys::pending_key(metadata.index_id);
if db.indexer.is_pending(&old_pending).into_diagnostic()? {
return Ok(Queued::Skipped);
}
if let Some(cause) = cause
&& refetch::charge(state, txn, did, &mut repo, &mut metadata, cause)?
== Charged::Deferred
{
return Ok(Queued::Deferred);
}
db.indexer
.stage_pending_remove(&mut txn.batch, did, &old_pending);
metadata
}
None => RepoMetadata::backfilling(0),
};
metadata.tracked = true;
metadata.refetch_deferred = false;
metadata.index_id = rand::Rng::next_u64(&mut rand::rng());
db.indexer
.stage_pending_insert(&mut txn.batch, did, &keys::pending_key(metadata.index_id));
txn.batch.remove(&db.indexer.resync, &did_key);
txn.batch.insert(
&db.repo_metadata,
&metadata_key,
db::ser_repo_meta(&metadata)?,
);
txn.transition_lifecycle(did, GaugeState::Pending)?;
Ok(Queued::Queued)
}
fn track(state: &AppState, did: &Did, txn: &mut Txn<'_>) -> Result<bool> {
let db = &state.db;
if txn.is_excluded(did)? {
return Ok(false);
}
let metadata_key = keys::repo_metadata_key(did);
let existing_metadata = db
.repo_metadata
.get(&metadata_key)
.into_diagnostic()?
.map(|bytes| db::deser_repo_meta(&bytes))
.transpose()?;
if let Some(metadata) = existing_metadata {
return Ok(!metadata.tracked && queue_backfill(state, did, txn, None)? == Queued::Queued);
}
if txn.read_repo(did)?.is_some() {
// the relay or the crawler made this row and counted the repo, so it only needs
// queueing
return Ok(queue_backfill(state, did, txn, None)? == Queued::Queued);
}
let metadata = RepoMetadata::backfilling(rand::random());
txn.put_repo(did, &RepoState::backfilling())?;
txn.batch.insert(
&db.repo_metadata,
&metadata_key,
db::ser_repo_meta(&metadata)?,
);
db.indexer
.stage_pending_insert(&mut txn.batch, did, &keys::pending_key(metadata.index_id));
txn.counts.add_repos(1);
txn.transition_lifecycle(did, GaugeState::Pending)?;
Ok(true)
}
fn untrack(db: &db::Db, did: &Did, txn: &mut Txn<'_>) -> Result<bool> {
let did_key = keys::repo_key(did);
let metadata_key = keys::repo_metadata_key(did);
if txn.read_repo(did)?.is_none() {
return Ok(false);
}
let existing_metadata = db
.repo_metadata
.get(&metadata_key)
.into_diagnostic()?
.map(|bytes| db::deser_repo_meta(&bytes))
.transpose()?;
let Some(mut metadata) = existing_metadata.filter(|metadata| metadata.tracked) else {
return Ok(false);
};
metadata.tracked = false;
metadata.refetch_deferred = false;
txn.batch.insert(
&db.repo_metadata,
&metadata_key,
db::ser_repo_meta(&metadata)?,
);
db.indexer
.stage_pending_remove(&mut txn.batch, did, &keys::pending_key(metadata.index_id));
txn.batch.remove(&db.indexer.resync, &did_key);
txn.transition_lifecycle(did, GaugeState::Synced)?;
Ok(true)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ingest::indexer::FirehoseWorker;
/// what shutdown leaves: the sender is gone, but a shard can still be running its last step
#[tokio::test]
async fn nothing_writes_in_place_while_a_shard_still_runs() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let state = Arc::new(AppState::new(&crate::config::Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?);
let did = Did::new_static("did:plc:ewvi7nxzyoun6zhxrhs64oiz").into_diagnostic()?;
let (tx, shard) = FirehoseWorker::spawn_one_shard(state.clone());
// the shard answers once it runs, so it's counted by then
tx.on_owner(&did, |_, _| ()).await?;
// before the stop, a test's own shard and an in place write may run side by side
write_repo(&state, &did, |_| Ok(())).await?;
state.stop.trigger();
let wrote = write_repo(&state, &did, |_| Ok(())).await;
assert!(wrote.is_err(), "wrote in place next to a running shard");
drop(tx);
tokio::task::spawn_blocking(move || shard.join())
.await
.into_diagnostic()?
.map_err(|_| miette::miette!("indexer shard panicked"))?;
write_repo(&state, &did, |_| Ok(())).await
}
#[test]
fn an_operators_resync_is_free_where_a_crawls_costs_a_refetch() -> Result<()> {
use super::refetch::tests as refetch;
let (_tmp, state) = refetch::state(Some(1));
let did = refetch::did();
let used = Some((chrono::Utc::now().timestamp(), 1));
refetch::seed(&state, true, used);
let pending = || state.db.indexer.pending_iter().count();
assert!(!repo_op(&state, None, RepoOp::ResyncCrawled, &did)?);
assert_eq!(pending(), 0);
assert_eq!(
refetch::row(&state).status,
crate::types::RepoStatus::Throttled
);
assert!(refetch::metadata_of(&state, &did).refetch_deferred);
assert!(repo_op(&state, None, RepoOp::Resync, &did)?);
assert_eq!(pending(), 1);
let metadata = refetch::metadata_of(&state, &did);
assert!(
!metadata.refetch_deferred,
"the operator's fetch replaces the wait"
);
assert_eq!(metadata.refetch.map(|allowance| allowance.used), Some(1));
assert!(refetch::entry_of(&state, &did).is_none());
Ok(())
}
}