tantivy-cache 0.7.0

use std::{
    io,
    ops::Range,
    path::{Path, PathBuf},
    sync::Arc,
};

use async_trait::async_trait;
use block_on_place::HandleExt;
use deadpool_redis::{Connection, Pool, redis::AsyncCommands};
use eyre::{OptionExt, Result};
use gxhash::GxBuildHasher;
use scc::HashIndex;
use tantivy::{
    Directory, HasLen,
    directory::{
        DirectoryLock, FileHandle, FileSlice, Lock, WatchCallback, WatchHandle, WritePtr,
        error::{DeleteError, LockError, OpenReadError, OpenWriteError},
    },
};
use tantivy_common::OwnedBytes;
use tokio::runtime::Handle;

use self::{error::WrapIoErrorExt, footer::Footer};

pub use self::warmer::Warmer;

mod error;
mod footer;
mod keys;
mod warmer;

#[cfg(test)]
mod tests;

/// A [`Directory`] implementation wrapping another one and caching some data using
/// [`deadpool_redis`].
#[derive(Clone, Debug)]
pub struct CachingDirectory<D> {
    rt: Handle,
    redis: Pool,
    dir: D,

    // In-memory cache used to not fetch footers more than once and to not hold more
    // than one instance of a file's footer in-memory.
    cache: Cache,
}

type Cache = Arc<HashIndex<PathBuf, Footer, GxBuildHasher>>;

#[derive(Debug)]
struct CachingHandle {
    file: FileSlice,
    footer: Footer,

    rt: Handle,
    redis: Pool,
    path: PathBuf,
}

impl<D> CachingDirectory<D> {
    /// Creates a new [`CachingDirectory`] from the given directory and pool.
    ///
    /// The provided [`Directory`] must return a [`FileHandle`] that implements
    /// [`read_bytes_async()`][1] without panicking.
    ///
    /// This must be called from within a `tokio` context.
    ///
    /// [1]: FileHandle::read_bytes_async()
    pub fn new(dir: D, redis: Pool) -> Self {
        let rt = Handle::current();
        let cache = Arc::default();

        Self {
            rt,
            redis,
            cache,
            dir,
        }
    }

    async fn open(&self, path: impl Into<PathBuf>, file: FileSlice) -> Result<CachingHandle> {
        let path = path.into();

        loop {
            if let Some(footer) = self.cache.get_async(&path).await {
                let footer = footer.get().clone();
                let rt = self.rt.clone();
                let redis = self.redis.clone();

                return Ok(CachingHandle {
                    file,
                    footer,
                    rt,
                    redis,
                    path,
                });
            }

            let footer = if let Some(offset) = offset(&path, &self.redis).await {
                Footer::with_offset(offset)
            } else {
                let footer = Footer::read(&path, &file).await?;
                update(&path, &footer, &self.redis).await;

                footer
            };

            if self
                .cache
                .insert_async(path.clone(), footer.clone())
                .await
                .is_err()
            {
                // A footer was inserted by another task – fetch it so that we don't end up with
                // more than one footer living for the same file.
                continue;
            }

            let rt = self.rt.clone();
            let redis = self.redis.clone();

            return Ok(CachingHandle {
                file,
                footer,
                rt,
                redis,
                path,
            });
        }
    }
}

impl<D: Directory + Clone> CachingDirectory<D> {
    /// Creates a [`Warmer`] that allows warming this directory's cache.
    ///
    /// The warmer defaults to the host's available parallelism; configure it with
    /// [`Warmer::with_concurrency`].
    ///
    /// [1]: tantivy::Index::open
    pub fn warmer(&self) -> Warmer<D> {
        Warmer::new(self.clone())
    }

