use crate::config::{Source, Storage};
use anyhow::{Result, bail, ensure};
use futures::{StreamExt, TryStreamExt};
use object_store::{
ClientOptions, GetOptions, ObjectStore, ObjectStoreExt, RetryConfig, aws::AmazonS3Builder,
path::Path,
};
use std::{
collections::HashMap,
io::{Read, Seek, SeekFrom},
sync::{Arc, Mutex},
time::{Duration, Instant},
};
#[derive(Debug)]
struct Incomplete;
impl std::fmt::Display for Incomplete {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("incomplete response")
}
}
impl std::error::Error for Incomplete {}
async fn retry(error: anyhow::Error, failures: u32, progress: Instant) -> Result<()> {
let transient = error.is::<Incomplete>()
|| error
.downcast_ref::<object_store::Error>()
.is_some_and(|e| {
matches!(
e,
object_store::Error::Generic { .. } | object_store::Error::JoinError { .. }
)
});
if !transient {
return Err(error);
}
if failures >= 64 || progress.elapsed().as_secs() >= 900 {
bail!("storage no-progress retry budget exhausted: {error}");
}
let mut random = [0u8; 8];
getrandom::fill(&mut random).map_err(|e| anyhow::anyhow!("retry jitter: {e}"))?;
let fraction = (u64::from_le_bytes(random) >> 11) as f64 / ((1u64 << 53) as f64);
let cap = (2f64.powi((failures - 1).min(5) as i32)).min(30.0);
let delay = (fraction * cap).min((900.0 - progress.elapsed().as_secs_f64()).max(0.0));
tokio::time::sleep(Duration::from_secs_f64(delay)).await;
Ok(())
}
pub struct Client {
config: Storage,
pub runtime: Arc<tokio::runtime::Runtime>,
stores: Mutex<HashMap<String, Arc<dyn ObjectStore>>>,
}
impl Client {
pub fn new(config: Storage, runtime: Arc<tokio::runtime::Runtime>) -> Self {
Self {
config,
runtime,
stores: Mutex::new(HashMap::new()),
}
}
pub fn store(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>> {
let mut stores = self.stores.lock().unwrap();
if let Some(store) = stores.get(bucket) {
return Ok(store.clone());
}
let (Some(base_url), Some(region), Some(key), Some(secret)) = (
self.config.base_url.as_ref(),
self.config.region.as_ref(),
self.config.access_key_id.as_ref(),
self.config.secret_access_key.as_ref(),
) else {
bail!("s3 URIs need storage baseUrl, region, accessKeyId and secretAccessKey");
};
let builder = AmazonS3Builder::new()
.with_bucket_name(bucket)
.with_region(region)
.with_endpoint(base_url)
.with_access_key_id(key)
.with_secret_access_key(secret)
.with_virtual_hosted_style_request(self.config.virtual_hosted_style_request)
.with_client_options(
ClientOptions::new()
.with_allow_http(true)
.with_pool_idle_timeout(Duration::from_secs(90))
.with_pool_max_idle_per_host(16)
.with_connect_timeout(Duration::from_secs(10))
.with_timeout(Duration::from_secs(60)),
)
.with_retry(RetryConfig {
max_retries: 3,
retry_timeout: Duration::from_secs(60),
..Default::default()
});
let store: Arc<dyn ObjectStore> = Arc::new(builder.build()?);
stores.insert(bucket.into(), store.clone());
Ok(store)
}
}
pub struct Object {
local: Option<std::fs::File>,
remote: Option<(Arc<dyn ObjectStore>, Path)>,
runtime: Arc<tokio::runtime::Runtime>,
pub size: u64,
etag: Option<String>,
}
impl Object {
pub fn open(source: &Source, client: &Client) -> Result<Self> {
let mut object = Self {
local: None,
remote: None,
runtime: client.runtime.clone(),
size: 0,
etag: source.etag.clone().filter(|e| !e.is_empty()),
};
if let Some(uri) = source.uri.strip_prefix("s3://") {
let (bucket, key) = uri
.split_once('/')
.ok_or_else(|| anyhow::anyhow!("invalid S3 URI"))?;
let store = client.store(bucket)?;
let path = Path::from(key);
object.size = if let Some(n) = source.size_bytes {
n
} else {
object.runtime.block_on(async {
let progress = Instant::now();
let mut failures = 0;
loop {
match store.head(&path).await {
Ok(meta) => return Ok::<_, anyhow::Error>(meta.size),
Err(error) => {
failures += 1;
retry(error.into(), failures, progress).await?;
}
}
}
})?
};
object.remote = Some((store, path));
} else {
let file =
std::fs::File::open(source.uri.strip_prefix("file://").unwrap_or(&source.uri))?;
object.size = file.metadata()?.len();
if let Some(n) = source.size_bytes {
ensure!(n == object.size, "local size mismatch");
}
object.local = Some(file);
}
Ok(object)
}
pub fn range(&mut self, start: u64, end: u64) -> Result<Vec<u8>> {
ensure!(start <= end && end <= self.size, "invalid byte range");
let mut body = Vec::with_capacity((end - start).try_into()?);
if start == end {
return Ok(body);
}
if let Some(file) = &mut self.local {
file.seek(SeekFrom::Start(start))?;
file.take(end - start).read_to_end(&mut body)?;
ensure!(body.len() as u64 == end - start, "truncated local object");
return Ok(body);
}
let (store, path) = self.remote.as_ref().unwrap();
let etag = self.etag.clone();
self.runtime.block_on(async {
let ranges = (start..end).step_by(2 * 1024 * 1024).map(|offset| {
fetch(
store.as_ref(),
path,
etag.clone(),
offset,
(offset + 2 * 1024 * 1024).min(end),
)
});
let mut pending = futures::stream::iter(ranges).buffered(2);
while let Some((part, observed)) = pending.try_next().await? {
if let (Some(expected), Some(actual)) = (&self.etag, &observed) {
ensure!(expected == actual, "object changed between ranges");
}
if self.etag.is_none() {
self.etag = observed;
}
body.extend_from_slice(&part);
}
Ok(body)
})
}
}
async fn fetch(
store: &dyn ObjectStore,
path: &Path,
mut etag: Option<String>,
start: u64,
end: u64,
) -> Result<(Vec<u8>, Option<String>)> {
let _timer = crate::profile::Timer::new(0);
let mut body = Vec::with_capacity((end - start) as usize);
let mut failures = 0u32;
let mut progress = Instant::now();
while start + (body.len() as u64) < end {
let result: Result<()> = async {
let options = GetOptions {
range: Some((start + body.len() as u64..end).into()),
if_match: etag.clone(),
..Default::default()
};
let request = crate::profile::Request::new(end - start - body.len() as u64);
let response = store.get_opts(path, options).await?;
request.headers();
if etag.is_none() {
etag = response.meta.e_tag.clone();
}
let mut stream = response.into_stream();
while let Some(chunk) = stream.try_next().await? {
if chunk.is_empty() {
continue;
}
ensure!(
body.len() as u64 + chunk.len() as u64 <= end - start,
"response exceeded range"
);
request.bytes(chunk.len());
body.extend_from_slice(&chunk);
failures = 0;
progress = Instant::now();
}
if body.len() as u64 != end - start {
return Err(Incomplete.into());
}
Ok(())
}
.await;
if let Err(error) = result {
failures += 1;
retry(error, failures, progress).await?;
}
}
Ok((body, etag))
}
pub struct Input {
object: Arc<Mutex<Object>>,
position: u64,
buffer: std::io::Cursor<Vec<u8>>,
}
impl Input {
pub fn new(object: Object) -> Self {
Self::at(Arc::new(Mutex::new(object)), 0)
}
pub fn at(object: Arc<Mutex<Object>>, position: u64) -> Self {
Self {
object,
position,
buffer: std::io::Cursor::new(Vec::new()),
}
}
}
impl Read for Input {
fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
if out.is_empty() {
return Ok(0);
}
if self.buffer.position() == self.buffer.get_ref().len() as u64 {
let mut object = self.object.lock().unwrap();
let end = (self.position + 4 * 1024 * 1024).min(object.size);
let data = object
.range(self.position, end)
.map_err(std::io::Error::other)?;
self.position = end;
self.buffer = std::io::Cursor::new(data);
}
self.buffer.read(out)
}
}