o41_training_data/

storage.rs

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()),
        }
    }
    /// The object store of one bucket, built once per client with the
    /// configured endpoint, credentials, timeouts and retries.
    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))
}

// Bounded range requests preserve decoder state across transient network errors.
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)
    }
}