clawhdf5-remote: object stores through object_store (S3, GCS, Azure)
ObjectStoreStorage (feature `object-store`, pure Rust) reads one object of any object_store store by ranged get_opts, pinned at open by a head request: If-Match with its ETag (and the ETag and size of every response compared), else its version or modification time. A change is RemoteError::FileChanged. object_store is async and Storage is not, so the storage owns a small multi-threaded tokio runtime (two workers) and blocks the calling thread on it; the ranges of one read_ranges call are fetched concurrently (up to 8). From inside another tokio runtime it refuses with RemoteError::Usage instead of blocking a worker, and it shuts its runtime down in the background on drop so dropping it in async code does not panic. open_object(store, path, options) opens a file through a block cache (first block prefetched); open_url accepts s3://, gs:// and az:// with the `s3`, `gcs` and `azure` features, configured from the environment by object_store's from_env builders. Those pull object_store's cloud clients and aws-lc-rs (C), so they are opt-in; without them the URL is a clean UnsupportedScheme error naming the feature. Tests against object_store's in-memory and local-file stores (no cloud): every fixture's transcript equals File::open's, a multi-block object is fetched in coalesced block runs, an object replaced while open is an error, and a missing object or a read from inside a runtime is a clean error. ci-test.sh lints all backends, runs these tests (with s3 for its URL parsing test) and checks object-store for C in the no-C step. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,285 @@
|
||||
//! 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) and blocks the calling
|
||||
//! thread on it for each read, so it can be used from ordinary threads —
|
||||
//! several at once. Calling it from inside another tokio runtime would
|
||||
//! block that runtime's worker, so it refuses with
|
||||
//! [`RemoteError::Usage`]: from async code, read in
|
||||
//! `tokio::task::spawn_blocking`.
|
||||
//!
|
||||
//! 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),
|
||||
)
|
||||
}
|
||||
|
||||
fn block_on<T>(
|
||||
&self,
|
||||
fut: impl std::future::Future<Output = Result<T, RemoteError>>,
|
||||
) -> Result<T, RemoteError> {
|
||||
if tokio::runtime::Handle::try_current().is_ok() {
|
||||
return Err(RemoteError::Usage(
|
||||
"ObjectStoreStorage blocks on its own runtime and cannot be read from inside \
|
||||
an async runtime; read in tokio::task::spawn_blocking"
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
let rt = self.runtime.as_ref().expect("runtime lives until drop");
|
||||
rt.block_on(fut)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
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 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!("{url}: {e}")))?;
|
||||
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(os_error)?,
|
||||
),
|
||||
#[cfg(feature = "gcs")]
|
||||
Some("gs") => Arc::new(
|
||||
object_store::gcp::GoogleCloudStorageBuilder::from_env()
|
||||
.with_url(url)
|
||||
.build()
|
||||
.map_err(os_error)?,
|
||||
),
|
||||
#[cfg(feature = "azure")]
|
||||
Some("az" | "azure" | "abfs" | "abfss" | "adl") => Arc::new(
|
||||
object_store::azure::MicrosoftAzureBuilder::from_env()
|
||||
.with_url(url)
|
||||
.build()
|
||||
.map_err(os_error)?,
|
||||
),
|
||||
_ => return Err(RemoteError::UnsupportedScheme(url.to_string())),
|
||||
};
|
||||
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}");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user