boreholeio 0.1.0

mod reference_counting_hash_map;
mod resource;

#[cfg(test)]
mod tests;

use crate::{
    schema::array_uri_reference::{AbsoluteUri, ArrayUriReference, Scheme, Uri},
    strided_array_file::{extract_array_index, ArrayProxy, Error, StridedArrayFile},
};
use reference_counting_hash_map::ReferenceCountingHashMap;
use resource::FileResource;
use serde::Deserialize;
use std::{
    collections::HashMap,
    fs::{create_dir, File},
    io::BufWriter,
    mem::transmute,
    path::{Path, PathBuf},
    sync::{Arc, Mutex, OnceLock, Weak},
};
use uuid::Uuid;
use zip::ZipArchive;

type FileArchive = ZipArchive<File>;
type SharedFileArchive = Arc<Mutex<FileArchive>>;

/// Configures `ArrayCache`'s temporary directory.
///
/// This method must be called at the very beginning of the program because
/// otherwise, if `ArrayCache` has already been instantiated, the specified
/// path won't be applied.
///
/// If this method is not called, the default value from `std::env::temp_dir()`
/// will be used instead.
///
/// Returns `true` if `temporary_directory` points to a valid directory and the
/// configuration succeeded, or `false` otherwise
pub fn configure(temporary_directory: PathBuf) -> bool {
    if !temporary_directory.is_dir() {
        return false;
    }
    INSTANCE
        .set(Arc::new(Mutex::new(ArrayCache {
            resource_views: ReferenceCountingHashMap::new(),
            file_resources: Vec::new(),
            resources: Vec::new(),
            file_archives: HashMap::new(),
            temporary_directory,
        })))
        .is_ok()
}

static INSTANCE: OnceLock<Arc<Mutex<ArrayCache>>> = OnceLock::new();

/// A singleton class that allows associating `schema::Uri` with a read-only
/// `ArrayProxy` backed by `StridedArrayFile` or (in-memory) `TypedArray`. This
/// class manages lifetimes under the hood: returned data views borrow their
/// calling `Uri` objects and are guaranteed to be valid for (at least) as long
/// as these `Uri` objects live. Caching the same data but with different
/// headers (such as slicing and reshaping `ArrayCache`'s results) is also
/// possible: `ArrayCache` can relate already cached resources to their views,
/// allowing it to avoid redundant data copying.
#[derive(Debug)]
pub(crate) struct ArrayCache {
    /// All the cached URIs.
    resource_views: ReferenceCountingHashMap<Uri, Arc<dyn resource::View<'static>>>,
    /// This is an auxiliary member that keeps track of all unique Strided
    /// Array Files we have cached. It is used as a LUT every time we add or
    /// remove a file view to relate this view with the corresponding file
    /// resource.
    file_resources: Vec<Weak<StridedArrayFile<'static>>>,
    /// Same as `files` - keeps track of (not only files but) all the resources
    /// we have. Used as a LUT every time we add or remove a resource.
    resources: Vec<Weak<dyn resource::Resource + 'static>>,
    /// Similarly to `file_resources`, this one keeps track of all currently
    /// opened Borehole Data Archives. It is used whenever a resource is being
    /// created or deleted:
    ///  - on creation, we specify that the Strided Array File has to be
    ///    extracted from the BDA file first;
    ///  - on deletion, archives whose reference counter has reached `1` are
    ///    removed, along with their temporary directories.
    ///
    /// However, there are some key differences:
    /// 1. Each archive is associated with its extraction directory, which is
    ///    used as the key in this variable.
    /// 2. This object itself increases the reference counter (`Arc` is used
    ///    instead of `WeakPtr`) because archives must exist before resources,
    ///    as resources are created from them.
    ///
    /// Compile-time type requirements: why `Arc<Mutex>`?
    /// 1. Being a thread-safe singleton, `ArrayCache` must be `Send + Sync`,
    ///    which in turn requires all its fields to be `Send + Sync`.
    /// 2. When multiple files from the same archive are needed, `ZipArchive`
    ///    is shared via `Arc` (not just `Rc` because of thread-safety).
    /// 3. Extracting a file requires `&mut ZipArchive`, so `Mutex` is used to
    ///    ensure safe access.
    file_archives: HashMap<AbsoluteUri, SharedFileArchive>,
    /// Specifies where to extract file archives, if any.
    temporary_directory: PathBuf,
}

/// This trait indicates that a data structure implementing it allows its URI
/// fields to be either relative references or absolute URIs.
pub(crate) trait UriResolvable: for<'de> Deserialize<'de> {
    /// Sets every `Uri.base` to the given value (leaving `Uri.relative`
    /// untouched).
    fn set_base_for_relative_uris(&mut self, base_uri: &AbsoluteUri);
}

