icechunk_storage/
storage.rs

1//! Core storage trait and shared types.
2
3use chrono::{DateTime, Utc};
4use core::fmt;
5use futures::{
6    Stream, StreamExt as _, TryStreamExt as _,
7    future::ready,
8    stream::{self, BoxStream, FuturesOrdered},
9};
10use itertools::Itertools as _;
11use serde::{Deserialize, Serialize};
12use std::{
13    cmp::{max, min},
14    collections::HashSet,
15    ffi::OsString,
16    fmt::Display,
17    iter,
18    num::{NonZeroU16, NonZeroU64},
19    ops::Range,
20    pin::Pin,
21    sync::{Arc, Mutex, OnceLock},
22};
23use tokio::io::AsyncBufRead;
24use tokio_util::io::StreamReader;
25use tracing::{instrument, warn};
26
27use async_trait::async_trait;
28use bytes::Bytes;
29use thiserror::Error;
30
31use crate::ICError;
32use crate::sealed;
33
34/// Storage operation error types.
35#[derive(Debug, Error)]
36pub enum StorageErrorKind {
37    #[error("object not found")]
38    ObjectNotFound,
39    #[error("bad object store prefix {0:?}")]
40    BadPrefix(OsString),
41    #[error("object store error {0}")]
42    ObjectStore(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
43    #[error("I/O error: {0}")]
44    IOError(#[from] std::io::Error),
45    #[error("storage configuration error: {0}")]
46    R2ConfigurationError(String),
47    #[error("error parsing URL: {url:?}")]
48    CannotParseUrl {
49        #[source]
50        cause: url::ParseError,
51        url: String,
52    },
53    #[error("Redirect Storage error: {0}")]
54    BadRedirect(String),
55    #[error("storage error: {0}")]
56    Other(String),
57}
58pub type StorageError = ICError<StorageErrorKind>;
59
60pub type StorageResult<A> = Result<A, StorageError>;
61
62pub fn obj_store_error_res<T>(
63    err: impl std::error::Error + Send + Sync + 'static,
64) -> StorageResult<T> {
65    Err(obj_store_error(err))
66}
67
68pub fn obj_store_error(
69    err: impl std::error::Error + Send + Sync + 'static,
70) -> StorageError {
71    StorageError::capture(StorageErrorKind::ObjectStore(Box::new(err)))
72}
73
74pub fn obj_not_found_res<T>() -> StorageResult<T> {
75    Err(StorageError::capture(StorageErrorKind::ObjectNotFound))
76}
77
78pub fn other_error(s: impl Into<String>) -> StorageError {
79    StorageError::capture(StorageErrorKind::Other(s.into()))
80}
81
82#[derive(Debug)]
83pub struct ListInfo<A> {
84    pub id: A,
85    pub created_at: DateTime<Utc>,
86    pub size_bytes: u64,
87}
88
89pub use icechunk_types::ETag;
90#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default)]
91pub struct Generation(pub String);
92
93#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone)]
94pub struct VersionInfo {
95    pub etag: Option<ETag>,
96    pub generation: Option<Generation>,
97}
98
99impl VersionInfo {
100    pub fn for_creation() -> Self {
101        Self { etag: None, generation: None }
102    }
103
104    pub fn from_etag_only(etag: String) -> Self {
105        Self { etag: Some(ETag(etag)), generation: None }
106    }
107
108    pub fn is_create(&self) -> bool {
109        self.etag.is_none() && self.generation.is_none()
110    }
111
112    pub fn etag(&self) -> Option<&String> {
113        self.etag.as_ref().map(|e| &e.0)
114    }
115
116    pub fn generation(&self) -> Option<&String> {
117        self.generation.as_ref().map(|e| &e.0)
118    }
119}
120
121impl Display for VersionInfo {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        match (&self.etag, &self.generation) {
124            (Some(etag), Some(generation)) => {
125                write!(f, "etag={}, generation={}", etag.0, generation.0)
126            }
127            (Some(etag), None) => write!(f, "etag={}", etag.0),
128            (None, Some(generation)) => write!(f, "generation={}", generation.0),
129            (None, None) => write!(f, "new"),
130        }
131    }
132}
133
134#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy, Default)]
135pub struct RetriesSettings {
136    pub max_tries: Option<NonZeroU16>,
137    pub initial_backoff_ms: Option<u32>,
138    pub max_backoff_ms: Option<u32>,
139}
140
141impl RetriesSettings {
142    pub fn max_tries(&self) -> NonZeroU16 {
143        self.max_tries.unwrap_or_else(|| NonZeroU16::new(10).unwrap_or(NonZeroU16::MIN))
144    }
145
146    pub fn initial_backoff_ms(&self) -> u32 {
147        self.initial_backoff_ms.unwrap_or(100)
148    }
149
150    pub fn max_backoff_ms(&self) -> u32 {
151        self.max_backoff_ms.unwrap_or(3 * 60 * 1000)
152    }
153
154    pub fn merge(&self, other: Self) -> Self {
155        Self {
156            max_tries: other.max_tries.or(self.max_tries),
157            initial_backoff_ms: other.initial_backoff_ms.or(self.initial_backoff_ms),
158            max_backoff_ms: other.max_backoff_ms.or(self.max_backoff_ms),
159        }
160    }
161}
162
163#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy, Default)]
164pub struct TimeoutSettings {
165    pub connect_timeout_ms: Option<u32>,
166    pub read_timeout_ms: Option<u32>,
167    pub operation_timeout_ms: Option<u32>,
168    pub operation_attempt_timeout_ms: Option<u32>,
169}
170
171impl TimeoutSettings {
172    pub fn merge(&self, other: Self) -> Self {
173        Self {
174            connect_timeout_ms: other.connect_timeout_ms.or(self.connect_timeout_ms),
175            read_timeout_ms: other.read_timeout_ms.or(self.read_timeout_ms),
176            operation_timeout_ms: other
177                .operation_timeout_ms
178                .or(self.operation_timeout_ms),
179            operation_attempt_timeout_ms: other
180                .operation_attempt_timeout_ms
181                .or(self.operation_attempt_timeout_ms),
182        }
183    }
184}
185
186#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy, Default)]
187pub struct ConcurrencySettings {
188    pub max_concurrent_requests_for_object: Option<NonZeroU16>,
189    pub ideal_concurrent_request_size: Option<NonZeroU64>,
190}
191
192impl ConcurrencySettings {
193    // AWS recommendations: https://docs.aws.amazon.com/whitepapers/latest/s3-optimizing-performance-best-practices/horizontal-scaling-and-request-parallelization-for-high-throughput.html
194    // 8-16 MB requests
195    // 85-90 MB/s per request
196    // these numbers would saturate a 12.5 Gbps network
197
198    pub fn max_concurrent_requests_for_object(&self) -> NonZeroU16 {
199        self.max_concurrent_requests_for_object
200            .unwrap_or_else(|| NonZeroU16::new(18).unwrap_or(NonZeroU16::MIN))
201    }
202    pub fn ideal_concurrent_request_size(&self) -> NonZeroU64 {
203        self.ideal_concurrent_request_size.unwrap_or_else(|| {
204            NonZeroU64::new(12 * 1024 * 1024).unwrap_or(NonZeroU64::MIN)
205        })
206    }
207
208    pub fn merge(&self, other: Self) -> Self {
209        Self {
210            max_concurrent_requests_for_object: other
211                .max_concurrent_requests_for_object
212                .or(self.max_concurrent_requests_for_object),
213            ideal_concurrent_request_size: other
214                .ideal_concurrent_request_size
215                .or(self.ideal_concurrent_request_size),
216        }
217    }
218}
219
220/// Configuration for storage operations (retries, concurrency, storage classes).
221#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default)]
222pub struct Settings {
223    #[serde(default)]
224    pub concurrency: Option<ConcurrencySettings>,
225
226    #[serde(default)]
227    pub retries: Option<RetriesSettings>,
228
229    #[serde(default)]
230    pub timeouts: Option<TimeoutSettings>,
231
232    #[serde(default)]
233    pub unsafe_use_conditional_update: Option<bool>,
234
235    #[serde(default)]
236    pub unsafe_use_conditional_create: Option<bool>,
237
238    /// Whether to stamp objects with user metadata. Defaults to `true`.
239    ///
240    /// Disabling this while leaving `unsafe_use_conditional_*` enabled
241    /// silently neutralises the lost-response recovery for conditional
242    /// PUTs: the conditional headers still go out, but with no write-id
243    /// stamped a transient PUT failure can surface as a spurious
244    /// `NotOnLatestVersion` even when the write landed. Only disable for
245    /// backends that genuinely don't support user metadata.
246    #[serde(default)]
247    pub unsafe_use_metadata: Option<bool>,
248
249    #[serde(default)]
250    pub storage_class: Option<String>,
251
252    #[serde(default)]
253    pub metadata_storage_class: Option<String>,
254
255    #[serde(default)]
256    pub chunks_storage_class: Option<String>,
257
258    #[serde(default)]
259    pub minimum_size_for_multipart_upload: Option<u64>,
260}
261
262static DEFAULT_CONCURRENCY: OnceLock<ConcurrencySettings> = OnceLock::new();
263static DEFAULT_RETRIES: OnceLock<RetriesSettings> = OnceLock::new();
264
265impl Settings {
266    pub fn concurrency(&self) -> &ConcurrencySettings {
267        self.concurrency
268            .as_ref()
269            .unwrap_or_else(|| DEFAULT_CONCURRENCY.get_or_init(Default::default))
270    }
271
272    pub fn retries(&self) -> &RetriesSettings {
273        self.retries
274            .as_ref()
275            .unwrap_or_else(|| DEFAULT_RETRIES.get_or_init(Default::default))
276    }
277
278    pub fn timeouts(&self) -> Option<&TimeoutSettings> {
279        self.timeouts.as_ref()
280    }
281
282    pub fn unsafe_use_conditional_create(&self) -> bool {
283        self.unsafe_use_conditional_create.unwrap_or(true)
284    }
285
286    pub fn unsafe_use_conditional_update(&self) -> bool {
287        self.unsafe_use_conditional_update.unwrap_or(true)
288    }
289
290    pub fn unsafe_use_metadata(&self) -> bool {
291        self.unsafe_use_metadata.unwrap_or(true)
292    }
293
294    pub fn metadata_storage_class(&self) -> Option<&String> {
295        self.metadata_storage_class.as_ref().or(self.storage_class.as_ref())
296    }
297
298    pub fn storage_class(&self) -> Option<&String> {
299        self.storage_class.as_ref()
300    }
301
302    pub fn chunks_storage_class(&self) -> Option<&String> {
303        self.chunks_storage_class.as_ref().or(self.storage_class.as_ref())
304    }
305
306    pub fn minimum_size_for_multipart_upload(&self) -> u64 {
307        // per AWS  recommendation: 100 MB
308        self.minimum_size_for_multipart_upload.unwrap_or(100 * 1024 * 1024)
309    }
310
311    pub fn merge(&self, other: Self) -> Self {
312        Self {
313            concurrency: match (&self.concurrency, other.concurrency) {
314                (None, None) => None,
315                (None, Some(c)) => Some(c),
316                (Some(c), None) => Some(*c),
317                (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
318            },
319            retries: match (&self.retries, other.retries) {
320                (None, None) => None,
321                (None, Some(c)) => Some(c),
322                (Some(c), None) => Some(*c),
323                (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
324            },
325            timeouts: match (&self.timeouts, other.timeouts) {
326                (None, None) => None,
327                (None, Some(c)) => Some(c),
328                (Some(c), None) => Some(*c),
329                (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
330            },
331            unsafe_use_conditional_create: match (
332                &self.unsafe_use_conditional_create,
333                other.unsafe_use_conditional_create,
334            ) {
335                (None, None) => None,
336                (None, Some(c)) => Some(c),
337                (Some(c), None) => Some(*c),
338                (Some(_), Some(theirs)) => Some(theirs),
339            },
340            unsafe_use_conditional_update: match (
341                &self.unsafe_use_conditional_update,
342                other.unsafe_use_conditional_update,
343            ) {
344                (None, None) => None,
345                (None, Some(c)) => Some(c),
346                (Some(c), None) => Some(*c),
347                (Some(_), Some(theirs)) => Some(theirs),
348            },
349            unsafe_use_metadata: match (
350                &self.unsafe_use_metadata,
351                other.unsafe_use_metadata,
352            ) {
353                (None, None) => None,
354                (None, Some(c)) => Some(c),
355                (Some(c), None) => Some(*c),
356                (Some(_), Some(theirs)) => Some(theirs),
357            },
358            storage_class: match (&self.storage_class, other.storage_class) {
359                (None, None) => None,
360                (None, Some(c)) => Some(c),
361                (Some(c), None) => Some(c.clone()),
362                (Some(_), Some(theirs)) => Some(theirs),
363            },
364            metadata_storage_class: match (
365                &self.metadata_storage_class,
366                other.metadata_storage_class,
367            ) {
368                (None, None) => None,
369                (None, Some(c)) => Some(c),
370                (Some(c), None) => Some(c.clone()),
371                (Some(_), Some(theirs)) => Some(theirs),
372            },
373            chunks_storage_class: match (
374                &self.chunks_storage_class,
375                other.chunks_storage_class,
376            ) {
377                (None, None) => None,
378                (None, Some(c)) => Some(c),
379                (Some(c), None) => Some(c.clone()),
380                (Some(_), Some(theirs)) => Some(theirs),
381            },
382            minimum_size_for_multipart_upload: match (
383                &self.minimum_size_for_multipart_upload,
384                other.minimum_size_for_multipart_upload,
385            ) {
386                (None, None) => None,
387                (None, Some(c)) => Some(c),
388                (Some(c), None) => Some(*c),
389                (Some(_), Some(theirs)) => Some(theirs),
390            },
391        }
392    }
393}
394
395#[derive(Debug, Clone, PartialEq, Eq)]
396pub enum VersionedUpdateResult {
397    Updated { new_version: VersionInfo },
398    NotOnLatestVersion,
399}
400
401impl VersionedUpdateResult {
402    pub fn must_write(self) -> StorageResult<VersionInfo> {
403        match self {
404            VersionedUpdateResult::Updated { new_version } => Ok(new_version),
405            VersionedUpdateResult::NotOnLatestVersion => {
406                Err(StorageError::capture(StorageErrorKind::ObjectNotFound))
407            }
408        }
409    }
410}
411
412#[derive(Debug, Clone, PartialEq, Eq, Default)]
413pub struct DeleteObjectsResult {
414    pub deleted_objects: u64,
415    pub deleted_bytes: u64,
416}
417
418impl DeleteObjectsResult {
419    pub fn merge(&mut self, other: &Self) {
420        self.deleted_objects += other.deleted_objects;
421        self.deleted_bytes += other.deleted_bytes;
422    }
423}
424
425pub enum GetModifiedResult {
426    Modified { data: Pin<Box<dyn AsyncBufRead + Send>>, new_version: VersionInfo },
427    OnLatestVersion,
428}
429
430impl fmt::Debug for GetModifiedResult {
431    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
432        match self {
433            Self::Modified { new_version, .. } => {
434                f.debug_struct("Modified").field("new_version", new_version).finish()
435            }
436            Self::OnLatestVersion => write!(f, "OnLatestVersion"),
437        }
438    }
439}
440
441impl Display for GetModifiedResult {
442    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
443        match self {
444            Self::Modified { new_version, .. } => {
445                write!(f, "Modified(new_version={new_version})")
446            }
447            Self::OnLatestVersion => write!(f, "OnLatestVersion"),
448        }
449    }
450}
451
452/// Fetch and write the parquet files that represent the repository in object store
453///
454/// Different implementation can cache the files differently, or not at all.
455/// Structured metadata about a storage backend, for display/repr purposes.
456#[derive(Debug, Clone)]
457pub struct StorageInfo {
458    /// Human-readable backend type, e.g. "S3", "GCS", "in-memory".
459    pub backend_type: &'static str,
460    /// Key-value pairs of relevant configuration (bucket, prefix, path, etc.).
461    pub fields: Vec<(&'static str, String)>,
462}
463
464/// Whether a repository may be created at a storage location.
465#[derive(Debug, Clone, PartialEq, Eq)]
466pub enum RepositoryCreation {
467    /// A new repository may be created here.
468    Allowed,
469    /// Refused: this is a cloud object store addressed at an empty prefix (the
470    /// bucket root). New empty-prefix repositories are no longer supported,
471    /// pre-existing ones can still be opened and updated.
472    RefusedEmptyPrefix,
473}
474
475/// Keeps the listed objects whose id starts with a character in `first_chars`.
476pub fn filter_ids_by_first_char<'a>(
477    listing: BoxStream<'a, StorageResult<ListInfo<String>>>,
478    first_chars: &HashSet<char>,
479) -> BoxStream<'a, StorageResult<ListInfo<String>>> {
480    let first_chars = first_chars.clone();
481    listing
482        .try_filter(move |info| {
483            ready(info.id.chars().next().is_some_and(|c| first_chars.contains(&c)))
484        })
485        .boxed()
486}
487
488/// Implementations are free to assume files are never overwritten.
489#[async_trait]
490#[typetag::serde(tag = "type")]
491pub trait Storage: fmt::Debug + Display + sealed::Sealed + Sync + Send {
492    /// Return structured metadata about this storage backend for display/repr.
493    fn storage_info(&self) -> StorageInfo;
494
495    async fn default_settings(&self) -> StorageResult<Settings> {
496        Ok(Default::default())
497    }
498
499    async fn can_write(&self) -> StorageResult<bool>;
500
501    /// Whether a repository may be created at this storage location.
502    async fn can_create_repository(&self) -> StorageResult<RepositoryCreation> {
503        Ok(RepositoryCreation::Allowed)
504    }
505
506    /// Ensure the storage location is ready to receive writes.
507    ///
508    /// Called by [`Repository::create`] before any other I/O. The default
509    /// implementation is a no-op; backends that need to materialize the
510    /// location (e.g. the local filesystem creating the directory) override
511    /// this. Read paths like `open` never call it, so a missing location
512    /// stays missing.
513    async fn create_location_if_needed(&self) -> StorageResult<()> {
514        Ok(())
515    }
516
517    async fn get_object(
518        &self,
519        settings: &Settings,
520        path: &str,
521        range: Option<&Range<u64>>,
522    ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
523        if let Some(range) = range {
524            self.get_object_concurrently(settings, path, range).await
525        } else {
526            self.get_object_range_read(settings, path, range).await
527        }
528    }
529
530    async fn get_object_range_read(
531        &self,
532        settings: &Settings,
533        path: &str,
534        range: Option<&Range<u64>>,
535    ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
536        let (stream, version) = self.get_object_range(settings, path, range).await?;
537        let reader = StreamReader::new(stream.map_err(std::io::Error::other));
538        Ok((Box::pin(reader), version))
539    }
540
541    async fn get_object_range(
542        &self,
543        settings: &Settings,
544        path: &str,
545        range: Option<&Range<u64>>,
546    ) -> StorageResult<(
547        Pin<Box<dyn Stream<Item = Result<Bytes, StorageError>> + Send>>,
548        VersionInfo,
549    )>;
550
551    async fn put_object(
552        &self,
553        settings: &Settings,
554        path: &str,
555        bytes: Bytes,
556        content_type: Option<&str>,
557        metadata: Vec<(String, String)>,
558        previous_version: Option<&VersionInfo>,
559    ) -> StorageResult<VersionedUpdateResult>;
560
561    async fn copy_object(
562        &self,
563        settings: &Settings,
564        from: &str,
565        to: &str,
566        content_type: Option<&str>,
567        version: &VersionInfo,
568    ) -> StorageResult<VersionedUpdateResult>;
569
570    /// List objects in storage whose keys start with the given prefix.
571    ///
572    /// Returns a stream of [`ListInfo`] entries, each containing the object's key and size in bytes.
573    /// Pass an empty prefix to list all objects in the repository's storage root.
574    async fn list_objects<'a>(
575        &'a self,
576        settings: &Settings,
577        prefix: &str,
578    ) -> StorageResult<BoxStream<'a, StorageResult<ListInfo<String>>>>;
579
580    /// Like [`Storage::list_objects`], limited to ids whose first character is in `first_chars`.
581    /// Unordered. Override to list each character concurrently; the default filters.
582    async fn list_objects_with_id_first_chars<'a>(
583        &'a self,
584        settings: &Settings,
585        prefix: &str,
586        first_chars: &HashSet<char>,
587    ) -> StorageResult<BoxStream<'a, StorageResult<ListInfo<String>>>> {
588        Ok(filter_ids_by_first_char(
589            self.list_objects(settings, prefix).await?,
590            first_chars,
591        ))
592    }
593
594    async fn delete_batch(
595        &self,
596        settings: &Settings,
597        prefix: &str,
598        batch: Vec<(String, u64)>,
599    ) -> StorageResult<DeleteObjectsResult>;
600
601    async fn get_object_last_modified(
602        &self,
603        path: &str,
604        settings: &Settings,
605    ) -> StorageResult<DateTime<Utc>>;
606
607    async fn get_object_conditional(
608        &self,
609        settings: &Settings,
610        path: &str,
611        previous_version: Option<&VersionInfo>,
612    ) -> StorageResult<GetModifiedResult>;
613
614    /// Delete a stream of objects, by their id string representations
615    /// Input stream includes sizes to get as result the total number of bytes deleted
616    #[instrument(skip(self, settings, ids))]
617    async fn delete_objects(
618        &self,
619        settings: &Settings,
620        prefix: &str,
621        ids: BoxStream<'_, (String, u64)>,
622    ) -> StorageResult<DeleteObjectsResult> {
623        let res = Arc::new(Mutex::new(DeleteObjectsResult::default()));
624        ids.chunks(1_000)
625            // FIXME: configurable concurrency
626            .for_each_concurrent(10, |batch| {
627                let res = Arc::clone(&res);
628                async move {
629                    let new_deletes = self
630                        .delete_batch(settings, prefix, batch)
631                        .await
632                        .unwrap_or_else(|_| {
633                            // FIXME: handle error instead of skipping
634                            warn!("ignoring error in Storage::delete_batch");
635                            Default::default()
636                        });
637                    #[expect(clippy::expect_used)]
638                    res.lock().expect("Bug in delete objects").merge(&new_deletes);
639                }
640            })
641            .await;
642        #[expect(clippy::expect_used)]
643        let res = res.lock().expect("Bug in delete objects");
644        Ok(res.clone())
645    }
646
647    async fn root_is_clean(&self, settings: &Settings) -> StorageResult<bool> {
648        match self.list_objects(settings, "").await {
649            Ok(mut stream) => match stream.next().await {
650                None => Ok(true),
651                Some(Ok(_)) => Ok(false),
652                Some(Err(err)) => Err(err),
653            },
654            Err(StorageError { kind: StorageErrorKind::ObjectNotFound, .. }) => Ok(true),
655            Err(err) => Err(err),
656        }
657    }
658
659    async fn get_object_concurrently_multiple(
660        &self,
661        settings: &Settings,
662        key: &str,
663        parts: Vec<Range<u64>>,
664    ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
665        let settings2 = settings.clone();
666        let key2 = key.to_string();
667        let results = parts
668            .into_iter()
669            .map(move |range| {
670                let key = key2.clone();
671                let settings = settings2.clone();
672                async move {
673                    let (stream, version) = self
674                        .get_object_range(&settings, key.as_ref(), Some(&range))
675                        .await?;
676                    let all_bytes: Vec<_> = stream.try_collect().await?;
677                    Ok::<_, StorageError>((all_bytes, version))
678                }
679            })
680            .collect::<FuturesOrdered<_>>();
681
682        let results = results.peekable();
683        tokio::pin!(results);
684        let version = match results.as_mut().peek().await {
685            Some(Ok((_, version))) => version.clone(),
686            _ => VersionInfo::for_creation(),
687        };
688        let all_bytes = results
689            .map_ok(|(all_bytes, _)| stream::iter(all_bytes).map(Ok::<_, StorageError>))
690            .try_flatten()
691            .try_collect::<Vec<_>>()
692            .await?;
693
694        let res = StreamReader::new(stream::iter(all_bytes).map(Ok::<_, std::io::Error>));
695        Ok((Box::pin(res), version))
696    }
697
698    async fn get_object_concurrently(
699        &self,
700        settings: &Settings,
701        key: &str,
702        range: &Range<u64>,
703    ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
704        let parts = split_in_multiple_requests(
705            range,
706            settings.concurrency().ideal_concurrent_request_size().get(),
707            settings.concurrency().max_concurrent_requests_for_object().get(),
708        )
709        .collect::<Vec<_>>();
710
711        let res: (Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo) = match parts.len() {
712            0 => (Box::pin(tokio::io::empty()), VersionInfo::for_creation()),
713            1 => self.get_object_range_read(settings, key, Some(range)).await?,
714            _ => self.get_object_concurrently_multiple(settings, key, parts).await?,
715        };
716        Ok(res)
717    }
718}
719
720/// Split an object request into multiple byte range requests
721///
722/// Returns tuples of Range for each request.
723///
724/// It generates requests that are as similar as possible in size, this means no more than 1 byte
725/// difference between the requests.
726///
727/// It tries to generate `ceil(size/ideal_req_size)` requests, but never exceeds `max_requests`.
728///
729/// `ideal_req_size` and `max_requests` must be > 0
730pub fn split_in_multiple_requests(
731    range: &Range<u64>,
732    ideal_req_size: u64,
733    max_requests: u16,
734) -> impl Iterator<Item = Range<u64>> + use<> {
735    let size = max(0, range.end - range.start);
736    // we do a ceiling division, rounding always up
737    let num_parts = size.div_ceil(ideal_req_size);
738    // no more than max_parts, so we limit
739    let num_parts = max(1, min(num_parts, max_requests as u64));
740
741    // we split the total size into request that are as similar as possible in size
742    // this means, we are going to have a few requests that are 1 byte larger
743    let big_parts = size % num_parts;
744    let small_parts_size = size / num_parts;
745    let big_parts_size = small_parts_size + 1;
746
747    iter::successors(Some((1, range.start..range.start)), move |(index, prev_range)| {
748        let size = if *index <= big_parts { big_parts_size } else { small_parts_size };
749        Some((index + 1, prev_range.end..prev_range.end + size))
750    })
751    .dropping(1)
752    .take(num_parts as usize)
753    .map(|(_, range)| range)
754}
755
756/// Split an object request into multiple byte range requests ensuring only the last request is
757/// smaller
758///
759/// Returns tuples of Range for each request.
760///
761/// It tries to generate `ceil(size/ideal_req_size)` requests, but never exceeds `max_requests`.
762///
763/// `ideal_req_size` and `max_requests` must be > 0
764pub fn split_in_multiple_equal_requests(
765    range: &Range<u64>,
766    ideal_req_size: u64,
767    max_requests: u16,
768) -> impl Iterator<Item = Range<u64>> + use<> {
769    let size = max(0, range.end - range.start);
770    // we do a ceiling division, rounding always up
771    let num_parts = size.div_ceil(ideal_req_size);
772    // no more than max_parts, so we limit
773    let num_parts = max(1, min(num_parts, max_requests as u64));
774
775    let big_parts = num_parts - 1;
776    let big_parts_size = size / max(1, big_parts);
777    let small_part_size = size - big_parts_size * big_parts;
778
779    iter::successors(Some((1, range.start..range.start)), move |(index, prev_range)| {
780        let size = if *index <= big_parts { big_parts_size } else { small_part_size };
781        Some((index + 1, prev_range.end..prev_range.end + size))
782    })
783    .dropping(1)
784    .take(num_parts as usize)
785    .map(|(_, range)| range)
786}
787
788pub fn strip_quotes(s: &str) -> &str {
789    s.strip_prefix('"').and_then(|s| s.strip_suffix('"')).unwrap_or(s)
790}
icechunk_storage/storage.rs

icechunk_storage/
storage.rs