    /// Warms the cache for a single file.
    ///
    /// Opens the file and loads its footer into both the in-memory cache and Redis, so
    /// that a later read of the footer region is served entirely from cache instead of
    /// from the underlying [`Directory`].
    ///
    /// Non-existent files are silently ignored.
    pub(crate) async fn warm(&self, path: &Path) -> Result<()> {
        let dir = self.dir.clone();
        let owned = path.to_path_buf();

        // `open_read` is blocking, so run it off the async runtime to avoid
        // serialising concurrent warms on the runtime's worker threads.
        let opened = self
            .rt
            .spawn_blocking(move || dir.open_read(&owned))
            .await?;

        let file = match opened {
            Ok(file) => file,
            Err(OpenReadError::FileDoesNotExist(_)) => return Ok(()),
            Err(error) => return Err(error.into()),
        };

        // Caches the footer offset (in Redis and in-memory).
        let handle = self.open(path, file).await?;

        // Eagerly loads the footer data into memory (and Redis), so later reads are served
        // without touching Redis or the directory.
        handle.footer().await;

        Ok(())
    }
}

impl CachingHandle {
    /// Fetches the file's footer.
    ///
    /// If the footer has already been fetched, this simply returns it. If it is stored
    /// in Redis, it fetches it from there. Otherwsie, it reads it and stores it in
    /// Redis.
    async fn footer(&self) -> Option<OwnedBytes> {
        let fetch = async {
            if let Some(data) = footer(&self.path, &self.redis).await {
                return Ok(OwnedBytes::new(data));
            }

            let footer = Footer::read(&self.path, &self.file).await?;
            update(&self.path, &footer, &self.redis).await;

            footer.data().ok_or_eyre("missing data")
        };

        match self.footer.get_or_fetch(fetch).await {
            Ok(footer) => Some(footer),
            Err(error) => {
                tracing::warn!("Failed to fetch footer: {error:?}");

                None
            }
        }
    }
}

impl<D: Clone + Directory> Directory for CachingDirectory<D> {
    #[inline]
    fn get_file_handle(&self, path: &Path) -> Result<Arc<dyn FileHandle>, OpenReadError> {
        let file = self.dir.open_read(path)?;
        let handle = self
            .rt
            .block_on_place(self.open(path, file))
            .map_err(OpenReadError::wrapper(path))?;

        Ok(Arc::new(handle))
    }

    #[inline]
    fn delete(&self, path: &Path) -> Result<(), DeleteError> {
        // TODO(MLB): remove from cache
        self.dir.delete(path)
    }

    #[inline]
    fn exists(&self, path: &Path) -> Result<bool, OpenReadError> {
        // TODO(MLB): check the cache
        self.dir.exists(path)
    }

    #[inline]
    fn open_read(&self, path: &Path) -> Result<FileSlice, OpenReadError> {
        self.get_file_handle(path).map(FileSlice::new)
    }

    #[inline]
    fn open_write(&self, path: &Path) -> Result<WritePtr, OpenWriteError> {
        // TODO(MLB): wrap and insert into the cache
        self.dir.open_write(path)
    }

    #[inline]
    fn atomic_read(&self, path: &Path) -> Result<Vec<u8>, OpenReadError> {
        // TODO(MLB): read from cache
        self.dir.atomic_read(path)
    }

    #[inline]
    fn atomic_write(&self, path: &Path, data: &[u8]) -> io::Result<()> {
        // TODO(MLB): write to cache
        self.dir.atomic_write(path, data)
    }

    #[inline]
    fn sync_directory(&self) -> io::Result<()> {
        self.dir.sync_directory()
    }

    #[inline]
    fn watch(&self, cb: WatchCallback) -> tantivy::Result<WatchHandle> {
        self.dir.watch(cb)
    }

    #[inline]
    fn acquire_lock(&self, lock: &Lock) -> Result<DirectoryLock, LockError> {
        self.dir.acquire_lock(lock)
    }
}