impl ArrayCache {
    /// Increases reference count of the given URI. Does nothing to malformed
    /// URIs.
    pub(crate) fn register(uri: &ArrayUriReference) {
        // Do nothing for malformed URIs.
        let resolved = match uri.resolved() {
            Ok(resolved) => resolved,
            Err(_) => return,
        };

        let instance = Self::instance();
        let mut cache = instance.lock().unwrap();

        // If URI is already cached, increment its reference count.
        if cache.resource_views.increment(&resolved).is_some() {
        }
        // Otherwise, if URI points to a Strided Array File, register it.
        else if let Some(file_view) = cache.create_file_view(uri) {
            cache.resource_views.insert(resolved, file_view).unwrap();
        }
    }

    /// Decreases reference count of the given URI, if it is already cached.
    /// Does nothing to non-cached or malformed URIs.
    pub(crate) fn deregister(uri: &ArrayUriReference) {
        // Do nothing for relative URIs.
        let absolute = match uri.resolved() {
            Ok(value) => value,
            _ => return,
        };

        let instance = Self::instance();
        let mut cache = instance.lock().unwrap();

        // Do nothing unless reference counter reaches `0`.
        if cache.resource_views.decrement(&absolute) != Some(0) {
            return;
        }
        // Remove the newly destroyed resource from the LUTs.
        cache.resources.retain(|r| r.strong_count() > 0);
        cache.file_resources.retain(|r| r.strong_count() > 0);
        // If the resource is backed by an archive which is no longer in use,
        // remove the archive.
        cache.consider_removing_archive(uri);
    }

    /// Returns the requested array from cache. Errors out if the URI is
    /// malformed or not cached, except for files, which are read and added to
    /// cache inside of this method.
    pub(crate) fn get(uri: &ArrayUriReference) -> Result<ArrayProxy<'_>, Error> {
        let resolved = uri.resolved()?;

        let instance = Self::instance();
        let mut cache = instance.lock().unwrap();

        // If `uri` points to a file which does not yet exist, tries extracting
        // it from an archive, if any.
        cache.consider_fetching_file(uri)?;

