1use 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#[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 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
92pub 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 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#[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 #[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 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#[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 pub fn without_node(settings: &'a Settings, labels: &'a AttributionLabels) -> Self {
431 Self::new(settings, RequestAttribution::without_node(labels))
432 }
433
434 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#[derive(Debug, Clone)]
508pub struct StorageInfo {
509 pub backend_type: &'static str,
511 pub fields: Vec<(&'static str, String)>,
513}
514
515#[derive(Debug, Clone, PartialEq, Eq)]
517pub enum RepositoryCreation {
518 Allowed,
520 RefusedEmptyPrefix,
524}
525
526pub 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#[async_trait]
541#[typetag::serde(tag = "type")]
542pub trait Storage: fmt::Debug + Display + sealed::Sealed + Sync + Send {
543 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 async fn can_create_repository(&self) -> StorageResult<RepositoryCreation> {
554 Ok(RepositoryCreation::Allowed)
555 }
556
557 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 async fn list_objects<'a>(
626 &'a self,
627 ctx: &StorageContext<'_>,
628 prefix: &str,
629 ) -> StorageResult<BoxStream<'a, StorageResult<ListInfo<String>>>>;
630
631 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 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 #[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 .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 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
772pub 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 let num_parts = size.div_ceil(ideal_req_size);
790 let num_parts = max(1, min(num_parts, max_requests as u64));
792
793 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
808pub 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 let num_parts = size.div_ceil(ideal_req_size);
824 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 #[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 assert_eq!(filtered(&ids, &["0B"]), vec!["0B"]);
893 assert!(filtered(&ids, &[]).is_empty());
894 assert!(filtered(&ids, &["1"]).is_empty());
895 }
896}