1use chrono::{DateTime, Utc};
4use core::fmt;
5use futures::{
6 Stream, StreamExt as _, TryStreamExt as _,
7 stream::{self, BoxStream, FuturesOrdered},
8};
9use itertools::Itertools as _;
10use serde::{Deserialize, Serialize};
11use std::{
12 cmp::{max, min},
13 ffi::OsString,
14 fmt::Display,
15 iter,
16 num::{NonZeroU16, NonZeroU64},
17 ops::Range,
18 pin::Pin,
19 sync::{Arc, Mutex, OnceLock},
20};
21use tokio::io::AsyncBufRead;
22use tokio_util::io::StreamReader;
23use tracing::{instrument, warn};
24
25use async_trait::async_trait;
26use bytes::Bytes;
27use thiserror::Error;
28
29use crate::ICError;
30use crate::sealed;
31
32#[derive(Debug, Error)]
34pub enum StorageErrorKind {
35 #[error("object not found")]
36 ObjectNotFound,
37 #[error("bad object store prefix {0:?}")]
38 BadPrefix(OsString),
39 #[error("object store error {0}")]
40 ObjectStore(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
41 #[error("I/O error: {0}")]
42 IOError(#[from] std::io::Error),
43 #[error("storage configuration error: {0}")]
44 R2ConfigurationError(String),
45 #[error("error parsing URL: {url:?}")]
46 CannotParseUrl {
47 #[source]
48 cause: url::ParseError,
49 url: String,
50 },
51 #[error("Redirect Storage error: {0}")]
52 BadRedirect(String),
53 #[error("storage error: {0}")]
54 Other(String),
55}
56pub type StorageError = ICError<StorageErrorKind>;
57
58pub type StorageResult<A> = Result<A, StorageError>;
59
60pub fn obj_store_error_res<T>(
61 err: impl std::error::Error + Send + Sync + 'static,
62) -> StorageResult<T> {
63 Err(obj_store_error(err))
64}
65
66pub fn obj_store_error(
67 err: impl std::error::Error + Send + Sync + 'static,
68) -> StorageError {
69 StorageError::capture(StorageErrorKind::ObjectStore(Box::new(err)))
70}
71
72pub fn obj_not_found_res<T>() -> StorageResult<T> {
73 Err(StorageError::capture(StorageErrorKind::ObjectNotFound))
74}
75
76pub fn other_error(s: impl Into<String>) -> StorageError {
77 StorageError::capture(StorageErrorKind::Other(s.into()))
78}
79
80#[derive(Debug)]
81pub struct ListInfo<A> {
82 pub id: A,
83 pub created_at: DateTime<Utc>,
84 pub size_bytes: u64,
85}
86
87pub use icechunk_types::ETag;
88#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default)]
89pub struct Generation(pub String);
90
91#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone)]
92pub struct VersionInfo {
93 pub etag: Option<ETag>,
94 pub generation: Option<Generation>,
95}
96
97impl VersionInfo {
98 pub fn for_creation() -> Self {
99 Self { etag: None, generation: None }
100 }
101
102 pub fn from_etag_only(etag: String) -> Self {
103 Self { etag: Some(ETag(etag)), generation: None }
104 }
105
106 pub fn is_create(&self) -> bool {
107 self.etag.is_none() && self.generation.is_none()
108 }
109
110 pub fn etag(&self) -> Option<&String> {
111 self.etag.as_ref().map(|e| &e.0)
112 }
113
114 pub fn generation(&self) -> Option<&String> {
115 self.generation.as_ref().map(|e| &e.0)
116 }
117}
118
119impl Display for VersionInfo {
120 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121 match (&self.etag, &self.generation) {
122 (Some(etag), Some(generation)) => {
123 write!(f, "etag={}, generation={}", etag.0, generation.0)
124 }
125 (Some(etag), None) => write!(f, "etag={}", etag.0),
126 (None, Some(generation)) => write!(f, "generation={}", generation.0),
127 (None, None) => write!(f, "new"),
128 }
129 }
130}
131
132#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy, Default)]
133pub struct RetriesSettings {
134 pub max_tries: Option<NonZeroU16>,
135 pub initial_backoff_ms: Option<u32>,
136 pub max_backoff_ms: Option<u32>,
137}
138
139impl RetriesSettings {
140 pub fn max_tries(&self) -> NonZeroU16 {
141 self.max_tries.unwrap_or_else(|| NonZeroU16::new(10).unwrap_or(NonZeroU16::MIN))
142 }
143
144 pub fn initial_backoff_ms(&self) -> u32 {
145 self.initial_backoff_ms.unwrap_or(100)
146 }
147
148 pub fn max_backoff_ms(&self) -> u32 {
149 self.max_backoff_ms.unwrap_or(3 * 60 * 1000)
150 }
151
152 pub fn merge(&self, other: Self) -> Self {
153 Self {
154 max_tries: other.max_tries.or(self.max_tries),
155 initial_backoff_ms: other.initial_backoff_ms.or(self.initial_backoff_ms),
156 max_backoff_ms: other.max_backoff_ms.or(self.max_backoff_ms),
157 }
158 }
159}
160
161#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy, Default)]
162pub struct TimeoutSettings {
163 pub connect_timeout_ms: Option<u32>,
164 pub read_timeout_ms: Option<u32>,
165 pub operation_timeout_ms: Option<u32>,
166 pub operation_attempt_timeout_ms: Option<u32>,
167}
168
169impl TimeoutSettings {
170 pub fn merge(&self, other: Self) -> Self {
171 Self {
172 connect_timeout_ms: other.connect_timeout_ms.or(self.connect_timeout_ms),
173 read_timeout_ms: other.read_timeout_ms.or(self.read_timeout_ms),
174 operation_timeout_ms: other
175 .operation_timeout_ms
176 .or(self.operation_timeout_ms),
177 operation_attempt_timeout_ms: other
178 .operation_attempt_timeout_ms
179 .or(self.operation_attempt_timeout_ms),
180 }
181 }
182}
183
184#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy, Default)]
185pub struct ConcurrencySettings {
186 pub max_concurrent_requests_for_object: Option<NonZeroU16>,
187 pub ideal_concurrent_request_size: Option<NonZeroU64>,
188}
189
190impl ConcurrencySettings {
191 pub fn max_concurrent_requests_for_object(&self) -> NonZeroU16 {
197 self.max_concurrent_requests_for_object
198 .unwrap_or_else(|| NonZeroU16::new(18).unwrap_or(NonZeroU16::MIN))
199 }
200 pub fn ideal_concurrent_request_size(&self) -> NonZeroU64 {
201 self.ideal_concurrent_request_size.unwrap_or_else(|| {
202 NonZeroU64::new(12 * 1024 * 1024).unwrap_or(NonZeroU64::MIN)
203 })
204 }
205
206 pub fn merge(&self, other: Self) -> Self {
207 Self {
208 max_concurrent_requests_for_object: other
209 .max_concurrent_requests_for_object
210 .or(self.max_concurrent_requests_for_object),
211 ideal_concurrent_request_size: other
212 .ideal_concurrent_request_size
213 .or(self.ideal_concurrent_request_size),
214 }
215 }
216}
217
218#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default)]
220pub struct Settings {
221 #[serde(default)]
222 pub concurrency: Option<ConcurrencySettings>,
223
224 #[serde(default)]
225 pub retries: Option<RetriesSettings>,
226
227 #[serde(default)]
228 pub timeouts: Option<TimeoutSettings>,
229
230 #[serde(default)]
231 pub unsafe_use_conditional_update: Option<bool>,
232
233 #[serde(default)]
234 pub unsafe_use_conditional_create: Option<bool>,
235
236 #[serde(default)]
245 pub unsafe_use_metadata: Option<bool>,
246
247 #[serde(default)]
248 pub storage_class: Option<String>,
249
250 #[serde(default)]
251 pub metadata_storage_class: Option<String>,
252
253 #[serde(default)]
254 pub chunks_storage_class: Option<String>,
255
256 #[serde(default)]
257 pub minimum_size_for_multipart_upload: Option<u64>,
258}
259
260static DEFAULT_CONCURRENCY: OnceLock<ConcurrencySettings> = OnceLock::new();
261static DEFAULT_RETRIES: OnceLock<RetriesSettings> = OnceLock::new();
262
263impl Settings {
264 pub fn concurrency(&self) -> &ConcurrencySettings {
265 self.concurrency
266 .as_ref()
267 .unwrap_or_else(|| DEFAULT_CONCURRENCY.get_or_init(Default::default))
268 }
269
270 pub fn retries(&self) -> &RetriesSettings {
271 self.retries
272 .as_ref()
273 .unwrap_or_else(|| DEFAULT_RETRIES.get_or_init(Default::default))
274 }
275
276 pub fn timeouts(&self) -> Option<&TimeoutSettings> {
277 self.timeouts.as_ref()
278 }
279
280 pub fn unsafe_use_conditional_create(&self) -> bool {
281 self.unsafe_use_conditional_create.unwrap_or(true)
282 }
283
284 pub fn unsafe_use_conditional_update(&self) -> bool {
285 self.unsafe_use_conditional_update.unwrap_or(true)
286 }
287
288 pub fn unsafe_use_metadata(&self) -> bool {
289 self.unsafe_use_metadata.unwrap_or(true)
290 }
291
292 pub fn metadata_storage_class(&self) -> Option<&String> {
293 self.metadata_storage_class.as_ref().or(self.storage_class.as_ref())
294 }
295
296 pub fn storage_class(&self) -> Option<&String> {
297 self.storage_class.as_ref()
298 }
299
300 pub fn chunks_storage_class(&self) -> Option<&String> {
301 self.chunks_storage_class.as_ref().or(self.storage_class.as_ref())
302 }
303
304 pub fn minimum_size_for_multipart_upload(&self) -> u64 {
305 self.minimum_size_for_multipart_upload.unwrap_or(100 * 1024 * 1024)
307 }
308
309 pub fn merge(&self, other: Self) -> Self {
310 Self {
311 concurrency: match (&self.concurrency, other.concurrency) {
312 (None, None) => None,
313 (None, Some(c)) => Some(c),
314 (Some(c), None) => Some(*c),
315 (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
316 },
317 retries: match (&self.retries, other.retries) {
318 (None, None) => None,
319 (None, Some(c)) => Some(c),
320 (Some(c), None) => Some(*c),
321 (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
322 },
323 timeouts: match (&self.timeouts, other.timeouts) {
324 (None, None) => None,
325 (None, Some(c)) => Some(c),
326 (Some(c), None) => Some(*c),
327 (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
328 },
329 unsafe_use_conditional_create: match (
330 &self.unsafe_use_conditional_create,
331 other.unsafe_use_conditional_create,
332 ) {
333 (None, None) => None,
334 (None, Some(c)) => Some(c),
335 (Some(c), None) => Some(*c),
336 (Some(_), Some(theirs)) => Some(theirs),
337 },
338 unsafe_use_conditional_update: match (
339 &self.unsafe_use_conditional_update,
340 other.unsafe_use_conditional_update,
341 ) {
342 (None, None) => None,
343 (None, Some(c)) => Some(c),
344 (Some(c), None) => Some(*c),
345 (Some(_), Some(theirs)) => Some(theirs),
346 },
347 unsafe_use_metadata: match (
348 &self.unsafe_use_metadata,
349 other.unsafe_use_metadata,
350 ) {
351 (None, None) => None,
352 (None, Some(c)) => Some(c),
353 (Some(c), None) => Some(*c),
354 (Some(_), Some(theirs)) => Some(theirs),
355 },
356 storage_class: match (&self.storage_class, other.storage_class) {
357 (None, None) => None,
358 (None, Some(c)) => Some(c),
359 (Some(c), None) => Some(c.clone()),
360 (Some(_), Some(theirs)) => Some(theirs),
361 },
362 metadata_storage_class: match (
363 &self.metadata_storage_class,
364 other.metadata_storage_class,
365 ) {
366 (None, None) => None,
367 (None, Some(c)) => Some(c),
368 (Some(c), None) => Some(c.clone()),
369 (Some(_), Some(theirs)) => Some(theirs),
370 },
371 chunks_storage_class: match (
372 &self.chunks_storage_class,
373 other.chunks_storage_class,
374 ) {
375 (None, None) => None,
376 (None, Some(c)) => Some(c),
377 (Some(c), None) => Some(c.clone()),
378 (Some(_), Some(theirs)) => Some(theirs),
379 },
380 minimum_size_for_multipart_upload: match (
381 &self.minimum_size_for_multipart_upload,
382 other.minimum_size_for_multipart_upload,
383 ) {
384 (None, None) => None,
385 (None, Some(c)) => Some(c),
386 (Some(c), None) => Some(*c),
387 (Some(_), Some(theirs)) => Some(theirs),
388 },
389 }
390 }
391}
392
393#[derive(Debug, Clone, PartialEq, Eq)]
394pub enum VersionedUpdateResult {
395 Updated { new_version: VersionInfo },
396 NotOnLatestVersion,
397}
398
399impl VersionedUpdateResult {
400 pub fn must_write(self) -> StorageResult<VersionInfo> {
401 match self {
402 VersionedUpdateResult::Updated { new_version } => Ok(new_version),
403 VersionedUpdateResult::NotOnLatestVersion => {
404 Err(StorageError::capture(StorageErrorKind::ObjectNotFound))
405 }
406 }
407 }
408}
409
410#[derive(Debug, Clone, PartialEq, Eq, Default)]
411pub struct DeleteObjectsResult {
412 pub deleted_objects: u64,
413 pub deleted_bytes: u64,
414}
415
416impl DeleteObjectsResult {
417 pub fn merge(&mut self, other: &Self) {
418 self.deleted_objects += other.deleted_objects;
419 self.deleted_bytes += other.deleted_bytes;
420 }
421}
422
423pub enum GetModifiedResult {
424 Modified { data: Pin<Box<dyn AsyncBufRead + Send>>, new_version: VersionInfo },
425 OnLatestVersion,
426}
427
428impl fmt::Debug for GetModifiedResult {
429 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
430 match self {
431 Self::Modified { new_version, .. } => {
432 f.debug_struct("Modified").field("new_version", new_version).finish()
433 }
434 Self::OnLatestVersion => write!(f, "OnLatestVersion"),
435 }
436 }
437}
438
439impl Display for GetModifiedResult {
440 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
441 match self {
442 Self::Modified { new_version, .. } => {
443 write!(f, "Modified(new_version={new_version})")
444 }
445 Self::OnLatestVersion => write!(f, "OnLatestVersion"),
446 }
447 }
448}
449
450#[derive(Debug, Clone)]
455pub struct StorageInfo {
456 pub backend_type: &'static str,
458 pub fields: Vec<(&'static str, String)>,
460}
461
462#[derive(Debug, Clone, PartialEq, Eq)]
464pub enum RepositoryCreation {
465 Allowed,
467 RefusedEmptyPrefix,
471}
472
473#[async_trait]
475#[typetag::serde(tag = "type")]
476pub trait Storage: fmt::Debug + Display + sealed::Sealed + Sync + Send {
477 fn storage_info(&self) -> StorageInfo;
479
480 async fn default_settings(&self) -> StorageResult<Settings> {
481 Ok(Default::default())
482 }
483
484 async fn can_write(&self) -> StorageResult<bool>;
485
486 async fn can_create_repository(&self) -> StorageResult<RepositoryCreation> {
488 Ok(RepositoryCreation::Allowed)
489 }
490
491 async fn create_location_if_needed(&self) -> StorageResult<()> {
499 Ok(())
500 }
501
502 async fn get_object(
503 &self,
504 settings: &Settings,
505 path: &str,
506 range: Option<&Range<u64>>,
507 ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
508 if let Some(range) = range {
509 self.get_object_concurrently(settings, path, range).await
510 } else {
511 self.get_object_range_read(settings, path, range).await
512 }
513 }
514
515 async fn get_object_range_read(
516 &self,
517 settings: &Settings,
518 path: &str,
519 range: Option<&Range<u64>>,
520 ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
521 let (stream, version) = self.get_object_range(settings, path, range).await?;
522 let reader = StreamReader::new(stream.map_err(std::io::Error::other));
523 Ok((Box::pin(reader), version))
524 }
525
526 async fn get_object_range(
527 &self,
528 settings: &Settings,
529 path: &str,
530 range: Option<&Range<u64>>,
531 ) -> StorageResult<(
532 Pin<Box<dyn Stream<Item = Result<Bytes, StorageError>> + Send>>,
533 VersionInfo,
534 )>;
535
536 async fn put_object(
537 &self,
538 settings: &Settings,
539 path: &str,
540 bytes: Bytes,
541 content_type: Option<&str>,
542 metadata: Vec<(String, String)>,
543 previous_version: Option<&VersionInfo>,
544 ) -> StorageResult<VersionedUpdateResult>;
545
546 async fn copy_object(
547 &self,
548 settings: &Settings,
549 from: &str,
550 to: &str,
551 content_type: Option<&str>,
552 version: &VersionInfo,
553 ) -> StorageResult<VersionedUpdateResult>;
554
555 async fn list_objects<'a>(
560 &'a self,
561 settings: &Settings,
562 prefix: &str,
563 ) -> StorageResult<BoxStream<'a, StorageResult<ListInfo<String>>>>;
564
565 async fn delete_batch(
566 &self,
567 settings: &Settings,
568 prefix: &str,
569 batch: Vec<(String, u64)>,
570 ) -> StorageResult<DeleteObjectsResult>;
571
572 async fn get_object_last_modified(
573 &self,
574 path: &str,
575 settings: &Settings,
576 ) -> StorageResult<DateTime<Utc>>;
577
578 async fn get_object_conditional(
579 &self,
580 settings: &Settings,
581 path: &str,
582 previous_version: Option<&VersionInfo>,
583 ) -> StorageResult<GetModifiedResult>;
584
585 #[instrument(skip(self, settings, ids))]
588 async fn delete_objects(
589 &self,
590 settings: &Settings,
591 prefix: &str,
592 ids: BoxStream<'_, (String, u64)>,
593 ) -> StorageResult<DeleteObjectsResult> {
594 let res = Arc::new(Mutex::new(DeleteObjectsResult::default()));
595 ids.chunks(1_000)
596 .for_each_concurrent(10, |batch| {
598 let res = Arc::clone(&res);
599 async move {
600 let new_deletes = self
601 .delete_batch(settings, prefix, batch)
602 .await
603 .unwrap_or_else(|_| {
604 warn!("ignoring error in Storage::delete_batch");
606 Default::default()
607 });
608 #[expect(clippy::expect_used)]
609 res.lock().expect("Bug in delete objects").merge(&new_deletes);
610 }
611 })
612 .await;
613 #[expect(clippy::expect_used)]
614 let res = res.lock().expect("Bug in delete objects");
615 Ok(res.clone())
616 }
617
618 async fn root_is_clean(&self, settings: &Settings) -> StorageResult<bool> {
619 match self.list_objects(settings, "").await {
620 Ok(mut stream) => match stream.next().await {
621 None => Ok(true),
622 Some(Ok(_)) => Ok(false),
623 Some(Err(err)) => Err(err),
624 },
625 Err(StorageError { kind: StorageErrorKind::ObjectNotFound, .. }) => Ok(true),
626 Err(err) => Err(err),
627 }
628 }
629
630 async fn get_object_concurrently_multiple(
631 &self,
632 settings: &Settings,
633 key: &str,
634 parts: Vec<Range<u64>>,
635 ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
636 let settings2 = settings.clone();
637 let key2 = key.to_string();
638 let results = parts
639 .into_iter()
640 .map(move |range| {
641 let key = key2.clone();
642 let settings = settings2.clone();
643 async move {
644 let (stream, version) = self
645 .get_object_range(&settings, key.as_ref(), Some(&range))
646 .await?;
647 let all_bytes: Vec<_> = stream.try_collect().await?;
648 Ok::<_, StorageError>((all_bytes, version))
649 }
650 })
651 .collect::<FuturesOrdered<_>>();
652
653 let results = results.peekable();
654 tokio::pin!(results);
655 let version = match results.as_mut().peek().await {
656 Some(Ok((_, version))) => version.clone(),
657 _ => VersionInfo::for_creation(),
658 };
659 let all_bytes = results
660 .map_ok(|(all_bytes, _)| stream::iter(all_bytes).map(Ok::<_, StorageError>))
661 .try_flatten()
662 .try_collect::<Vec<_>>()
663 .await?;
664
665 let res = StreamReader::new(stream::iter(all_bytes).map(Ok::<_, std::io::Error>));
666 Ok((Box::pin(res), version))
667 }
668
669 async fn get_object_concurrently(
670 &self,
671 settings: &Settings,
672 key: &str,
673 range: &Range<u64>,
674 ) -> StorageResult<(Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo)> {
675 let parts = split_in_multiple_requests(
676 range,
677 settings.concurrency().ideal_concurrent_request_size().get(),
678 settings.concurrency().max_concurrent_requests_for_object().get(),
679 )
680 .collect::<Vec<_>>();
681
682 let res: (Pin<Box<dyn AsyncBufRead + Send>>, VersionInfo) = match parts.len() {
683 0 => (Box::pin(tokio::io::empty()), VersionInfo::for_creation()),
684 1 => self.get_object_range_read(settings, key, Some(range)).await?,
685 _ => self.get_object_concurrently_multiple(settings, key, parts).await?,
686 };
687 Ok(res)
688 }
689}
690
691pub fn split_in_multiple_requests(
702 range: &Range<u64>,
703 ideal_req_size: u64,
704 max_requests: u16,
705) -> impl Iterator<Item = Range<u64>> + use<> {
706 let size = max(0, range.end - range.start);
707 let num_parts = size.div_ceil(ideal_req_size);
709 let num_parts = max(1, min(num_parts, max_requests as u64));
711
712 let big_parts = size % num_parts;
715 let small_parts_size = size / num_parts;
716 let big_parts_size = small_parts_size + 1;
717
718 iter::successors(Some((1, range.start..range.start)), move |(index, prev_range)| {
719 let size = if *index <= big_parts { big_parts_size } else { small_parts_size };
720 Some((index + 1, prev_range.end..prev_range.end + size))
721 })
722 .dropping(1)
723 .take(num_parts as usize)
724 .map(|(_, range)| range)
725}
726
727pub fn split_in_multiple_equal_requests(
736 range: &Range<u64>,
737 ideal_req_size: u64,
738 max_requests: u16,
739) -> impl Iterator<Item = Range<u64>> + use<> {
740 let size = max(0, range.end - range.start);
741 let num_parts = size.div_ceil(ideal_req_size);
743 let num_parts = max(1, min(num_parts, max_requests as u64));
745
746 let big_parts = num_parts - 1;
747 let big_parts_size = size / max(1, big_parts);
748 let small_part_size = size - big_parts_size * big_parts;
749
750 iter::successors(Some((1, range.start..range.start)), move |(index, prev_range)| {
751 let size = if *index <= big_parts { big_parts_size } else { small_part_size };
752 Some((index + 1, prev_range.end..prev_range.end + size))
753 })
754 .dropping(1)
755 .take(num_parts as usize)
756 .map(|(_, range)| range)
757}
758
759pub fn strip_quotes(s: &str) -> &str {
760 s.strip_prefix('"').and_then(|s| s.strip_suffix('"')).unwrap_or(s)
761}