o41-training-data 0.1.6

use crate::config::{Source, Storage};
use anyhow::{Result, bail, ensure};
use bytes::Bytes;
use futures::{StreamExt, TryStreamExt};
use object_store::{
    ClientOptions, GetOptions, ObjectStore, ObjectStoreExt, RetryConfig, aws::AmazonS3Builder,
    path::Path,
};
use std::{
    collections::HashMap,
    io::Read,
    ops::Range,
    os::unix::fs::FileExt,
    sync::{Arc, Mutex},
    time::{Duration, Instant},
};

// Object stores answer a range request after ~100-200 ms, so throughput comes
// from many requests in flight. Tuned from GCP us-central1 against Tigris.
/// Largest single range request. A range is split into parts this size.
const PART: u64 = 4 * 1024 * 1024;
/// Part requests in flight for one list of ranges.
const PARALLEL: usize = 8;
/// Lists of ranges a prefetch fetches ahead of its reader.
const AHEAD: usize = 4;
/// Most lists a prefetch holds: those in flight, one queued and one read.
pub const HELD: usize = AHEAD + 2;
/// Bytes of in-flight parts not yet appended to their list.
pub const PARTS: u64 = (AHEAD * PARALLEL) as u64 * PART;
/// Window a streamed object is fetched in.
pub const WINDOW: u64 = 8 * 1024 * 1024;

#[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(64)
                    .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)
    }
}

#[derive(Clone)]
enum Location {
    Local(Arc<std::fs::File>),
    Remote(Arc<dyn ObjectStore>, Path),
}
pub struct Object {
    location: Location,
    runtime: Arc<tokio::runtime::Runtime>,
    pub size: u64,
    etag: Option<String>,
}
impl Object {
    pub fn open(source: &Source, client: &Client) -> Result<Self> {
        let runtime = client.runtime.clone();
        let 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);
            let size = if let Some(n) = source.size_bytes {
                n
            } else {
                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?;
                            }
                        }
                    }
                })?
            };
            return Ok(Self {
                location: Location::Remote(store, path),
                runtime,
                size,
                etag,
            });
        }
        let file = std::fs::File::open(source.uri.strip_prefix("file://").unwrap_or(&source.uri))?;
        let size = file.metadata()?.len();
        if let Some(n) = source.size_bytes {
            ensure!(n == size, "local size mismatch");
        }
        Ok(Self {
            location: Location::Local(Arc::new(file)),
            runtime,
            size,
            etag,
        })
    }
    pub fn range(&mut self, start: u64, end: u64) -> Result<Bytes> {
        ensure!(start <= end && end <= self.size, "invalid byte range");
        let (mut bodies, observed) = self.runtime.block_on(fetch_ranges(
            self.location.clone(),
            self.etag.clone(),
            std::iter::once(start..end).collect(),
        ))?;
        if self.etag.is_none() {
            self.etag = observed;
        }
        Ok(bodies.pop().unwrap())
    }
    /// Fetches each list of ranges in the background, in order, at most
    /// `AHEAD` lists ahead of the reader. Dropping the prefetch cancels it.
    pub fn prefetch(&self, lists: Vec<Vec<Range<u64>>>) -> Result<Prefetch> {
        for range in lists.iter().flatten() {
            ensure!(
                range.start <= range.end && range.end <= self.size,
                "invalid byte range"
            );
        }
        let (send, receiver) = tokio::sync::mpsc::channel(1);
        let (location, etag) = (self.location.clone(), self.etag.clone());
        let task = self.runtime.spawn(async move {
            let mut fetched = futures::stream::iter(lists)
                .map(|ranges| fetch_ranges(location.clone(), etag.clone(), ranges))
                .buffered(AHEAD);
            while let Some(result) = fetched.next().await {
                let failed = result.is_err();
                if send.send(result).await.is_err() || failed {
                    return;
                }
            }
        });
        Ok(Prefetch {
            receiver,
            task,
            etag: self.etag.clone(),
        })
    }
}

