Skip to main content

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