o41_training_data/
storage.rs

1use crate::config::{Source, Storage};
2use anyhow::{Result, bail, ensure};
3use bytes::Bytes;
4use futures::{StreamExt, TryStreamExt};
5use object_store::{
6    ClientOptions, GetOptions, ObjectStore, ObjectStoreExt, RetryConfig, aws::AmazonS3Builder,
7    path::Path,
8};
9use std::{
10    collections::HashMap,
11    io::Read,
12    ops::Range,
13    os::unix::fs::FileExt,
14    sync::{Arc, Mutex},
15    time::{Duration, Instant},
16};
17
18// Object stores answer a range request after ~100-200 ms, so throughput comes
19// from many requests in flight. Tuned from GCP us-central1 against Tigris.
20/// Largest single range request. A range is split into parts this size.
21const PART: u64 = 4 * 1024 * 1024;
22/// Part requests in flight for one list of ranges.
23const PARALLEL: usize = 8;
24/// Lists of ranges a prefetch fetches ahead of its reader.
25const AHEAD: usize = 4;
26/// Most lists a prefetch holds: those in flight, one queued and one read.
27pub const HELD: usize = AHEAD + 2;
28/// Bytes of in-flight parts not yet appended to their list.
29pub const PARTS: u64 = (AHEAD * PARALLEL) as u64 * PART;
30/// Window a streamed object is fetched in.
31pub const WINDOW: u64 = 8 * 1024 * 1024;
32
33#[derive(Debug)]
34struct Incomplete;
35impl std::fmt::Display for Incomplete {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.write_str("incomplete response")
38    }
39}
40impl std::error::Error for Incomplete {}
41async fn retry(error: anyhow::Error, failures: u32, progress: Instant) -> Result<()> {
42    let transient = error.is::<Incomplete>()
43        || error
44            .downcast_ref::<object_store::Error>()
45            .is_some_and(|e| {
46                matches!(
47                    e,
48                    object_store::Error::Generic { .. } | object_store::Error::JoinError { .. }
49                )
50            });
51    if !transient {
52        return Err(error);
53    }
54    if failures >= 64 || progress.elapsed().as_secs() >= 900 {
55        bail!("storage no-progress retry budget exhausted: {error}");
56    }
57    let mut random = [0u8; 8];
58    getrandom::fill(&mut random).map_err(|e| anyhow::anyhow!("retry jitter: {e}"))?;
59    let fraction = (u64::from_le_bytes(random) >> 11) as f64 / ((1u64 << 53) as f64);
60    let cap = (2f64.powi((failures - 1).min(5) as i32)).min(30.0);
61    let delay = (fraction * cap).min((900.0 - progress.elapsed().as_secs_f64()).max(0.0));
62    tokio::time::sleep(Duration::from_secs_f64(delay)).await;
63    Ok(())
64}
65
66pub struct Client {
67    config: Storage,
68    pub runtime: Arc<tokio::runtime::Runtime>,
69    stores: Mutex<HashMap<String, Arc<dyn ObjectStore>>>,
70}
71impl Client {
72    pub fn new(config: Storage, runtime: Arc<tokio::runtime::Runtime>) -> Self {
73        Self {
74            config,
75            runtime,
76            stores: Mutex::new(HashMap::new()),
77        }
78    }
79    /// The object store of one bucket, built once per client with the
80    /// configured endpoint, credentials, timeouts and retries.
81    pub fn store(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>> {
82        let mut stores = self.stores.lock().unwrap();
83        if let Some(store) = stores.get(bucket) {
84            return Ok(store.clone());
85        }
86        let (Some(base_url), Some(region), Some(key), Some(secret)) = (
87            self.config.base_url.as_ref(),
88            self.config.region.as_ref(),
89            self.config.access_key_id.as_ref(),
90            self.config.secret_access_key.as_ref(),
91        ) else {
92            bail!("s3 URIs need storage baseUrl, region, accessKeyId and secretAccessKey");
93        };
94        let builder = AmazonS3Builder::new()
95            .with_bucket_name(bucket)
96            .with_region(region)
97            .with_endpoint(base_url)
98            .with_access_key_id(key)
99            .with_secret_access_key(secret)
100            .with_virtual_hosted_style_request(self.config.virtual_hosted_style_request)
101            .with_client_options(
102                ClientOptions::new()
103                    .with_allow_http(true)
104                    .with_pool_idle_timeout(Duration::from_secs(90))
105                    .with_pool_max_idle_per_host(64)
106                    .with_connect_timeout(Duration::from_secs(10))
107                    .with_timeout(Duration::from_secs(60)),
108            )
109            .with_retry(RetryConfig {
110                max_retries: 3,
111                retry_timeout: Duration::from_secs(60),
112                ..Default::default()
113            });
114        let store: Arc<dyn ObjectStore> = Arc::new(builder.build()?);
115
116        stores.insert(bucket.into(), store.clone());
117        Ok(store)
118    }
119}
120
121#[derive(Clone)]
122enum Location {
123    Local(Arc<std::fs::File>),
124    Remote(Arc<dyn ObjectStore>, Path),
125}
126pub struct Object {
127    location: Location,
128    runtime: Arc<tokio::runtime::Runtime>,
129    pub size: u64,
130    etag: Option<String>,
131}
132impl Object {
133    pub fn open(source: &Source, client: &Client) -> Result<Self> {
134        let runtime = client.runtime.clone();
135        let etag = source.etag.clone().filter(|e| !e.is_empty());
136        if let Some(uri) = source.uri.strip_prefix("s3://") {
137            let (bucket, key) = uri
138                .split_once('/')
139                .ok_or_else(|| anyhow::anyhow!("invalid S3 URI"))?;
140            let store = client.store(bucket)?;
141            let path = Path::from(key);
142            let size = if let Some(n) = source.size_bytes {
143                n
144            } else {
145                runtime.block_on(async {
146                    let progress = Instant::now();
147                    let mut failures = 0;
148                    loop {
149                        match store.head(&path).await {
150                            Ok(meta) => return Ok::<_, anyhow::Error>(meta.size),
151                            Err(error) => {
152                                failures += 1;
153                                retry(error.into(), failures, progress).await?;
154                            }
155                        }
156                    }
157                })?
158            };
159            return Ok(Self {
160                location: Location::Remote(store, path),
161                runtime,
162                size,
163                etag,
164            });
165        }
166        let file = std::fs::File::open(source.uri.strip_prefix("file://").unwrap_or(&source.uri))?;
167        let size = file.metadata()?.len();
168        if let Some(n) = source.size_bytes {
169            ensure!(n == size, "local size mismatch");
170        }
171        Ok(Self {
172            location: Location::Local(Arc::new(file)),
173            runtime,
174            size,
175            etag,
176        })
177    }
178    pub fn range(&mut self, start: u64, end: u64) -> Result<Bytes> {
179        ensure!(start <= end && end <= self.size, "invalid byte range");
180        let (mut bodies, observed) = self.runtime.block_on(fetch_ranges(
181            self.location.clone(),
182            self.etag.clone(),
183            std::iter::once(start..end).collect(),
184        ))?;
185        if self.etag.is_none() {
186            self.etag = observed;
187        }
188        Ok(bodies.pop().unwrap())
189    }
190    /// Fetches each list of ranges in the background, in order, at most
191    /// `AHEAD` lists ahead of the reader. Dropping the prefetch cancels it.
192    pub fn prefetch(&self, lists: Vec<Vec<Range<u64>>>) -> Result<Prefetch> {
193        for range in lists.iter().flatten() {
194            ensure!(
195                range.start <= range.end && range.end <= self.size,
196                "invalid byte range"
197            );
198        }
199        let (send, receiver) = tokio::sync::mpsc::channel(1);
200        let (location, etag) = (self.location.clone(), self.etag.clone());
201        let task = self.runtime.spawn(async move {
202            let mut fetched = futures::stream::iter(lists)
203                .map(|ranges| fetch_ranges(location.clone(), etag.clone(), ranges))
204                .buffered(AHEAD);
205            while let Some(result) = fetched.next().await {
206                let failed = result.is_err();
207                if send.send(result).await.is_err() || failed {
208                    return;
209                }
210            }
211        });
212        Ok(Prefetch {
213            receiver,
214            task,
215            etag: self.etag.clone(),
216        })
217    }
218}
219
220type Fetched = Result<(Vec<Bytes>, Option<String>)>;
221pub struct Prefetch {
222    receiver: tokio::sync::mpsc::Receiver<Fetched>,
223    task: tokio::task::JoinHandle<()>,
224    etag: Option<String>,
225}
226impl Iterator for Prefetch {
227    type Item = Result<Vec<Bytes>>;
228    fn next(&mut self) -> Option<Self::Item> {
229        let wait = crate::profile::Timer::new(11);
230        let fetched = self.receiver.blocking_recv()?;
231        drop(wait);
232        Some(fetched.and_then(|(bodies, observed)| {
233            if let (Some(expected), Some(actual)) = (&self.etag, &observed) {
234                ensure!(expected == actual, "object changed between ranges");
235            }
236            if self.etag.is_none() {
237                self.etag = observed;
238            }
239            Ok(bodies)
240        }))
241    }
242}
243impl Drop for Prefetch {
244    fn drop(&mut self) {
245        self.task.abort();
246    }
247}
248
249/// Fetches byte ranges as concurrent parts of at most `PART` bytes. Every part
250/// comes from one object version.
251async fn fetch_ranges(
252    location: Location,
253    mut etag: Option<String>,
254    ranges: Vec<Range<u64>>,
255) -> Fetched {
256    let (store, path) = match location {
257        Location::Remote(store, path) => (store, path),
258        Location::Local(file) => {
259            let bodies = tokio::task::spawn_blocking(move || {
260                ranges
261                    .iter()
262                    .map(|range| {
263                        let mut body = vec![0; (range.end - range.start).try_into()?];
264                        file.read_exact_at(&mut body, range.start)?;
265                        Ok(Bytes::from(body))
266                    })
267                    .collect::<Result<Vec<_>>>()
268            })
269            .await??;
270            return Ok((bodies, None));
271        }
272    };
273    let parts: Vec<_> = ranges
274        .iter()
275        .flat_map(|range| {
276            (range.start..range.end)
277                .step_by(PART as usize)
278                .map(move |offset| (offset, (offset + PART).min(range.end)))
279        })
280        .collect();
281    let requested = etag.clone();
282    let mut pending = futures::stream::iter(parts)
283        .map(move |(start, end)| {
284            let (store, path, etag) = (store.clone(), path.clone(), requested.clone());
285            async move { fetch(store.as_ref(), &path, etag, start, end).await }
286        })
287        .buffered(PARALLEL);
288    let mut bodies = Vec::with_capacity(ranges.len());
289    for range in &ranges {
290        let mut body = Vec::with_capacity((range.end - range.start).try_into()?);
291        while (body.len() as u64) < range.end - range.start {
292            let (part, observed) = pending.try_next().await?.unwrap();
293            if let (Some(expected), Some(actual)) = (&etag, &observed) {
294                ensure!(expected == actual, "object changed between ranges");
295            }
296            if etag.is_none() {
297                etag = observed;
298            }
299            body.extend_from_slice(&part);
300        }
301        bodies.push(Bytes::from(body));
302    }
303    Ok((bodies, etag))
304}
305
306async fn fetch(
307    store: &dyn ObjectStore,
308    path: &Path,
309    mut etag: Option<String>,
310    start: u64,
311    end: u64,
312) -> Result<(Vec<u8>, Option<String>)> {
313    let _timer = crate::profile::Timer::new(0);
314    let mut body = Vec::with_capacity((end - start) as usize);
315    let mut failures = 0u32;
316    let mut progress = Instant::now();
317    while start + (body.len() as u64) < end {
318        let result: Result<()> = async {
319            let options = GetOptions {
320                range: Some((start + body.len() as u64..end).into()),
321                if_match: etag.clone(),
322                ..Default::default()
323            };
324            let request = crate::profile::Request::new(end - start - body.len() as u64);
325            let response = store.get_opts(path, options).await?;
326            request.headers();
327            if etag.is_none() {
328                etag = response.meta.e_tag.clone();
329            }
330            let mut stream = response.into_stream();
331            while let Some(chunk) = stream.try_next().await? {
332                if chunk.is_empty() {
333                    continue;
334                }
335                ensure!(
336                    body.len() as u64 + chunk.len() as u64 <= end - start,
337                    "response exceeded range"
338                );
339                request.bytes(chunk.len());
340                body.extend_from_slice(&chunk);
341                failures = 0;
342                progress = Instant::now();
343            }
344            if body.len() as u64 != end - start {
345                return Err(Incomplete.into());
346            }
347            Ok(())
348        }
349        .await;
350        if let Err(error) = result {
351            failures += 1;
352            retry(error, failures, progress).await?;
353        }
354    }
355    Ok((body, etag))
356}
357
358/// A sequential reader over an object: `stream` fetches the whole object
359/// ahead in the background; `at` fetches one window per refill, on demand.
360/// Bounded range requests preserve decoder state across transient network
361/// errors.
362pub struct Input {
363    object: Arc<Mutex<Object>>,
364    position: u64,
365    size: u64,
366    window: u64,
367    buffer: std::io::Cursor<Bytes>,
368    ahead: Option<Prefetch>,
369}
370impl Input {
371    pub fn stream(object: Object) -> Result<Self> {
372        let windows = (0..object.size)
373            .step_by(WINDOW as usize)
374            .map(|start| std::iter::once(start..(start + WINDOW).min(object.size)).collect())
375            .collect();
376        let ahead = Some(object.prefetch(windows)?);
377        let mut input = Self::at(Arc::new(Mutex::new(object)), 0, WINDOW);
378        input.ahead = ahead;
379        Ok(input)
380    }
381    pub fn at(object: Arc<Mutex<Object>>, position: u64, window: u64) -> Self {
382        let size = object.lock().unwrap().size;
383        Self {
384            object,
385            position,
386            size,
387            window,
388            buffer: std::io::Cursor::new(Bytes::new()),
389            ahead: None,
390        }
391    }
392}
393impl Read for Input {
394    fn read(&mut self, out: &mut [u8]) -> std::io::Result<usize> {
395        if out.is_empty() {
396            return Ok(0);
397        }
398        if self.buffer.position() == self.buffer.get_ref().len() as u64 && self.position < self.size
399        {
400            let data = match &mut self.ahead {
401                Some(ahead) => ahead
402                    .next()
403                    .ok_or_else(|| std::io::Error::other("object prefetch stopped"))?
404                    .map_err(std::io::Error::other)?
405                    .pop()
406                    .unwrap(),
407                None => {
408                    let end = (self.position + self.window).min(self.size);
409                    self.object
410                        .lock()
411                        .unwrap()
412                        .range(self.position, end)
413                        .map_err(std::io::Error::other)?
414                }
415            };
416            self.position += data.len() as u64;
417            self.buffer = std::io::Cursor::new(data);
418        }
419        self.buffer.read(out)
420    }
421}
o41_training_data/storage.rs

o41_training_data/
storage.rs