        let proxy = cache.resource_views.get(&resolved)?.proxy()?;
        // This is safe because nothing outlives `<'static>`.
        unsafe { Ok(transmute::<ArrayProxy<'static>, ArrayProxy<'_>>(proxy)) }
    }

    /// Associates given `proxy` with a new URI.
    pub(crate) fn insert(proxy: ArrayProxy<'_>) -> Result<ArrayUriReference, Error> {
        let uuid = Uuid::new_v4();
        let uri: ArrayUriReference = format!("memory:{uuid}").as_str().try_into()?;
        let resolved = uri.resolved()?;

        let instance = Self::instance();
        let mut cache = instance.lock().unwrap();

        // Check if `proxy` points to one of the already cached resources.
        for resource in &cache.resources {
            let resource = resource.upgrade().unwrap();
            if resource.is_parent_of(proxy.clone())? {
                // It's guaranteed that `proxy` has `'static` lifetime because
                // it points to an already cached resource (which is `'static`).
                let proxy = unsafe { transmute::<ArrayProxy<'_>, ArrayProxy<'static>>(proxy) };
                let resource_view = resource::create_view(proxy, resource)?;
                cache.resource_views.insert(resolved, resource_view)?;
                return Ok(uri);
            }
        }

        // If `proxy` does not point to any of the already cached resources,
        // copy it to a new "in-memory" resource. Note that this new resource
        // points to a different memory location!
        let array_copy = proxy.to_owned()?;
        let resource_view = resource::create_new_array(array_copy)?;
        let resource = Arc::downgrade(&resource_view.resource());
        cache.resources.push(resource);
        cache.resource_views.insert(resolved, resource_view)?;
        Ok(uri)
    }

    /// Generates a temporary subdirectory where the the given `archive` will
    /// be lazily extracted to. Any `file:///` URI whose base equals the
    /// generated directory, will be associated with `archive` in the future.
    pub(crate) fn insert_archive(archive: FileArchive) -> Result<AbsoluteUri, Error> {
        let archive = Arc::new(Mutex::new(archive));

        let instance = Self::instance();
        let mut cache = instance.lock().unwrap();

        let path = cache.temporary_directory().join(Uuid::new_v4().to_string());
        let uri = AbsoluteUri::try_from(path.as_path())?;

        cache.file_archives.insert(uri.clone(), archive);
        Ok(uri)
    }

    fn instance() -> Arc<Mutex<Self>> {
        INSTANCE
            .get_or_init(|| {
                Arc::new(Mutex::new(ArrayCache {
                    resource_views: ReferenceCountingHashMap::new(),
                    file_resources: Vec::new(),
                    resources: Vec::new(),
                    file_archives: HashMap::new(),
                    temporary_directory: std::env::temp_dir(),
                }))
            })
            .clone()
    }

    /// Returns a temporary root directory, where multiple subdirectories may
    /// be created in the future (one for each Borehole Data Archive).
    fn temporary_directory(&self) -> &Path {
        &self.temporary_directory
    }

    /// If URI base points to the temporary directory of a cached archive,
    /// extracts the file from that archive if it has not already been
    /// extracted. Otherwise, does nothing.
    fn consider_fetching_file(&mut self, uri: &ArrayUriReference) -> Result<(), Error> {
        // Do nothing if URI does not point to a file.
        let destination_path: PathBuf = if uri.scheme() == Some(Scheme::FILE) {
            let (absolute, _fragment) = uri.resolved()?.to_absolute_and_fragment();
            (&absolute).try_into()?
        } else {
            return Ok(());
        };
        // Do nothing if the file already exists.
        if destination_path.is_file() {
            return Ok(());
        }
        // Do nothing if URI does not have a base component (because it points
        // to the extraction directory).
        let base = match uri.base() {
            Some(value) => value,
            None => return Ok(()),
        };
        // Do nothing if the archive is not cached.
        let mut archive = match self.file_archives.get(base) {
            Some(value) => value.lock().unwrap(),
            None => return Ok(()),
        };
        // Find the file to be extracted from the archive.
        let source_path = {
            let mut relative = uri.relative().to_owned();
            relative.set_fragment(None);
            relative
        };
        let mut source_file = archive
            .by_name(source_path.as_str())
            .map_err(|e| e.to_string())?;
        // Make sure destination directory exists.
        let destination_directory = destination_path.parent().unwrap();
        if !destination_directory.is_dir() {
            create_dir(destination_directory)
                .map_err(|e| format!("Failed to create archive subdirectory {e}"))?;
        }
        // Copy the contents to the new file.
        let destination_file = File::create(destination_path).map_err(|e| e.to_string())?;
        let mut writer = BufWriter::new(destination_file);
        let _bytes_copied =
            std::io::copy(&mut source_file, &mut writer).map_err(|e| e.to_string())?;
        Ok(())
    }

    /// If URI is backed by an archive that is no longer in use, tries
    /// deleting its extraction directory.
    fn consider_removing_archive(&mut self, uri: &ArrayUriReference) {
        // Do nothing if URI does not have a base component (because it
        // identifies the cached archive).
        let base = match uri.base() {
            Some(value) => value,
            None => return,
        };
        // Do nothing if URI is not backed by any archive.
        if !self.file_archives.contains_key(base) {
            return;
        }
        // Do nothing if URI does not point a valid file.
        let extraction_directory: PathBuf = match base.try_into() {
            Ok(path) => path,
            Err(_) => return,
        };
        // Do nothing if any of the files still points to this same archive.
        for resource in &self.file_resources {
            let file = resource.upgrade().unwrap();
            if file.path().starts_with(&extraction_directory) {
                return;
            }
        }
        // Remove both the temporary directory and the cached archive.
        self.file_archives.remove(base).unwrap();
        remove_directory_gently(&extraction_directory);
    }

    /// Creates a Strided Array File view. If the file has been cached
    /// previously, reuses the existing resource; otherwise, creates a new one.
    /// Returns `None` if `uri` does not point to an array from a Strided Array
    /// File.
    fn create_file_view(
        &mut self,
        uri: &ArrayUriReference,
    ) -> Option<Arc<dyn resource::View<'static>>> {
        // Do nothing for URIs which do not form a valid path to a Strided
        // Array File.
        let path = extract_strided_array_file_path(uri)?;
        let array_index = match extract_array_index(uri) {
            Ok(value) => value.unwrap_or(0),
            Err(_) => return None,
        };

        // See if the file is already cached.
        for weak_ptr in &self.file_resources {
            let file = weak_ptr.upgrade().unwrap();
            if file.path() == path {
                return Some(file.create_view(array_index));
            }
        }

        // Since the file hasn't been cached before, create new resource for it.
        let file = Arc::new(StridedArrayFile::new(path));
        let weak_ptr = Arc::downgrade(&file);
        self.file_resources.push(weak_ptr.clone());
        self.resources.push(weak_ptr);
        Some(file.create_view(array_index))
    }

    /// Returns reference count for the given URI. If the URI is not cached,
    /// returns `0`.
    #[cfg(test)]
    fn reference_count(uri: &ArrayUriReference) -> Result<usize, Error> {
        let instance = Self::instance();
        let cache = instance.lock().unwrap();

        Ok(cache.resource_views.reference_count(&uri.resolved()?))
    }
}

/// Tries to extract an absolute file path from the input URI, ignoring its
/// fragment if present. Returns `None` if `uri` does not contain an absolute
/// file path.
///
/// NOTE This method is called while the `ArrayCache`'s mutex is locked, so it
/// cannot create any temporary `ArrayUriReference` objects, because it would
/// cause deadlocks.
fn extract_strided_array_file_path(uri: &ArrayUriReference) -> Option<PathBuf> {
    let resolved = match uri.resolved() {
        Ok(value) => value,
        Err(_) => return None,
    };
    let (absolute, _fragment) = resolved.to_absolute_and_fragment();
    (&absolute).try_into().ok()
}

/// Tries to delete the specified directory with all its contents. This should
/// normally succeed, but if, for example, another process is using some of the
/// files, deletion may fail. We don't want such a failure to stop the whole
/// program, so we silently ignore the error (if any), hoping that some other
/// process will clean it up later.
fn remove_directory_gently(path: &Path) {
    let is_directory = std::fs::metadata(path).is_ok_and(|m| m.is_dir());
    if is_directory {
        let _result = std::fs::remove_dir_all(path);
    }
}