use std::sync::Arc;
use crate::ingest::firehose_stats::StatsInstant;
use jacquard_common::types::did::Did;
use miette::{IntoDiagnostic, Result};
use tokio::runtime::Handle;
use tracing::{debug, error, info, info_span, trace, warn};
use url::Url;
use crate::db::CountDeltas;
use crate::ingest::cursor::{self, Ticket};
use crate::ingest::stream::SubscribeReposMessage;
use crate::ingest::validation::ValidationOptions;
use crate::ingest::{BufferRx, BufferTx, IngestMessage};
use crate::state::AppState;
use super::context::{Answered, Checked, RefreshOrigin};
use super::handlers::Handled;
use super::lookup::{self, Answer, Dropped, LookupError, LookupKind, Priority, Request, Submitted};
use super::parked::Held;
use super::sink::{EventSink, SinkSeed};
use super::{AuthorityOutcome, WorkerContext, relay_message_kind};
use crate::net::PublicHttpClient;
use crate::resolver::ResolveFailure;
enum Next {
Message(IngestMessage),
Answer(Answer),
}
struct Committed {
timings: crate::db::TxnCommitTimings,
broadcast: std::time::Duration,
cursor: std::time::Duration,
}
/// how far a message got
enum Step {
Done,
/// it needs a lookup for its did first, and staged nothing
Wait(LookupKind, WorkerMessage),
}
pub struct WorkerMessage {
pub(crate) is_pds: bool,
pub(crate) firehose: Url,
pub(crate) msg: SubscribeReposMessage<'static>,
}
impl WorkerMessage {
fn did(&self) -> Did {
self.msg
.did()
.expect("the reader forwards no info frames")
.clone()
}
}
pub struct RelayWorker {
pub(crate) state: Arc<AppState>,
pub(crate) rxs: Vec<BufferRx>,
pub(crate) seed: SinkSeed,
pub(crate) verify_signatures: bool,
pub(crate) num_shards: usize,
pub(crate) validation_opts: Arc<ValidationOptions>,
pub(crate) http: PublicHttpClient,
}
impl RelayWorker {
pub fn new(
state: Arc<AppState>,
seed: SinkSeed,
verify_signatures: bool,
num_shards: usize,
validation_opts: ValidationOptions,
) -> (BufferTx, Self) {
let (tx, rxs) = BufferTx::channel(num_shards);
(
tx,
Self {
state,
rxs,
seed,
verify_signatures,
num_shards,
validation_opts: Arc::new(validation_opts),
http: crate::net::public_http().clone(),
},
)
}
pub fn run(self, handle: Handle) -> Result<()> {
let (exit_tx, exit_rx) = std::sync::mpsc::channel();
for (i, rx) in self.rxs.into_iter().enumerate() {
let state = Arc::clone(&self.state);
let seed = self.seed.clone();
let verify = self.verify_signatures;
let h = handle.clone();
let opts = self.validation_opts.clone();
let http = self.http.clone();
let exit_tx = exit_tx.clone();
std::thread::Builder::new()
.name(format!("relay-shard-{i}"))
.spawn(move || {
let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
Self::shard(i, rx, state, seed, verify, h, opts, http);
}));
let _ = exit_tx.send((i, res));
})
.into_diagnostic()?;
}
drop(exit_tx);
info!(num = self.num_shards, "relay worker: started shards");
match exit_rx.recv() {
Ok((id, Ok(()))) => Err(miette::miette!("relay worker shard {id} shut down")),
Ok((id, Err(_))) => Err(miette::miette!("relay worker shard {id} panicked")),
Err(_) => Err(miette::miette!("relay worker shards shut down")),
}
}
#[allow(clippy::too_many_arguments)]
fn shard(
id: usize,
mut rx: BufferRx,
state: Arc<AppState>,
seed: SinkSeed,
verify_signatures: bool,
handle: Handle,
validation_opts: Arc<ValidationOptions>,
http: PublicHttpClient,
) {
let _guard = handle.enter();
let span = info_span!("worker_shard", shard = id);
let _entered = span.clone().entered();
debug!("relay shard started");
let shard_stats = state.firehose_stats.relay_shard(id);
let identity_epoch = crate::db::counts::identity_ttl_epoch(&state.db)
.inspect_err(|e| error!(err = %e, "relay shard: can't read the identity ttl epoch"))
.unwrap_or_else(|_| chrono::Utc::now().timestamp());
let (answers_tx, mut answers) = tokio::sync::mpsc::unbounded_channel();
let mut ctx = WorkerContext {
verify_signatures,
state: &state,
vctx: crate::ingest::validation::ValidationContext {
opts: &validation_opts,
},
stats: shard_stats.clone(),
batch: state.db.inner.batch(),
count_deltas: CountDeltas::default(),
sink: EventSink::new(&seed, shard_stats.clone()),
http,
error_counts: Default::default(),
wrong_host_authority: Default::default(),
loaded: None,
staged: None,
answers: answers_tx,
identity_epoch,
refresh_origins: Default::default(),
answered: None,
parked: Default::default(),
};
loop {
// no timer for this: the held did's own lookup answers within its deadline, which
// wakes the shard, and a timer polled after the runtime shut down would panic
Self::sweep(&mut ctx, id);
let next = handle.block_on(async {
tokio::select! {
// answers first, they're what lets held messages go
biased;
Some(answer) = answers.recv() => Some(Next::Answer(answer)),
msg = rx.recv() => msg.map(Next::Message),
}
});
let Some(next) = next else {
break;
};
// past this the runtime may be gone, and the firehose cursor hasn't moved for what's
// left or held, so it's read again next start
if state.stop.is_set() {
break;
}
match next {
Next::Message(msg) => Self::ingest(&mut ctx, id, msg),
Next::Answer(answer) => Self::answer(&mut ctx, id, answer),
}
}
}
fn ingest(ctx: &mut WorkerContext, id: usize, msg: IngestMessage) {
let IngestMessage::Firehose {
url,
is_pds,
msg,
ticket,
} = msg;
let held = Held {
msg: WorkerMessage {
is_pds,
firehose: url,
msg,
},
ticket,
};
let did = held.msg.did();
ctx.stats.record_received(held.ticket.seq());
if let Some(kind) = ctx.parked.kind(&did) {
if let Err(dropped) = ctx.parked.push(&ctx.state.lookups, &did, held) {
warn!(%did, "too much held for one did, dropping it all");
ctx.state.lookups.back_off(&did);
Self::drop_held(ctx, id, kind, dropped, Dropped::Cap);
}
return;
}
Self::step(ctx, id, held);
}
/// runs one message, committing what it did, or holds it while its did's lookup runs
fn step(ctx: &mut WorkerContext, id: usize, held: Held) {
let message_started = StatsInstant::now();
let Held { msg, ticket } = held;
let did = msg.did();
let firehose = msg.firehose.clone();
let _span =
info_span!("relay", did = %did, firehose = %firehose, seq = %ticket.seq()).entered();
let process_started = StatsInstant::now();
match Self::process_message(ctx, msg) {
Ok(Step::Done) => {}
Ok(Step::Wait(kind, msg)) => {
ctx.discard();
Self::park(ctx, id, kind, Held { msg, ticket });
return;
}
Err(e) => {
ctx.stats.record_process_error();
error!(did = %did, err = %e, "relay shard: error processing message");
}
}
let process_message = process_started.elapsed();
let Some(Committed {
timings,
broadcast,
cursor,
}) = Self::commit(ctx, id, &did, Some((&firehose, ticket)))
else {
return;
};
ctx.stats
.record_processed(crate::ingest::firehose_stats::RelayShardTimings {
process_message,
stage_counts: timings.stage_counts,
stage_and_commit: timings.stage_and_commit,
apply_counts: timings.apply_counts,
broadcast,
cursor,
total: message_started.elapsed(),
});
}
/// holds a message whose did needs a lookup first, and starts the lookup
fn park(ctx: &mut WorkerContext, id: usize, kind: LookupKind, held: Held) {
let did = held.msg.did();
if ctx
.answered
.as_ref()
.is_some_and(|answered| answered.did == did)
{
// every step takes a replayed answer as its own, so this is a bug, and holding the
// message again could wait forever
error!(%did, ?kind, "relay shard: a replayed message asked for another lookup");
Self::drop_held(ctx, id, kind, [held], Dropped::Cap);
return;
}
let host = lookup::identity_host(&did);
let lookups = &ctx.state.lookups;
if let Err(held) = ctx.parked.start(lookups, did.clone(), kind, held, host) {
debug!(%did, parked = ctx.parked.len(), "can't hold more, dropping");
Self::drop_held(ctx, id, kind, [held], Dropped::Cap);
return;
}
let request = Request {
did: did.clone(),
kind,
priority: Priority::Park,
resolver: ctx.state.resolver.clone(),
probe: ctx.probe_client(),
};
match lookups.submit(request, &ctx.answers) {
Submitted::Started | Submitted::Joined => {}
Submitted::Failing | Submitted::Full => {
let held = ctx.parked.take(lookups, &did).expect("just started");
Self::drop_held(ctx, id, kind, held, Dropped::Cap);
}
}
}
/// lets go of one did's held messages without running them, each done for the cursor all
/// the same. for a did with a row the newest #account among them still runs, as if the
/// lookup `kind` failed, since it's the account's last word and a takedown dropped here
/// wouldn't come again
fn drop_held(
ctx: &mut WorkerContext,
id: usize,
kind: LookupKind,
held: impl IntoIterator<Item = Held>,
why: Dropped,
) {
let mut account = None;
for held in held {
if !matches!(held.msg.msg, SubscribeReposMessage::Account(_)) {
Self::let_go(ctx, held, why);
} else if let Some(older) = account.replace(held) {
Self::let_go(ctx, older, why);
}
}
let Some(account) = account else {
return;
};
let did = account.msg.did();
// a new account has no row to land it on, and a replay with this did's answer is what
// asked to hold it, so it would only ask again
if kind == LookupKind::NewAccount
|| ctx
.answered
.as_ref()
.is_some_and(|answered| answered.did == did)
{
Self::let_go(ctx, account, why);
return;
}
let replaying = ctx.answered.replace(Answered {
did,
kind,
result: Err(LookupError::Abandoned),
resolved_at: chrono::Utc::now().timestamp(),
});
Self::step(ctx, id, account);
ctx.answered = replaying;
}
fn let_go(ctx: &WorkerContext, Held { msg, ticket }: Held, why: Dropped) {
ctx.state.lookups.count_dropped(why);
cursor::done(ctx.state, &msg.firehose, ticket);
}
/// replays the messages a did held with its lookup's answer, or puts a background
/// refresh's doc on the row
fn answer(ctx: &mut WorkerContext, id: usize, answer: Answer) {
let origin = ctx.refresh_origins.remove(&answer.did);
if let Ok(resolved) = &answer.result {
ctx.stats.record_resolve_doc(resolved.doc_elapsed);
if let Some(elapsed) = resolved.probe_elapsed {
ctx.stats.record_repo_status_probe(elapsed);
}
}
let kind = ctx.parked.kind(&answer.did);
let (Some(kind), Some(held)) = (kind, ctx.parked.take(&ctx.state.lookups, &answer.did))
else {
Self::apply_answer(ctx, id, answer, origin);
return;
};
if let Err(e) = &answer.result {
debug!(did = %answer.did, err = %e, kind = ?answer.kind, "lookup failed, replaying without it");
}
ctx.answered = Some(Answered {
did: answer.did,
kind,
result: answer.result,
resolved_at: answer.resolved_at,
});
for held in held {
Self::step(ctx, id, held);
}
ctx.answered = None;
}
/// drops what was held too long, and backs the did off
fn sweep(ctx: &mut WorkerContext, id: usize) {
for (did, kind, held) in ctx.parked.expired(&ctx.state.lookups) {
debug!(%did, "held too long, dropping");
ctx.state.lookups.back_off(&did);
Self::drop_held(ctx, id, kind, held, Dropped::Deadline);
}
}
/// commits what was staged for `did` and hands its events on, `None` if the commit failed.
/// `frame` is the firehose frame the message came from, done here unless an event of it went
/// on to the indexer, which marks it done after its own commit
fn commit(
ctx: &mut WorkerContext,
id: usize,
did: &Did,
frame: Option<(&Url, Ticket)>,
) -> Option<Committed> {
let batch = std::mem::replace(&mut ctx.batch, ctx.state.db.inner.batch());
let count_deltas = std::mem::take(&mut ctx.count_deltas);
let mut txn = crate::db::Txn::from_parts(&ctx.state.db, batch, count_deltas);
let (loaded, staged) = (ctx.loaded.take(), ctx.staged.take());
let firehose = frame.as_ref().map(|(firehose, _)| *firehose);
if let Err(e) = ctx
.sink
.stage_row(ctx.state, &mut txn, did, loaded, staged, firehose)
{
ctx.stats.record_process_error();
error!(did = %did, err = %e, "relay shard: failed to stage repo state");
}
let (staged, timings) = match ctx.sink.commit_txn(ctx.state, txn) {
Ok(staged) => staged,
Err(e) => {
ctx.stats.record_commit_error();
error!(shard = id, err = %e, "relay shard: failed to commit batch");
// what it staged for the indexer goes too, or the next frame's ticket and row
// would ride on this frame's events
ctx.sink.discard();
// a later frame's commit would have moved the cursor past this one anyway
if let Some((firehose, ticket)) = frame {
cursor::done(ctx.state, firehose, ticket);
}
return None;
}
};
let cursor_started = StatsInstant::now();
if let Some((firehose, ticket)) = frame
&& let Some(ticket) = ctx.sink.hand_off(ticket)
{
cursor::done(ctx.state, firehose, ticket);
}
let cursor = cursor_started.elapsed();
let broadcast_started = StatsInstant::now();
ctx.sink.flush(ctx.state, staged);
Some(Committed {
timings,
broadcast: broadcast_started.elapsed(),
cursor,
})
}
/// puts a background refresh's doc on the row it was for, or the doc of a lookup whose
/// held messages were dropped already, it's still the newest. the pool cached it already.
/// an answer that the did's identity is broken takes the key off the row instead
fn apply_answer(
ctx: &mut WorkerContext,
id: usize,
answer: Answer,
origin: Option<RefreshOrigin>,
) {
let did = answer.did;
let resolved = match answer.result {
Ok(resolved) => resolved,
Err(LookupError::Resolve(
failure @ (ResolveFailure::MissingPds
| ResolveFailure::NoSigningKey
| ResolveFailure::Tombstoned),
)) => {
debug!(%did, ?failure, "identity broken, dropping its stored key");
Self::forget_key(ctx, id, &did);
return;
}
Err(e) => {
debug!(%did, err = %e, "identity lookup nobody waits on failed");
return;
}
};
let mut repo_state = match ctx.read_repo_state(&did) {
Ok(Some(repo_state)) => repo_state,
Ok(None) => return,
Err(e) => {
ctx.stats.record_process_error();
error!(%did, err = %e, "relay shard: failed to read repo state");
return;
}
};
ctx.loaded = Some(repo_state.clone());
let moved = repo_state.update_from_resolve(resolved.doc, answer.resolved_at);
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
if moved
&& let Some(origin) = origin
&& let Err(e) =
ctx.announce_identity(&did, &repo_state, &origin.firehose, origin.seq, origin.time)
{
error!(%did, err = %e, "relay shard: failed to announce a refreshed identity");
}
#[cfg(not(any(feature = "indexer_stream", feature = "relay")))]
let _ = (moved, origin);
ctx.stage_repo_state(repo_state);
Self::commit(ctx, id, &did, None);
}
/// takes the signing key off `did`'s row, so its commits and syncs wait for a lookup again
/// and are dropped while it fails, as the sync spec wants for an account without a working
/// identity. a not found isn't enough for that, a plc mirror can be behind
fn forget_key(ctx: &mut WorkerContext, id: usize, did: &Did) {
let repo_state = match ctx.read_repo_state(did) {
Ok(Some(repo_state)) if repo_state.signing_key.is_some() => repo_state,
Ok(_) => return,
Err(e) => {
ctx.stats.record_process_error();
error!(%did, err = %e, "relay shard: failed to read repo state");
return;
}
};
ctx.loaded = Some(repo_state.clone());
let mut repo_state = repo_state;
repo_state.signing_key = None;
repo_state.touch();
ctx.stage_repo_state(repo_state);
Self::commit(ctx, id, did, None);
}
fn process_message(ctx: &mut WorkerContext, msg: WorkerMessage) -> Result<Step> {
if let Some(kind) = relay_message_kind(&msg.msg) {
ctx.stats.record_message_kind(kind);
}
// an #identity says the doc may have changed, so a backed off did may ask again. one from
// a relay we chose can do that before anything else, one from a pds only once the host
// check took it, in handle_identity
if let SubscribeReposMessage::Identity(identity) = &msg.msg
&& !msg.is_pds
{
ctx.state.lookups.clear_for_identity(&identity.did);
}
let load_started = StatsInstant::now();
let repo_state_result = ctx.load_repo_state(&msg);
ctx.stats.record_repo_state_load(load_started.elapsed());
let mut repo_state = match repo_state_result? {
Checked::Done(Some(repo_state)) => repo_state,
Checked::Done(None) => return Ok(Step::Done),
Checked::Wait(kind) => return Ok(Step::Wait(kind, msg)),
};
let did = msg.msg.did().expect("already checked for did");
let mut doc_refreshed = false;
let mut authority_identity_moved = false;
if let Some(host) = msg.firehose.host_str()
&& msg.is_pds
{
let authority_started = StatsInstant::now();
let outcome_result = match ctx.check_host_authority(did, &mut repo_state, host) {
Ok(Checked::Wait(kind)) => return Ok(Step::Wait(kind, msg)),
Ok(Checked::Done(outcome)) => Ok(outcome),
Err(e) => Err(e),
};
ctx.stats.record_host_authority(
authority_started.elapsed(),
match &outcome_result {
Ok(AuthorityOutcome::Authorized) => {
crate::ingest::firehose_stats::HostAuthorityStatsOutcome::Authorized
}
Ok(AuthorityOutcome::WasStale { .. }) => {
crate::ingest::firehose_stats::HostAuthorityStatsOutcome::WasStale
}
Ok(AuthorityOutcome::WrongHost { .. }) => {
crate::ingest::firehose_stats::HostAuthorityStatsOutcome::WrongHost
}
Err(_) => crate::ingest::firehose_stats::HostAuthorityStatsOutcome::Error,
},
);
let outcome = outcome_result?;
match outcome {
AuthorityOutcome::WrongHost { expected } => {
if !ctx.inc_error(host) {
warn!(got = host, expected = %expected, "message rejected: wrong host");
}
return Ok(Step::Done);
}
AuthorityOutcome::WasStale { identity_moved } => {
doc_refreshed = true;
authority_identity_moved = identity_moved;
}
AuthorityOutcome::Authorized => {}
}
ctx.reset_error(host);
}
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
if authority_identity_moved {
// an identity event announces itself
if !matches!(msg.msg, SubscribeReposMessage::Identity(_))
&& let Some((seq, time)) = msg.msg.position()
{
ctx.announce_identity(did, &repo_state, &msg.firehose, seq, time.clone())?;
}
}
let WorkerMessage {
is_pds,
firehose,
msg,
} = msg;
let waiting = |kind, msg| {
Step::Wait(
kind,
WorkerMessage {
is_pds,
firehose: firehose.clone(),
msg,
},
)
};
match msg {
SubscribeReposMessage::Commit(commit) => {
trace!("processing commit");
let started = StatsInstant::now();
let result = Self::handle_commit(ctx, repo_state, &firehose, commit);
ctx.stats.record_handle_message(
crate::ingest::firehose_stats::RelayMessageKind::Commit,
started.elapsed(),
);
Ok(match result? {
Handled::Done => Step::Done,
Handled::Wait(kind, commit) => {
waiting(kind, SubscribeReposMessage::Commit(commit))
}
})
}
SubscribeReposMessage::Sync(sync) => {
debug!("processing sync");
let started = StatsInstant::now();
let result = Self::handle_sync(ctx, repo_state, &firehose, sync);
ctx.stats.record_handle_message(
crate::ingest::firehose_stats::RelayMessageKind::Sync,
started.elapsed(),
);
Ok(match result? {
Handled::Done => Step::Done,
Handled::Wait(kind, sync) => waiting(kind, SubscribeReposMessage::Sync(sync)),
})
}
SubscribeReposMessage::Identity(identity) => {
debug!("processing identity");
let started = StatsInstant::now();
let result = Self::handle_identity(
ctx,
repo_state,
&firehose,
identity,
is_pds,
doc_refreshed,
authority_identity_moved,
);
ctx.stats.record_handle_message(
crate::ingest::firehose_stats::RelayMessageKind::Identity,
started.elapsed(),
);
Ok(match result? {
Handled::Done => Step::Done,
Handled::Wait(kind, identity) => {
waiting(kind, SubscribeReposMessage::Identity(identity))
}
})
}
SubscribeReposMessage::Account(account) => {
debug!("processing account");
let started = StatsInstant::now();
let result = Self::handle_account(ctx, repo_state, &firehose, *account, is_pds);
ctx.stats.record_handle_message(
crate::ingest::firehose_stats::RelayMessageKind::Account,
started.elapsed(),
);
result.map(|()| Step::Done)
}
_ => Ok(Step::Done),
}
}
}