use std::{cmp, fmt::Debug, path::PathBuf, thread};
use eyre::{OptionExt, Result};
use tantivy::{Directory, Index};
use tokio::task::{self, JoinSet};
use crate::CachingDirectory;
#[derive(Clone, Debug)]
pub struct Warmer<D> {
dir: CachingDirectory<D>,
concurrency: usize,
}
impl<D: Directory + Clone> Warmer<D> {
pub(crate) fn new(dir: CachingDirectory<D>) -> Self {
let concurrency = thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
Self { dir, concurrency }
}
pub fn with_concurrency(mut self, concurrency: usize) -> Self {
self.concurrency = cmp::max(concurrency, 1);
self
}
pub fn concurrency(&self) -> usize {
self.concurrency
}
pub async fn index(&self, index: &Index) -> Result<()> {
let metas = task::block_in_place(|| index.searchable_segment_metas())?;
let paths = metas.into_iter().flat_map(|meta| meta.list_files());
self.files(paths).await
}
pub async fn files<I>(&self, paths: I) -> Result<()>
where
I: IntoIterator,
I::Item: Into<PathBuf>,
{
let mut tasks = JoinSet::new();
for path in paths {
let path = path.into();
if tasks.len() >= self.concurrency {
tasks.join_next().await.ok_or_eyre("no task running")???;
}
let dir = self.dir.clone();
tasks.spawn(async move { dir.warm(&path).await });
}
while let Some(joined) = tasks.join_next().await {
joined??;
}
Ok(())
}
}