tantivy-cache 0.7.0

tantivy caching for faster searching
Documentation
use std::{cmp, fmt::Debug, path::PathBuf, thread};

use eyre::{OptionExt, Result};
use tantivy::{Directory, Index};
use tokio::task::{self, JoinSet};

use crate::CachingDirectory;

/// A parallel cache warmer for a [`CachingDirectory`].
///
/// `tantivy` opens segments, and the component files within each segment,
/// *sequentially* when building its [`SegmentReader`s][1]. Over a high-latency
/// [`Directory`] (e.g. one backed by object storage) this makes opening a searcher
/// slow.
///
/// A [`Warmer`] pre-populates the cache for many files *in parallel*, so that the
/// subsequent sequential segment opening done by `tantivy` reads every footer from
/// the cache rather than from the slow directory.
#[derive(Clone, Debug)]
pub struct Warmer<D> {
    dir: CachingDirectory<D>,

    /// The maximum number of files warmed at the same time. Always `>= 1`.
    concurrency: usize,
}

impl<D: Directory + Clone> Warmer<D> {
    /// Builds a warmer for `dir`. This is the implementation of
    /// [`CachingDirectory::warmer`].
    pub(crate) fn new(dir: CachingDirectory<D>) -> Self {
        let concurrency = thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1);

        Self { dir, concurrency }
    }

    /// Sets the maximum number of files warmed concurrently (clamped to `>= 1`).
    pub fn with_concurrency(mut self, concurrency: usize) -> Self {
        self.concurrency = cmp::max(concurrency, 1);
        self
    }

    /// The maximum number of files this warmer warms concurrently.
    pub fn concurrency(&self) -> usize {
        self.concurrency
    }

    /// Warms the cache for every file of every searchable segment of `index`.
    ///
    /// The files of each segment are warmed in parallel, up to [`concurrency`][1] files
    /// at a time.
    ///
    /// `index` must have been opened on the [`CachingDirectory`] this warmer was
    /// created from (or a clone of it), so that warming populates the very cache
    /// `index` reads from.
    ///
    /// [1]: 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
    }

    /// Warms the cache for the given `paths`.
    ///
    /// The files are warmed in parallel, up to [`concurrency`][1] files at a time.
    ///
    /// Non-existent files are ignored. The first error encountered is returned,
    /// aborting the remaining work.
    ///
    /// [1]: Self::concurrency()
    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 });
        }

        // Propagate the first join error (panic) or warming error.
        while let Some(joined) = tasks.join_next().await {
            joined??;
        }

        Ok(())
    }
}