clawhdf5/crates/clawhdf5-remote/src/object.rs

//! Object stores (S3, GCS, Azure, local files, memory) through the
//! [`object_store`] crate: [`ObjectStoreStorage`].
//!
//! `object_store` is async and [`Storage`] is synchronous (parsing is CPU
//! work; `docs/design/range-reads.md` §3 (a)). The storage owns a small
//! multi-threaded tokio runtime (two worker threads): each read is spawned
//! on it and the calling thread waits for the result, so it can be used
//! from any thread — several at once, and from async code too. From async
//! code prefer `tokio::task::spawn_blocking` (a read blocks the thread it
//! is called on, which inside a runtime is one of its workers).
//!
//! The object is pinned when the storage is made: its size, and its ETag
//! (sent as `If-Match` with every read, and compared with every response)
//! or, without one, its version or modification time. A change while it is
//! open is [`RemoteError::FileChanged`]. The ranges of one `read_ranges`
//! call are fetched concurrently (at most 8 at a time).

use std::borrow::Cow;
use std::ops::Range;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use clawhdf5_format::error::FormatError;
use clawhdf5_format::storage::Storage;
use futures_util::{StreamExt, TryStreamExt};
use object_store::path::Path;
use object_store::{GetOptions, GetRange, ObjectMeta, ObjectStore, ObjectStoreExt};

use crate::error::RemoteError;

/// Concurrent requests of one `read_ranges` call.
const MAX_CONCURRENT: usize = 8;

/// One object of an [`ObjectStore`], read by ranged `get`s. See the
/// [module documentation](self).
pub struct ObjectStoreStorage {
    store: Arc<dyn ObjectStore>,
    path: Path,
    meta: ObjectMeta,
    runtime: Option<tokio::runtime::Runtime>,
    requests: AtomicU64,
    bytes: AtomicU64,
}

impl std::fmt::Debug for ObjectStoreStorage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ObjectStoreStorage")
            .field("store", &self.store.to_string())
            .field("path", &self.path)
            .field("size", &self.meta.size)
            .field("e_tag", &self.meta.e_tag)
            .finish()
    }
}

fn os_error(e: object_store::Error) -> RemoteError {
    match e {
        object_store::Error::Precondition { .. } | object_store::Error::NotModified { .. } => {
            RemoteError::FileChanged(e.to_string())
        }
        other => RemoteError::ObjectStore(other.to_string()),
    }
}

impl ObjectStoreStorage {
    /// Open the object at `path` of `store`: one `head` request for its
    /// size and validators.
    pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Result<Self, RemoteError> {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .thread_name("clawhdf5-remote")
            .enable_all()
            .build()
            .map_err(|e| RemoteError::Usage(format!("cannot start a tokio runtime: {e}")))?;
        let mut s = ObjectStoreStorage {
            store,
            path,
            meta: ObjectMeta {
                location: Path::default(),
                last_modified: Default::default(),
                size: 0,
                e_tag: None,
                version: None,
            },
            runtime: Some(runtime),
            requests: AtomicU64::new(1),
            bytes: AtomicU64::new(0),
        };
        let (store, path) = (s.store.clone(), s.path.clone());
        s.meta = s.block_on(async move { store.head(&path).await.map_err(os_error) })?;
        Ok(s)
    }

    /// The object's metadata as pinned at open.
    pub fn meta(&self) -> &ObjectMeta {
        &self.meta
    }

    /// Requests made (the `head` included) and bytes received.
    pub fn stats(&self) -> (u64, u64) {
        (
            self.requests.load(Ordering::Relaxed),
            self.bytes.load(Ordering::Relaxed),
        )
    }