#[async_trait]
impl FileHandle for CachingHandle {
    fn read_bytes(&self, range: Range<usize>) -> io::Result<OwnedBytes> {
        if range.end <= self.footer.offset {
            return self.file.read_bytes_slice(range);
        }

        let Some(footer) = self.rt.block_on_place(self.footer()) else {
            return self.file.read_bytes_slice(range);
        };

        if range.start >= self.footer.offset {
            let start = range.start - self.footer.offset;
            let end = range.end - self.footer.offset;
            let slice = footer.slice(start..end);

            return Ok(slice);
        }

        let start = range.start;
        let end = self.footer.offset;
        let data = self.file.read_bytes_slice(start..end)?;

        let start = 0;
        let end = range.end - self.footer.offset;
        let footer = footer.slice(start..end);

        let mut combined = Vec::with_capacity(data.len() + footer.len());
        combined.extend_from_slice(data.as_ref());
        combined.extend_from_slice(footer.as_ref());

        Ok(OwnedBytes::new(combined))
    }

    async fn read_bytes_async(&self, range: Range<usize>) -> io::Result<OwnedBytes> {
        if range.end <= self.footer.offset {
            return self.file.read_bytes_slice_async(range).await;
        }

        let Some(footer) = self.footer().await else {
            return self.file.read_bytes_slice_async(range).await;
        };

        if range.start >= self.footer.offset {
            let start = range.start - self.footer.offset;
            let end = range.end - self.footer.offset;
            let slice = footer.slice(start..end);

            return Ok(slice);
        }

        let start = range.start;
        let end = self.footer.offset;
        let data = self.file.read_bytes_slice_async(start..end).await?;

        let start = 0;
        let end = range.end - self.footer.offset;
        let footer = footer.slice(start..end);

        let mut combined = Vec::with_capacity(data.len() + footer.len());
        combined.extend_from_slice(data.as_ref());
        combined.extend_from_slice(footer.as_ref());

        Ok(OwnedBytes::new(combined))
    }
}

impl HasLen for CachingHandle {
    #[inline]
    fn len(&self) -> usize {
        self.file.len()
    }

    #[inline]
    fn is_empty(&self) -> bool {
        self.file.is_empty()
    }
}

/// Gets a connection from the given Redis pool.
///
/// If an error occurs, this returns `None` and logs the error message.
async fn conn(redis: &Pool) -> Option<Connection> {
    match redis.get().await {
        Ok(conn) => Some(conn),
        Err(error) => {
            tracing::warn!("Failed to get Redis connection: {error:?}");
            None
        }
    }
}

/// Fetches the footer offset for the given path from Redis.
///
/// If an error occurs, this returns `None` and logs the error message.
async fn offset(path: &Path, redis: &Pool) -> Option<usize> {
    let key = keys::offset(path).ok()?;
    let mut conn = conn(redis).await?;

    match conn.get(&key).await {
        Ok(offset) => offset,
        Err(error) => {
            tracing::warn!("Failed to fetch footer offset: {error:?}");
            None
        }
    }
}

/// Fetches the footer data for the given path from Redis.
///
/// If an error occurs, this returns `None` and logs the error message.
async fn footer(path: &Path, redis: &Pool) -> Option<Vec<u8>> {
    let key = keys::footer(path).ok()?;
    let mut conn = conn(redis).await?;

    match conn.get(&key).await {
        Ok(offset) => offset,
        Err(error) => {
            tracing::warn!("Failed to fetch footer: {error:?}");
            None
        }
    }
}

/// Updates the cached data for the given footer in Redis.
///
/// If an error occurs, this returns `None` and logs the error message.
async fn update(path: &Path, footer: &Footer, redis: &Pool) {
    let Some(mut conn) = conn(redis).await else {
        return;
    };

    let Ok(key) = keys::offset(path) else { return };
    match conn.set(&key, footer.offset).await {
        Ok(()) => (),
        Err(error) => {
            tracing::warn!("Failed to update footer offset: {error:?}");
        }
    }

    let Ok(key) = keys::footer(path) else { return };
    let Some(data) = footer.data() else { return };
    match conn.set(&key, data.as_slice()).await {
        Ok(()) => (),
        Err(error) => {
            tracing::warn!("Failed to update footer: {error:?}");
        }
    }
}