qdrant-edge 0.8.0

use std::fmt::Debug;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use super::config::DiskCacheConfig;
use super::file::{DiskCache, State};
use super::pipeline::REMOTE_READ_ALIGNMENT;
use super::{DiskCacheRemote, block_aligned_fetch};
use crate::common::generic_consts::Sequential;
use crate::common::universal_io::simple_disk_cache::REMOTE_OPEN_OPTIONS;
use crate::common::universal_io::simple_disk_cache::local_state::LocalState;
use crate::common::universal_io::{
    ListedFile, OpenExtra, OpenOptions, OwnedPipeline, Populate, UioResult, UniversalIoError,
    UniversalRead, UniversalReadFileOps, UniversalReadFs,
};

/// Construction context for [`DiskCacheFs`]: carries the
/// remote/local-mirror layout config alongside the remote filesystem's own
/// construction context. Always constructed explicitly — there is no
/// global / default lookup.
pub struct DiskCacheFsContext<C> {
    pub config: Arc<DiskCacheConfig>,
    pub remote: C,
}

#[derive(Default, Debug)]
pub struct DiskCacheFsOpenExtra<RemoteExtra: OpenExtra> {
    /// Extra options passed to the remote
    remote_extra: RemoteExtra,
    /// The length of the file, if known
    known_len: Option<u64>,
}

impl<RemoteExtra: OpenExtra> OpenExtra for DiskCacheFsOpenExtra<RemoteExtra> {
    fn with_prevent_caching(self, _prevent_caching: bool) -> Self {
        self
    }

    fn with_known_len(self, known_len: u64) -> Self {
        Self {
            remote_extra: self.remote_extra,
            known_len: Some(known_len),
        }
    }
}

/// Filesystem handle for the simple disk cache. Carries the remote-side
/// filesystem used to fetch missing blocks plus the shared
/// [`DiskCacheConfig`] that maps remote paths to local mirror paths.
///
/// `DiskCacheFs` is parameterized by the remote *file* type `R`; the
/// underlying remote filesystem is `R::Fs`. This indirection lets the
/// `UniversalRead::Fs = DiskCacheFs<R>` constraint on
/// [`DiskCache<R>`] line up without an extra generic parameter on the
/// file type.
pub struct DiskCacheFs<R>
where
    R: UniversalRead,
{
    config: Arc<DiskCacheConfig>,
    remote_fs: R::Fs,
}

impl<R> Clone for DiskCacheFs<R>
where
    R: UniversalRead,
    R::Fs: Clone,
{
    fn clone(&self) -> Self {
        let Self { config, remote_fs } = self;
        Self {
            config: config.clone(),
            remote_fs: remote_fs.clone(),
        }
    }
}

impl<R> Debug for DiskCacheFs<R>
where
    R: UniversalRead,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self { config, remote_fs } = self;
        f.debug_struct("DiskCacheFs")
            .field("config", config)
            .field("remote_fs", remote_fs)
            .finish()
    }
}

impl<R: UniversalRead> DiskCacheFs<R> {
    fn open_remote(
        &self,
        path: impl AsRef<Path>,
        extra: <R::Fs as UniversalReadFs>::OpenExtra,
    ) -> UioResult<R> {
        self.remote_fs
            .open(path.as_ref(), REMOTE_OPEN_OPTIONS, extra)
    }
}

impl<R> UniversalReadFileOps for DiskCacheFs<R>
where
    R: UniversalRead,
{
    type ContextConfig = DiskCacheFsContext<<R::Fs as UniversalReadFileOps>::ContextConfig>;

    fn from_context(ctx: Self::ContextConfig) -> UioResult<Self> {
        let DiskCacheFsContext { config, remote } = ctx;
        Ok(Self {
            config,
            remote_fs: R::Fs::from_context(remote)?,
        })
    }

    fn list_files(&self, prefix_path: &Path) -> UioResult<Vec<ListedFile>> {
        self.remote_fs.list_files(prefix_path)
    }

    fn exists(&self, path: &Path) -> UioResult<bool> {
        self.remote_fs.exists(path)
    }
}

// Deliberately no `UniversalWriteFileOps` impl: the disk cache is strictly
// read-only, at the filesystem level as much as at the file level (see the
// `assert_not_impl_any!` on `DiskCache`). Mutations — creating, removing and
// appending to files — go straight to the backing storage, whose handle the
// caller holds anyway to build this one.

