use std::{ops::Range, str::FromStr, sync::Arc, time::Duration};
use async_trait::async_trait;
use bytes::Bytes;
use futures::TryStreamExt;
use object_store::{
Certificate, ClientOptions, Error as ObjError, GetOptions, GetRange, MultipartUpload,
ObjectStore, ObjectStoreExt, PutMode, PutOptions, PutPayload, UpdateVersion,
aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, S3ConditionalPut},
path::Path as ObjPath,
};
use super::{
ObjectMeta, StorageError, StorageOptions, StorageProvider, counting, io_counters,
logical_list_key, options::apply, retry,
};
use crate::runtime_metrics::io::UsageMeter;
fn has_custom_endpoint(opts: &StorageOptions) -> bool {
opts.keys().any(|k| {
matches!(
AmazonS3ConfigKey::from_str(k),
Ok(AmazonS3ConfigKey::Endpoint | AmazonS3ConfigKey::S3Endpoint)
)
})
}
#[derive(Debug)]
pub struct S3StorageProvider {
bucket: String,
prefix: String,
store: Arc<AmazonS3>,
meter: Arc<UsageMeter>,
}
impl S3StorageProvider {
pub fn new(bucket: impl Into<String>) -> Result<Self, StorageError> {
Self::new_with_prefix(bucket, "", &StorageOptions::new())
}
pub fn new_with_prefix(
bucket: impl Into<String>,
prefix: impl Into<String>,
opts: &StorageOptions,
) -> Result<Self, StorageError> {
let bucket = bucket.into();
let uri = format!("s3://{bucket}");
let mut builder = AmazonS3Builder::new()
.with_bucket_name(&bucket)
.with_conditional_put(S3ConditionalPut::ETagMatch)
.with_retry(retry::config());
builder = if has_custom_endpoint(opts) {
builder.with_virtual_hosted_style_request(false)
} else {
builder.with_client_options(tuned_client_options())
};
let builder = apply::<AmazonS3ConfigKey, _>(builder, opts, &uri, |b, key, value| {
b.with_config(key, value)
})?;
let store = builder.build().map_err(|e| StorageError::Permanent {
uri,
source: Box::new(e),
})?;
Ok(Self {
bucket,
prefix: normalize_prefix(prefix),
store: Arc::new(store),
meter: UsageMeter::process_default(),
})
}
pub fn new_with_endpoint(
endpoint: impl Into<String>,
bucket: impl Into<String>,
access_key: impl Into<String>,
secret_key: impl Into<String>,
region: impl Into<String>,
trusted_ca_pem: Option<&[u8]>,
) -> Result<Self, StorageError> {
let bucket = bucket.into();
let endpoint = endpoint.into();
let store = build_custom_endpoint_store(
&endpoint,
&bucket,
access_key,
secret_key,
region,
trusted_ca_pem,
)?;
Ok(Self {
bucket,
prefix: String::new(),
store: Arc::new(store),
meter: UsageMeter::process_default(),
})
}
pub fn new_with_endpoint_and_prefix(
endpoint: impl Into<String>,
bucket: impl Into<String>,
access_key: impl Into<String>,
secret_key: impl Into<String>,
region: impl Into<String>,
prefix: impl Into<String>,
trusted_ca_pem: Option<&[u8]>,
) -> Result<Self, StorageError> {
let mut provider = Self::new_with_endpoint(
endpoint,
bucket,
access_key,
secret_key,
region,
trusted_ca_pem,
)?;
provider.prefix = normalize_prefix(prefix);
Ok(provider)
}
pub fn from_object_store(bucket: impl Into<String>, store: AmazonS3) -> Self {
Self {
bucket: bucket.into(),
prefix: String::new(),
store: Arc::new(store),
meter: UsageMeter::process_default(),
}
}
pub fn with_usage_meter(mut self, meter: Arc<UsageMeter>) -> Self {
self.meter = meter;
self
}
pub fn from_object_store_with_prefix(
bucket: impl Into<String>,
store: AmazonS3,
prefix: impl Into<String>,
) -> Self {
let mut provider = Self::from_object_store(bucket, store);
provider.prefix = normalize_prefix(prefix);
provider
}
pub fn bucket(&self) -> &str {
&self.bucket
}
pub fn prefix(&self) -> &str {
&self.prefix
}
fn key(&self, uri: &str) -> String {
let uri = uri.trim_start_matches('/');
if self.prefix.is_empty() {
uri.to_string()
} else {
format!("{}/{uri}", self.prefix)
}
}
fn path(&self, uri: &str) -> Result<ObjPath, StorageError> {
let key = self.key(uri);
ObjPath::parse(&key).map_err(|e| StorageError::Permanent {
uri: uri.into(),
source: Box::new(e),
})
}
}
fn normalize_prefix(prefix: impl Into<String>) -> String {
prefix.into().trim_matches('/').to_string()
}
fn build_custom_endpoint_store(
endpoint: &str,
bucket: &str,
access_key: impl Into<String>,
secret_key: impl Into<String>,
region: impl Into<String>,
trusted_ca_pem: Option<&[u8]>,
) -> Result<AmazonS3, StorageError> {
let mut builder = AmazonS3Builder::new()
.with_endpoint(endpoint)
.with_bucket_name(bucket)
.with_access_key_id(access_key.into())
.with_secret_access_key(secret_key.into())
.with_region(region.into())
.with_virtual_hosted_style_request(false)
.with_conditional_put(S3ConditionalPut::ETagMatch);
if let Some(ca_pem) = trusted_ca_pem {
let cert = Certificate::from_pem(ca_pem).map_err(|e| StorageError::Permanent {
uri: format!("s3://{bucket} @ {endpoint}"),
source: Box::new(e),
})?;
let client_options = ClientOptions::new().with_root_certificate(cert);
builder = builder.with_client_options(client_options);
} else {
builder = builder.with_allow_http(true);
}
builder.build().map_err(|e| StorageError::Permanent {
uri: format!("s3://{bucket} @ {endpoint}"),
source: Box::new(e),
})
}
const S3_POOL_MAX_IDLE_PER_HOST: usize = 1024;
const S3_POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(10);
const S3_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
fn tuned_client_options() -> ClientOptions {
ClientOptions::new()
.with_pool_max_idle_per_host(S3_POOL_MAX_IDLE_PER_HOST)
.with_pool_idle_timeout(S3_POOL_IDLE_TIMEOUT)
.with_connect_timeout(S3_CONNECT_TIMEOUT)
}
fn translate(uri: &str, e: ObjError) -> StorageError {
match e {
ObjError::NotFound { .. } => StorageError::NotFound { uri: uri.into() },
ObjError::AlreadyExists { .. } | ObjError::Precondition { .. } => {
StorageError::PreconditionFailed { uri: uri.into() }
}
ObjError::Generic { source, .. } => StorageError::TransientExhausted {
uri: uri.into(),
source,
},
other => StorageError::Permanent {
uri: uri.into(),
source: Box::new(other),
},
}
}
#[async_trait]
impl StorageProvider for S3StorageProvider {
async fn head(&self, uri: &str) -> Result<ObjectMeta, StorageError> {
let path = self.path(uri)?;
let meta = self
.store
.head(&path)
.await
.map_err(|e| translate(uri, e))?;
self.meter.record_head();
Ok(ObjectMeta {
size: meta.size as u64,
etag: meta.e_tag,
last_modified: meta.last_modified.into(),
})
}
async fn get(&self, uri: &str) -> Result<(Bytes, ObjectMeta), StorageError> {
let path = self.path(uri)?;
let tl = io_counters::timeline_start();
let out = retry::complete_get(uri, || async {
let result = self.store.get(&path).await.map_err(|e| translate(uri, e))?;
let meta = ObjectMeta {
size: result.meta.size as u64,
etag: result.meta.e_tag.clone(),
last_modified: result.meta.last_modified.into(),
};
let bytes = result.bytes().await.map_err(|e| translate(uri, e))?;
Ok((bytes, meta))
})
.await;
if let Ok((b, _)) = &out {
self.meter.record_get(uri, None, b.len() as u64);
io_counters::timeline_record("get", uri, 0, b.len() as u64, tl);
}
out
}
async fn get_if_none_match(
&self,
uri: &str,
etag: &str,
) -> Result<Option<(Bytes, ObjectMeta)>, StorageError> {
let path = self.path(uri)?;
let out = retry::with_reissue(|| async {
let options = GetOptions {
if_none_match: Some(etag.to_string()),
..GetOptions::default()
};
let result = match self.store.get_opts(&path, options).await {
Ok(result) => result,
Err(ObjError::NotModified { .. }) => return Ok(None),
Err(e) => return Err(translate(uri, e)),
};
let meta = ObjectMeta {
size: result.meta.size as u64,
etag: result.meta.e_tag.clone(),
last_modified: result.meta.last_modified.into(),
};
let bytes = result.bytes().await.map_err(|e| translate(uri, e))?;
Ok(Some((bytes, meta)))
})
.await;
match &out {
Ok(Some((b, _))) => self.meter.record_get(uri, None, b.len() as u64),
Ok(None) => self.meter.record_get(uri, None, 0),
Err(_) => {}
}
out
}
#[cfg_attr(
feature = "detailed-tracing",
tracing::instrument(skip_all, fields(uri = uri, len = range.end - range.start))
)]
async fn get_range(&self, uri: &str, range: Range<u64>) -> Result<Bytes, StorageError> {
let path = self.path(uri)?;
let requested = (range.start, range.end);
let off = range.start;
let tl = io_counters::timeline_start();
let out = retry::complete_range(uri, range, |r| async {
self.store
.get_range(&path, r)
.await
.map_err(|e| translate(uri, e))
})
.await;
if let Ok(b) = &out {
self.meter.record_get(uri, Some(requested), b.len() as u64);
io_counters::timeline_record("get_range", uri, off, b.len() as u64, tl);
}
out
}
async fn tail(&self, uri: &str, len: u64) -> Result<(Bytes, u64), StorageError> {
if len == 0 {
let meta = self.head(uri).await?;
return Ok((Bytes::new(), meta.size));
}
let path = self.path(uri)?;
let tl = io_counters::timeline_start();
let out = retry::with_reissue(|| async {
let opts = GetOptions {
range: Some(GetRange::Suffix(len)),
..Default::default()
};
let result = self
.store
.get_opts(&path, opts)
.await
.map_err(|e| translate(uri, e))?;
let size = result.meta.size as u64;
let bytes = result.bytes().await.map_err(|e| translate(uri, e))?;
Ok((bytes, size))
})
.await;
if let Ok((b, size)) = &out {
let start = size.saturating_sub(b.len() as u64);
self.meter
.record_get(uri, Some((start, *size)), b.len() as u64);
io_counters::timeline_record("tail", uri, start, b.len() as u64, tl);
}
out
}
async fn put_atomic(&self, uri: &str, bytes: Bytes) -> Result<Option<String>, StorageError> {
let path = self.path(uri)?;
let n = bytes.len() as u64;
let out = retry::with_reissue(|| {
let bytes = bytes.clone();
async {
let opts = PutOptions {
mode: PutMode::Create,
..Default::default()
};
self.store
.put_opts(&path, PutPayload::from_bytes(bytes), opts)
.await
.map(|r| r.e_tag)
.map_err(|e| translate(uri, e))
}
})
.await;
if out.is_ok() {
self.meter.record_put(n);
}
out
}
async fn put_if_match(
&self,
uri: &str,
bytes: Bytes,
expected_etag: Option<&str>,
) -> Result<Option<String>, StorageError> {
let path = self.path(uri)?;
let opts = match expected_etag {
None => PutOptions {
mode: PutMode::Create,
..Default::default()
},
Some(expected) => PutOptions {
mode: PutMode::Update(UpdateVersion {
e_tag: Some(expected.to_string()),
version: None,
}),
..Default::default()
},
};
let n = bytes.len() as u64;
let out = self
.store
.put_opts(&path, PutPayload::from_bytes(bytes), opts)
.await
.map(|r| r.e_tag)
.map_err(|e| translate(uri, e));
if out.is_ok() {
self.meter.record_put(n);
}
out
}
async fn put_multipart(&self, uri: &str) -> Result<Box<dyn MultipartUpload>, StorageError> {
let path = self.path(uri)?;
let upload = self
.store
.put_multipart(&path)
.await
.map_err(|e| translate(uri, e))?;
self.meter.record_put(0);
Ok(counting::wrap_multipart(upload, Arc::clone(&self.meter)))
}
async fn delete(&self, uri: &str) -> Result<(), StorageError> {
let path = self.path(uri)?;
match self.store.delete(&path).await {
Ok(()) => {
self.meter.record_delete();
Ok(())
}
Err(ObjError::NotFound { .. }) => {
self.meter.record_delete();
Ok(())
}
Err(e) => Err(translate(uri, e)),
}
}
async fn list_with_prefix_metadata(
&self,
prefix: &str,
) -> Result<Vec<(String, ObjectMeta)>, StorageError> {
let path = self.path(prefix)?;
let mut stream = self.store.list(Some(&path));
self.meter.record_list();
let mut out = Vec::new();
while let Some(meta) = stream.try_next().await.map_err(|e| translate(prefix, e))? {
let location = meta.location.to_string();
out.push((
logical_list_key(&self.prefix, &location),
ObjectMeta {
size: meta.size,
etag: meta.e_tag,
last_modified: meta.last_modified.into(),
},
));
}
Ok(out)
}
fn object_store_handle(&self, uri: &str) -> Option<(Arc<dyn ObjectStore>, ObjPath)> {
let path = self.path(uri).ok()?;
Some((
counting::wrap_object_store(
Arc::clone(&self.store) as Arc<dyn ObjectStore>,
Arc::clone(&self.meter),
),
path,
))
}
fn usage_meter(&self) -> Arc<UsageMeter> {
Arc::clone(&self.meter)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn translate_not_found_to_typed_variant() {
let err = translate(
"some/key",
ObjError::NotFound {
path: "some/key".into(),
source: "raw".into(),
},
);
match err {
StorageError::NotFound { uri } => assert_eq!(uri, "some/key"),
other => panic!("expected NotFound; got {other:?}"),
}
}
#[test]
fn translate_already_exists_to_precondition_failed() {
let err = translate(
"k",
ObjError::AlreadyExists {
path: "k".into(),
source: "raw".into(),
},
);
assert!(matches!(err, StorageError::PreconditionFailed { uri } if uri == "k"));
}
#[test]
fn translate_precondition_to_precondition_failed() {
let err = translate(
"k",
ObjError::Precondition {
path: "k".into(),
source: "raw".into(),
},
);
assert!(matches!(err, StorageError::PreconditionFailed { uri } if uri == "k"));
}
#[test]
fn translate_generic_to_transient_exhausted() {
let err = translate(
"k",
ObjError::Generic {
store: "S3",
source: "boom".into(),
},
);
match err {
StorageError::TransientExhausted { uri, .. } => assert_eq!(uri, "k"),
other => panic!("expected TransientExhausted; got {other:?}"),
}
}
#[test]
fn translate_other_variant_to_permanent() {
let err = translate(
"k",
ObjError::UnknownConfigurationKey {
store: "S3",
key: "foo".into(),
},
);
match err {
StorageError::Permanent { uri, .. } => assert_eq!(uri, "k"),
other => panic!("expected Permanent; got {other:?}"),
}
}
#[test]
fn path_parses_simple_uri() {
let p = endpoint_provider().path("foo/bar.txt").expect("parse");
assert_eq!(p.to_string(), "foo/bar.txt");
}
#[test]
fn path_parses_nested_uri() {
let p = endpoint_provider()
.path("manifest/manifest-000042.json")
.expect("parse");
assert_eq!(p.to_string(), "manifest/manifest-000042.json");
}
fn endpoint_provider() -> S3StorageProvider {
S3StorageProvider::new_with_endpoint(
"http://127.0.0.1:1",
"test-bucket",
"AKIATESTKEY",
"secret/example",
"us-east-1",
None,
)
.expect("construct with endpoint")
}
#[test]
fn new_with_endpoint_builds_succeeds_and_exposes_bucket() {
let p = endpoint_provider();
assert_eq!(p.bucket(), "test-bucket");
}
#[test]
fn from_object_store_preserves_bucket() {
let store = AmazonS3Builder::new()
.with_endpoint("http://127.0.0.1:1")
.with_bucket_name("hatch-bucket")
.with_access_key_id("AKIATESTKEY")
.with_secret_access_key("secret")
.with_region("us-east-1")
.with_allow_http(true)
.with_virtual_hosted_style_request(false)
.build()
.expect("build AmazonS3");
let p = S3StorageProvider::from_object_store("hatch-bucket", store);
assert_eq!(p.bucket(), "hatch-bucket");
}
#[test]
fn debug_impl_does_not_panic() {
let p = endpoint_provider();
let s = format!("{p:?}");
assert!(s.contains("S3StorageProvider"));
}
#[test]
fn normalize_prefix_trims_surrounding_slashes() {
assert_eq!(normalize_prefix("/tbl/"), "tbl");
assert_eq!(normalize_prefix("///a/b///"), "a/b");
assert_eq!(normalize_prefix("plain"), "plain");
assert_eq!(normalize_prefix(""), "");
}
#[test]
fn key_without_prefix_strips_leading_slash() {
let p = endpoint_provider();
assert_eq!(p.prefix(), "");
assert_eq!(p.key("/foo/bar"), "foo/bar");
assert_eq!(p.key("foo/bar"), "foo/bar");
}
#[test]
fn key_with_prefix_prepends_and_strips_leading_slash() {
let mut p = endpoint_provider();
p.prefix = "tbl".into();
assert_eq!(p.prefix(), "tbl");
assert_eq!(p.key("data/seg-1"), "tbl/data/seg-1");
assert_eq!(p.key("/data/seg-1"), "tbl/data/seg-1");
}
#[test]
fn new_with_endpoint_and_prefix_normalizes_and_applies_prefix() {
let p = S3StorageProvider::new_with_endpoint_and_prefix(
"http://127.0.0.1:1",
"b",
"AKIATESTKEY",
"secret",
"us-east-1",
"/scoped/tbl/",
None,
)
.expect("construct with endpoint + prefix");
assert_eq!(p.bucket(), "b");
assert_eq!(p.prefix(), "scoped/tbl");
assert_eq!(p.key("data/seg-1"), "scoped/tbl/data/seg-1");
}
#[test]
fn object_store_handle_returns_path_under_prefix() {
let mut p = endpoint_provider();
p.prefix = "tbl".into();
let (_, path) = p
.object_store_handle("data/seg-1")
.expect("handle for valid uri");
assert_eq!(path.to_string(), "tbl/data/seg-1");
}
}