use miette::{Diagnostic, IntoDiagnostic, Result};
use std::sync::Arc;
use thiserror::Error;
use tokio::runtime::Handle as TokioHandle;
use tracing::info;
use super::message::{BulkRx, IndexerRx, IndexerTx};
use crate::ingest::mailbox::ShardedSender;
use crate::ingest::stream::Commit;
use crate::resolver::{NoSigningKeyError, ResolverError};
use crate::state::AppState;
#[derive(Debug, Diagnostic, Error)]
pub(crate) enum IngestError {
#[error("{0}")]
Generic(miette::Report),
#[error(transparent)]
#[diagnostic(transparent)]
Resolver(#[from] ResolverError),
#[error(transparent)]
#[diagnostic(transparent)]
NoSigningKey(#[from] NoSigningKeyError),
}
impl From<miette::Report> for IngestError {
fn from(report: miette::Report) -> Self {
IngestError::Generic(report)
}
}
#[derive(Debug)]
pub(crate) enum RepoProcessResult<'c> {
Ok,
// needs backfill and carries the triggering commit to buffer
NeedsBackfill(&'c Commit<'c>),
}
pub struct FirehoseWorker {
pub(crate) state: Arc<AppState>,
pub(crate) rxs: Vec<IndexerRx>,
pub(crate) bulk_rxs: Vec<BulkRx>,
}
impl FirehoseWorker {
pub fn new(state: Arc<AppState>, num_shards: usize) -> (IndexerTx, Self) {
let (inner, rxs) = ShardedSender::channel(num_shards);
let (bulk, bulk_rxs) = ShardedSender::channel(num_shards);
(
IndexerTx { inner, bulk },
Self {
state,
rxs,
bulk_rxs,
},
)
}
/// one shard on its own thread, which stops once every sender is dropped
#[cfg(test)]
pub(crate) fn spawn_one_shard(
state: Arc<AppState>,
) -> (IndexerTx, std::thread::JoinHandle<()>) {
let (tx, mut worker) = Self::new(state.clone(), 1);
let rx = worker.rxs.pop().expect("single shard receiver");
let bulk_rx = worker.bulk_rxs.pop().expect("single shard bulk receiver");
let runtime = TokioHandle::current();
let shard = std::thread::spawn(move || Self::shard(0, 1, rx, bulk_rx, state, runtime));
(tx, shard)
}
/// [`Self::spawn_one_shard`] as the shard the api and the crawler hand their work to,
/// until the returned guard drops
#[cfg(test)]
pub(crate) fn serve_one_shard(state: Arc<AppState>) -> ServedShard {
Self::serve_shards(state, 1, None)
}
/// `num_shards` shards the api and the crawler hand their work to, until the returned guard
/// drops. `gone` names one that never runs, so whatever is sent to it fails
#[cfg(test)]
pub(crate) fn serve_shards(
state: Arc<AppState>,
num_shards: usize,
gone: Option<usize>,
) -> ServedShard {
let (tx, worker) = Self::new(state.clone(), num_shards);
let runtime = TokioHandle::current();
for (i, (rx, bulk_rx)) in worker.rxs.into_iter().zip(worker.bulk_rxs).enumerate() {
if gone == Some(i) {
continue;
}
let (state, runtime) = (state.clone(), runtime.clone());
std::thread::spawn(move || Self::shard(i, num_shards, rx, bulk_rx, state, runtime));
}
state.indexer.lock().replace(tx);
ServedShard(state)
}
pub fn run(self, handle: TokioHandle) -> Result<()> {
let num_shards = self.rxs.len();
let (exit_tx, exit_rx) = std::sync::mpsc::channel();
for (i, (rx, bulk_rx)) in self.rxs.into_iter().zip(self.bulk_rxs).enumerate() {
let state = Arc::clone(&self.state);
let handle = handle.clone();
let exit_tx = exit_tx.clone();
std::thread::Builder::new()
.name(format!("ingest-shard-{i}"))
.spawn(move || {
let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
Self::shard(i, num_shards, rx, bulk_rx, state, handle);
}));
let _ = exit_tx.send((i, res));
})
.into_diagnostic()?;
}
drop(exit_tx);
info!(num = num_shards, "started shards");
match exit_rx.recv() {
Ok((id, Ok(()))) => Err(miette::miette!("firehose worker shard {id} shut down")),
Ok((id, Err(_))) => Err(miette::miette!("firehose worker shard {id} panicked")),
Err(_) => Err(miette::miette!("firehose worker shards shut down")),
}
}
}
/// closes the shard [`FirehoseWorker::serve_one_shard`] started once it drops, which stops it
#[cfg(test)]
pub(crate) struct ServedShard(Arc<AppState>);
#[cfg(test)]
impl Drop for ServedShard {
fn drop(&mut self) {
self.0.indexer.lock().take();
}
}