/// Make the mirror path unique per open, so concurrently-alive [`DiskCache`]
/// instances for the same remote path never share (and truncate) each other's
/// local mirror. This lets callers refresh a file by opening a fresh handle
/// while the old one is still alive.
///
/// The name carries no state across opens: a mirror is truncated when its
/// [`LocalState`] materializes and removed when its `DiskCache` is dropped.
fn unique_local_path(mut path: PathBuf) -> PathBuf {
    static NEXT_MIRROR_ID: AtomicU64 = AtomicU64::new(0);
    let id = NEXT_MIRROR_ID.fetch_add(1, Ordering::Relaxed);
    // Process id disambiguates processes sharing a local cache dir.
    path.as_mut_os_string()
        .push(format!(".{:x}-{id:x}", std::process::id()));
    path
}

impl<R> UniversalReadFs for DiskCacheFs<R>
where
    R: DiskCacheRemote,
{
    type File = DiskCache<R>;
    type OpenExtra = DiskCacheFsOpenExtra<<R::Fs as UniversalReadFs>::OpenExtra>;

    fn open(
        &self,
        path: impl AsRef<Path>,
        options: OpenOptions,
        extra: Self::OpenExtra,
    ) -> UioResult<DiskCache<R>> {
        // The cache is strictly read-only: appends (and any other writes)
        // must go directly to the backing storage, not through the cache.
        if options.writeable {
            return Err(UniversalIoError::Uninitialized {
                description:
                    "DiskCache only supports immutable files, writeable option is not allowed"
                        .to_string(),
            });
        }

        let remote_extra = extra.remote_extra.with_prevent_caching(true);
        let local_path = unique_local_path(self.config.local_path_for(path.as_ref())?);

        let populate = if crate::common::low_memory::low_memory_mode().skip_populate() {
            Populate::No
        } else {
            options.populate
        };

        let state = match (extra.known_len, populate) {
            (None, Populate::Auto | Populate::No) => State::Uninit,
            // We know file length, initialize local state.
            (Some(len), Populate::Auto | Populate::No) => {
                let remote = self.open_remote(path.as_ref(), remote_extra.clone())?;

                let local = LocalState::new(&local_path, len, options)?;

                State::ready(remote, local)
            }
            // Even if we know length, we don't need it to do `schedule_whole`.
            (None | Some(_), Populate::Blocking | Populate::PreferBackground) => {
                let remote = self.open_remote(path.as_ref(), remote_extra.clone())?;

                let mut pipeline = OwnedPipeline::new(remote)?;

                // FIXME: check `can_schedule` in a loop first
                pipeline.schedule_whole((), 0)?;

                State::OpenPrefill { pipeline }
            }
            // Special case of no known length and empty range. Don't initialize.
            (None, Populate::Partial(range)) if range.into_byte_range::<u8>().is_empty() => {
                State::Uninit
            }
            // Schedule a partial block-aligned read.
            (known_len, Populate::Partial(range)) => {
                let remote = self.open_remote(path.as_ref(), remote_extra.clone())?;

                // `Partial` must create the local file, so we need the length up front.
                let file_len = match known_len {
                    Some(len) => len,
                    None => remote.len::<u8>()?,
                };

                let requested_byte_range = range.into_byte_range::<u8>();

                if let Some((blocks_range, byte_range)) =
                    block_aligned_fetch(requested_byte_range.clone(), file_len)
                {
                    let mut pipeline = OwnedPipeline::new(remote)?;

                    // FIXME: check `can_schedule` in a loop first
                    pipeline.schedule::<Sequential>(
                        blocks_range,
                        byte_range,
                        REMOTE_READ_ALIGNMENT,
                    )?;

                    State::PartialPrefill {
                        pipeline,
                        len: file_len,
                    }
                } else {
                    // empty byte range, just initialize with length.
                    let local = LocalState::new(&local_path, file_len, options)?;
                    State::ready(remote, local)
                }
            }
        };

        let cache = DiskCache::new(
            self.remote_fs.clone(),
            remote_extra,
            path.as_ref(),
            local_path,
            options,
            state,
        );

        if matches!(populate, Populate::Blocking) {
            // Force the prefill to resolve before returning.
            cache.state()?;
        }

        Ok(cache)
    }
}