type Fetched = Result<(Vec<Bytes>, Option<String>)>;
pub struct Prefetch {
    receiver: tokio::sync::mpsc::Receiver<Fetched>,
    task: tokio::task::JoinHandle<()>,
    etag: Option<String>,
}
impl Iterator for Prefetch {
    type Item = Result<Vec<Bytes>>;
    fn next(&mut self) -> Option<Self::Item> {
        let wait = crate::profile::Timer::new(11);
        let fetched = self.receiver.blocking_recv()?;
        drop(wait);
        Some(fetched.and_then(|(bodies, observed)| {
            if let (Some(expected), Some(actual)) = (&self.etag, &observed) {
                ensure!(expected == actual, "object changed between ranges");
            }
            if self.etag.is_none() {
                self.etag = observed;
            }
            Ok(bodies)
        }))
    }
}
impl Drop for Prefetch {
    fn drop(&mut self) {
        self.task.abort();
    }
}

/// Fetches byte ranges as concurrent parts of at most `PART` bytes. Every part
/// comes from one object version.
async fn fetch_ranges(
    location: Location,
    mut etag: Option<String>,
    ranges: Vec<Range<u64>>,
) -> Fetched {
    let (store, path) = match location {
        Location::Remote(store, path) => (store, path),
        Location::Local(file) => {
            let bodies = tokio::task::spawn_blocking(move || {
                ranges
                    .iter()
                    .map(|range| {
                        let mut body = vec![0; (range.end - range.start).try_into()?];
                        file.read_exact_at(&mut body, range.start)?;
                        Ok(Bytes::from(body))
                    })
                    .collect::<Result<Vec<_>>>()
            })
            .await??;
            return Ok((bodies, None));
        }
    };
    let parts: Vec<_> = ranges
        .iter()
        .flat_map(|range| {
            (range.start..range.end)
                .step_by(PART as usize)
                .map(move |offset| (offset, (offset + PART).min(range.end)))
        })
        .collect();
    let requested = etag.clone();
    let mut pending = futures::stream::iter(parts)
        .map(move |(start, end)| {
            let (store, path, etag) = (store.clone(), path.clone(), requested.clone());
            async move { fetch(store.as_ref(), &path, etag, start, end).await }
        })
        .buffered(PARALLEL);
    let mut bodies = Vec::with_capacity(ranges.len());
    for range in &ranges {
        let mut body = Vec::with_capacity((range.end - range.start).try_into()?);
        while (body.len() as u64) < range.end - range.start {
            let (part, observed) = pending.try_next().await?.unwrap();
            if let (Some(expected), Some(actual)) = (&etag, &observed) {
                ensure!(expected == actual, "object changed between ranges");
            }
            if etag.is_none() {
                etag = observed;
            }
            body.extend_from_slice(&part);
        }
        bodies.push(Bytes::from(body));
    }
    Ok((bodies, etag))
}

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))
}

/// A sequential reader over an object: `stream` fetches the whole object
/// ahead in the background; `at` fetches one window per refill, on demand.
/// Bounded range requests preserve decoder state across transient network
/// errors.
pub struct Input {
    object: Arc<Mutex<Object>>,
    position: u64,
    size: u64,
    window: u64,
    buffer: std::io::Cursor<Bytes>,
    ahead: Option<Prefetch>,
}
impl Input {
    pub fn stream(object: Object) -> Result<Self> {
        let windows = (0..object.size)
            .step_by(WINDOW as usize)
            .map(|start| std::iter::once(start..(start + WINDOW).min(object.size)).collect())
            .collect();
        let ahead = Some(object.prefetch(windows)?);
        let mut input = Self::at(Arc::new(Mutex::new(object)), 0, WINDOW);
        input.ahead = ahead;
        Ok(input)
    }
    pub fn at(object: Arc<Mutex<Object>>, position: u64, window: u64) -> Self {
        let size = object.lock().unwrap().size;
        Self {
            object,
            position,
            size,
            window,
            buffer: std::io::Cursor::new(Bytes::new()),
            ahead: None,
        }
    }
}
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 && self.position < self.size
        {
            let data = match &mut self.ahead {
                Some(ahead) => ahead
                    .next()
                    .ok_or_else(|| std::io::Error::other("object prefetch stopped"))?
                    .map_err(std::io::Error::other)?
                    .pop()
                    .unwrap(),
                None => {
                    let end = (self.position + self.window).min(self.size);
                    self.object
                        .lock()
                        .unwrap()
                        .range(self.position, end)
                        .map_err(std::io::Error::other)?
                }
            };
            self.position += data.len() as u64;
            self.buffer = std::io::Cursor::new(data);
        }
        self.buffer.read(out)
    }
}