clawhdf5/crates/clawhdf5-remote/tests/object_store.rs

//! `ObjectStoreStorage` against object_store's in-memory and local-file
//! backends (no cloud needed): values equal `File::open`'s, the block cache
//! coalesces, and an object replaced while open is an error.

#![cfg(feature = "object-store")]

mod common;

use std::sync::Arc;

use clawhdf5::File;
use clawhdf5_format::storage::Storage;
use clawhdf5_remote::object_store::memory::InMemory;
use clawhdf5_remote::object_store::path::Path as ObjectPath;
use clawhdf5_remote::object_store::{ObjectStore, ObjectStoreExt, PutPayload};
use clawhdf5_remote::{ObjectStoreStorage, Options, open_object};
use common::{multi_block_file, transcript};

fn put(store: &dyn ObjectStore, path: &str, bytes: Vec<u8>) {
    let rt = tokio_rt();
    rt.block_on(store.put(&ObjectPath::from(path), PutPayload::from(bytes)))
        .unwrap();
}

fn tokio_rt() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
}

#[test]
fn in_memory_store_reads_like_file_open() {
    let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
    let mut compared = 0;
    for (i, p) in common::fixtures().iter().enumerate() {
        let bytes = std::fs::read(p).unwrap();
        let key = format!("f{i}.h5");
        put(store.as_ref(), &key, bytes);
        let (Ok(local), Ok((mut remote, cache))) = (
            File::open(p),
            open_object(store.clone(), &key, &Options::default()),
        ) else {
            continue;
        };
        remote.set_vds_resolver(common::sibling_resolver(p.parent().unwrap().into()));
        assert!(remote.contiguous_bytes().is_none());
        assert_eq!(transcript(&remote), transcript(&local), "{}", p.display());
        assert!(cache.stats().reads > 0);
        compared += 1;
    }
    assert!(compared >= 40, "{compared}");
}

#[test]
fn multi_block_object_is_fetched_in_coalesced_blocks() {
    let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
    let bytes = multi_block_file();
    put(store.as_ref(), "m.h5", bytes.clone());
    let (remote, cache) = open_object(store, "m.h5", &Options::default()).unwrap();
    let local = File::from_bytes(bytes.clone()).unwrap();
    assert_eq!(
        remote.dataset("big").unwrap().read_f64().unwrap(),
        local.dataset("big").unwrap().read_f64().unwrap()
    );
    let s = cache.stats();
    let blocks = (bytes.len() as u64).div_ceil(1 << 20);
    assert!(s.bytes_fetched <= bytes.len() as u64);
    assert!(s.requests <= blocks, "{s:?}");
    assert_eq!(cache.inner().len(), bytes.len() as u64);
}

#[test]
fn local_file_store_reads_like_file_open() {
    let dir = tempfile::tempdir().unwrap();
    let bytes = multi_block_file();
    std::fs::write(dir.path().join("m.h5"), &bytes).unwrap();
    let store: Arc<dyn ObjectStore> = Arc::new(
        clawhdf5_remote::object_store::local::LocalFileSystem::new_with_prefix(dir.path()).unwrap(),
    );
    let (remote, _) = open_object(store, "m.h5", &Options::default()).unwrap();
    let local = File::from_bytes(bytes).unwrap();
    assert_eq!(transcript(&remote), transcript(&local));
}

#[test]
fn an_object_replaced_while_open_is_an_error() {
    let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
    let bytes = multi_block_file();
    put(store.as_ref(), "m.h5", bytes.clone());
    let (remote, _) = open_object(store.clone(), "m.h5", &Options::default()).unwrap();
    assert_eq!(remote.root().groups().unwrap(), ["grp"]);
    let mut other = bytes;
    let n = other.len();
    other[n / 2] ^= 0xff;
    put(store.as_ref(), "m.h5", other);
    let err = remote
        .dataset("big")
        .unwrap()
        .read_f64()
        .unwrap_err()
        .to_string();
    assert!(err.contains("changed while open"), "{err}");
}

#[test]
fn missing_objects_and_async_callers_are_clean_errors() {
    let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
    assert!(ObjectStoreStorage::new(store.clone(), ObjectPath::from("nope.h5")).is_err());
    put(store.as_ref(), "m.h5", multi_block_file());
    let storage = ObjectStoreStorage::new(store, ObjectPath::from("m.h5")).unwrap();
    // From inside a current-thread runtime: works (the read runs on the
    // storage's own runtime), no panic, no deadlock.
    let rt = tokio_rt();
    let r = rt.block_on(async { storage.read_at(0, 10).map(|b| b.len()) });
    assert_eq!(r.unwrap(), 10);
    // Dropping the storage inside a runtime does not panic.
    rt.block_on(async move { drop(storage) });
}

/// The advice for async callers works: a read in `spawn_blocking` of a
/// multi-threaded runtime (where `Handle::try_current` is Ok), and a read
/// straight inside a current-thread runtime's task.
#[test]
fn object_stores_read_from_spawn_blocking_and_inside_runtimes() {
    let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
    let bytes = multi_block_file();
    put(store.as_ref(), "m.h5", bytes.clone());
    let want = File::from_bytes(bytes)
        .unwrap()
        .dataset("big")
        .unwrap()
        .read_f64()
        .unwrap();
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(2)
        .build()
        .unwrap();
    let (s, w) = (store.clone(), want.clone());
    let got = rt
        .block_on(async move {
            tokio::task::spawn_blocking(move || {
                assert!(tokio::runtime::Handle::try_current().is_ok());
                let (f, _) =
                    open_object(s, "m.h5", &Options::default()).map_err(|e| e.to_string())?;
                let v = f.dataset("big").unwrap().read_f64().unwrap();
                Ok::<_, String>(v == w)
            })
            .await
        })
        .unwrap();
    assert_eq!(got, Ok(true));
    let rt = tokio_rt();
    let got = rt.block_on(async {
        let (f, _) = open_object(store, "m.h5", &Options::default()).unwrap();
        f.dataset("big").unwrap().read_f64().unwrap()
    });
    assert_eq!(got, want);
}