    /// Run `fut` on the storage's own runtime and wait for it. The future
    /// never runs on the caller's thread, so the caller's context does not
    /// matter: a plain thread, `spawn_blocking`, or even inside another
    /// runtime (whose thread is then blocked for the duration of the read,
    /// as by any blocking call, but nothing deadlocks or panics).
    fn block_on<T: Send + 'static>(
        &self,
        fut: impl std::future::Future<Output = Result<T, RemoteError>> + Send + 'static,
    ) -> Result<T, RemoteError> {
        let rt = self.runtime.as_ref().expect("runtime lives until drop");
        let (tx, rx) = std::sync::mpsc::sync_channel(1);
        rt.spawn(async move {
            let _ = tx.send(fut.await);
        });
        rx.recv().map_err(|_| {
            RemoteError::ObjectStore("the object store task ended without a result".into())
        })?
    }

    fn options(&self, range: Range<u64>) -> GetOptions {
        let mut o = GetOptions {
            range: Some(GetRange::Bounded(range)),
            ..GetOptions::default()
        };
        if let Some(e) = &self.meta.e_tag {
            o.if_match = Some(e.clone());
        } else if let Some(v) = &self.meta.version {
            o.version = Some(v.clone());
        } else {
            o.if_unmodified_since = Some(self.meta.last_modified);
        }
        o
    }

    /// The object's first `n` bytes (fewer if it is shorter).
    pub(crate) fn fetch_first(&self, n: u64) -> Result<Vec<u8>, RemoteError> {
        Ok(self.fetch_all(&[0..n])?.pop().unwrap_or_default())
    }

    fn fetch_all(&self, ranges: &[Range<u64>]) -> Result<Vec<Vec<u8>>, RemoteError> {
        let len = self.meta.size;
        let jobs: Vec<(usize, Range<u64>)> = ranges
            .iter()
            .enumerate()
            .filter_map(|(i, r)| {
                let end = r.end.min(len);
                (r.start < end).then_some((i, r.start..end))
            })
            .collect();
        let store = self.store.clone();
        let path = self.path.clone();
        let pinned = self.meta.e_tag.clone();
        let reqs: Vec<(usize, Range<u64>, GetOptions)> = jobs
            .into_iter()
            .map(|(i, r)| {
                let o = self.options(r.clone());
                (i, r, o)
            })
            .collect();
        self.requests
            .fetch_add(reqs.len() as u64, Ordering::Relaxed);
        let fetched: Vec<(usize, Vec<u8>)> = self.block_on(async move {
            futures_util::stream::iter(reqs)
                .map(|(i, r, o)| {
                    let (store, path, pinned) = (store.clone(), path.clone(), pinned.clone());
                    async move {
                        let got = store.get_opts(&path, o).await.map_err(os_error)?;
                        if let (Some(want), Some(have)) = (&pinned, &got.meta.e_tag)
                            && want != have
                        {
                            return Err(RemoteError::FileChanged(format!(
                                "{path}: ETag {have} instead of {want}"
                            )));
                        }
                        if got.meta.size != len {
                            return Err(RemoteError::FileChanged(format!(
                                "{path}: size {} instead of {len}",
                                got.meta.size
                            )));
                        }
                        let bytes = got.bytes().await.map_err(os_error)?;
                        if bytes.len() as u64 != r.end - r.start {
                            return Err(RemoteError::BadResponse(format!(
                                "{path}: {} bytes for range {r:?}",
                                bytes.len()
                            )));
                        }
                        Ok::<_, RemoteError>((i, bytes.to_vec()))
                    }
                })
                .buffer_unordered(MAX_CONCURRENT)
                .try_collect()
                .await
        })?;
        let mut out = vec![Vec::new(); ranges.len()];
        for (i, b) in fetched {
            self.bytes.fetch_add(b.len() as u64, Ordering::Relaxed);
            out[i] = b;
        }
        Ok(out)
    }
}

impl Drop for ObjectStoreStorage {
    fn drop(&mut self) {
        // Dropping a runtime blocks, which panics inside an async context.
        if let Some(rt) = self.runtime.take() {
            rt.shutdown_background();
        }
    }
}

impl Storage for ObjectStoreStorage {
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        let range = offset..offset.saturating_add(len as u64);
        let mut v = self.fetch_all(std::slice::from_ref(&range))?;
        Ok(Cow::Owned(v.pop().unwrap_or_default()))
    }

    fn len(&self) -> u64 {
        self.meta.size
    }

    fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
        if ranges.iter().any(|r| r.end < r.start) {
            return Err(FormatError::Storage(
                "read range ends before it starts".into(),
            ));
        }
        Ok(self
            .fetch_all(ranges)?
            .into_iter()
            .map(Cow::Owned)
            .collect())
    }
}

/// The store and object path for a cloud URL (`s3://bucket/key`,
/// `gs://bucket/key`, `az://container/key`, ...), configured from the
/// environment as `object_store`'s `from_env` builders do (`AWS_*`,
/// `GOOGLE_*`, `AZURE_*`).
#[cfg(any(feature = "s3", feature = "gcs", feature = "azure"))]
pub(crate) fn store_for_url(url: &str) -> Result<(Arc<dyn ObjectStore>, Path), RemoteError> {
    let redactor = crate::error::Redactor::new(url);
    let parsed = object_store::path::Path::parse(
        url.split_once("://")
            .and_then(|(_, rest)| rest.split_once('/'))
            .map(|(_, key)| key)
            .unwrap_or(""),
    )
    .map_err(|e| RemoteError::InvalidUrl(format!("{}: {e}", crate::redact_url(url))))?;
    let scheme = url.split_once("://").map(|(s, _)| s.to_ascii_lowercase());
    let store: Arc<dyn ObjectStore> = match scheme.as_deref() {
        #[cfg(feature = "s3")]
        Some("s3" | "s3a") => Arc::new(
            object_store::aws::AmazonS3Builder::from_env()
                .with_url(url)
                .build()
                .map_err(|e| os_error(e).scrubbed(&redactor))?,
        ),
        #[cfg(feature = "gcs")]
        Some("gs") => Arc::new(
            object_store::gcp::GoogleCloudStorageBuilder::from_env()
                .with_url(url)
                .build()
                .map_err(|e| os_error(e).scrubbed(&redactor))?,
        ),
        #[cfg(feature = "azure")]
        Some("az" | "azure" | "abfs" | "abfss" | "adl") => Arc::new(
            object_store::azure::MicrosoftAzureBuilder::from_env()
                .with_url(url)
                .build()
                .map_err(|e| os_error(e).scrubbed(&redactor))?,
        ),
        _ => return Err(RemoteError::UnsupportedScheme(crate::redact_url(url))),
    };
    Ok((store, parsed))
}

#[cfg(all(test, feature = "s3"))]
mod tests {
    #[test]
    fn s3_urls_name_the_bucket_and_key() {
        let (store, path) = super::store_for_url("s3://my-bucket/dir/file.h5").unwrap();
        assert_eq!(path.as_ref(), "dir/file.h5");
        assert!(store.to_string().contains("my-bucket"), "{store}");
    }
}