Line data Source code
1 : //! A set of generic storage abstractions for the page server to use when backing up and restoring its state from the external storage.
2 : //! No other modules from this tree are supposed to be used directly by the external code.
3 : //!
4 : //! [`RemoteStorage`] trait a CRUD-like generic abstraction to use for adapting external storages with a few implementations:
5 : //! * [`local_fs`] allows to use local file system as an external storage
6 : //! * [`s3_bucket`] uses AWS S3 bucket as an external storage
7 : //! * [`azure_blob`] allows to use Azure Blob storage as an external storage
8 : //!
9 : #![deny(unsafe_code)]
10 : #![deny(clippy::undocumented_unsafe_blocks)]
11 :
12 : mod azure_blob;
13 : mod error;
14 : mod local_fs;
15 : mod metrics;
16 : mod s3_bucket;
17 : mod simulate_failures;
18 : mod support;
19 :
20 : use std::{
21 : collections::HashMap,
22 : fmt::Debug,
23 : num::{NonZeroU32, NonZeroUsize},
24 : pin::Pin,
25 : str::FromStr,
26 : sync::Arc,
27 : time::{Duration, SystemTime},
28 : };
29 :
30 : use anyhow::{bail, Context};
31 : use aws_sdk_s3::types::StorageClass;
32 : use camino::{Utf8Path, Utf8PathBuf};
33 :
34 : use bytes::Bytes;
35 : use futures::stream::Stream;
36 : use serde::{Deserialize, Serialize};
37 : use tokio::sync::Semaphore;
38 : use tokio_util::sync::CancellationToken;
39 : use tracing::info;
40 :
41 : pub use self::{
42 : azure_blob::AzureBlobStorage, local_fs::LocalFs, s3_bucket::S3Bucket,
43 : simulate_failures::UnreliableWrapper,
44 : };
45 : use s3_bucket::RequestKind;
46 :
47 : /// Azure SDK's ETag type is a simple String wrapper: we use this internally instead of repeating it here.
48 : pub use azure_core::Etag;
49 :
50 : pub use error::{DownloadError, TimeTravelError, TimeoutOrCancel};
51 :
52 : /// Currently, sync happens with AWS S3, that has two limits on requests per second:
53 : /// ~200 RPS for IAM services
54 : /// <https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/UsingWithRDS.IAMDBAuth.html>
55 : /// ~3500 PUT/COPY/POST/DELETE or 5500 GET/HEAD S3 requests
56 : /// <https://aws.amazon.com/premiumsupport/knowledge-center/s3-request-limit-avoid-throttling/>
57 : pub const DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT: usize = 100;
58 : /// Set this limit analogously to the S3 limit
59 : ///
60 : /// Here, a limit of max 20k concurrent connections was noted.
61 : /// <https://learn.microsoft.com/en-us/answers/questions/1301863/is-there-any-limitation-to-concurrent-connections>
62 : pub const DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT: usize = 100;
63 : /// No limits on the client side, which currenltly means 1000 for AWS S3.
64 : /// <https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html#API_ListObjectsV2_RequestSyntax>
65 : pub const DEFAULT_MAX_KEYS_PER_LIST_RESPONSE: Option<i32> = None;
66 :
67 : /// As defined in S3 docs
68 : pub const MAX_KEYS_PER_DELETE: usize = 1000;
69 :
70 : const REMOTE_STORAGE_PREFIX_SEPARATOR: char = '/';
71 :
72 : /// Path on the remote storage, relative to some inner prefix.
73 : /// The prefix is an implementation detail, that allows representing local paths
74 : /// as the remote ones, stripping the local storage prefix away.
75 : #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
76 : pub struct RemotePath(Utf8PathBuf);
77 :
78 : impl Serialize for RemotePath {
79 0 : fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
80 0 : where
81 0 : S: serde::Serializer,
82 0 : {
83 0 : serializer.collect_str(self)
84 0 : }
85 : }
86 :
87 : impl<'de> Deserialize<'de> for RemotePath {
88 0 : fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
89 0 : where
90 0 : D: serde::Deserializer<'de>,
91 0 : {
92 0 : let str = String::deserialize(deserializer)?;
93 0 : Ok(Self(Utf8PathBuf::from(&str)))
94 0 : }
95 : }
96 :
97 : impl std::fmt::Display for RemotePath {
98 252 : fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 252 : std::fmt::Display::fmt(&self.0, f)
100 252 : }
101 : }
102 :
103 : impl RemotePath {
104 3469 : pub fn new(relative_path: &Utf8Path) -> anyhow::Result<Self> {
105 3469 : anyhow::ensure!(
106 3469 : relative_path.is_relative(),
107 6 : "Path {relative_path:?} is not relative"
108 : );
109 3463 : Ok(Self(relative_path.to_path_buf()))
110 3469 : }
111 :
112 3111 : pub fn from_string(relative_path: &str) -> anyhow::Result<Self> {
113 3111 : Self::new(Utf8Path::new(relative_path))
114 3111 : }
115 :
116 3342 : pub fn with_base(&self, base_path: &Utf8Path) -> Utf8PathBuf {
117 3342 : base_path.join(&self.0)
118 3342 : }
119 :
120 22 : pub fn object_name(&self) -> Option<&str> {
121 22 : self.0.file_name()
122 22 : }
123 :
124 99 : pub fn join(&self, path: impl AsRef<Utf8Path>) -> Self {
125 99 : Self(self.0.join(path))
126 99 : }
127 :
128 522 : pub fn get_path(&self) -> &Utf8PathBuf {
129 522 : &self.0
130 522 : }
131 :
132 0 : pub fn extension(&self) -> Option<&str> {
133 0 : self.0.extension()
134 0 : }
135 :
136 78 : pub fn strip_prefix(&self, p: &RemotePath) -> Result<&Utf8Path, std::path::StripPrefixError> {
137 78 : self.0.strip_prefix(&p.0)
138 78 : }
139 :
140 161 : pub fn add_trailing_slash(&self) -> Self {
141 161 : // Unwrap safety inputs are guararnteed to be valid UTF-8
142 161 : Self(format!("{}/", self.0).try_into().unwrap())
143 161 : }
144 : }
145 :
146 : /// We don't need callers to be able to pass arbitrary delimiters: just control
147 : /// whether listings will use a '/' separator or not.
148 : ///
149 : /// The WithDelimiter mode will populate `prefixes` and `keys` in the result. The
150 : /// NoDelimiter mode will only populate `keys`.
151 : pub enum ListingMode {
152 : WithDelimiter,
153 : NoDelimiter,
154 : }
155 :
156 : #[derive(Default)]
157 : pub struct Listing {
158 : pub prefixes: Vec<RemotePath>,
159 : pub keys: Vec<RemotePath>,
160 : }
161 :
162 : /// Storage (potentially remote) API to manage its state.
163 : /// This storage tries to be unaware of any layered repository context,
164 : /// providing basic CRUD operations for storage files.
165 : #[allow(async_fn_in_trait)]
166 : pub trait RemoteStorage: Send + Sync + 'static {
167 : /// List objects in remote storage, with semantics matching AWS S3's ListObjectsV2.
168 : /// (see `<https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html>`)
169 : ///
170 : /// Note that the prefix is relative to any `prefix_in_bucket` configured for the client, not
171 : /// from the absolute root of the bucket.
172 : ///
173 : /// `mode` configures whether to use a delimiter. Without a delimiter all keys
174 : /// within the prefix are listed in the `keys` of the result. With a delimiter, any "directories" at the top level of
175 : /// the prefix are returned in the `prefixes` of the result, and keys in the top level of the prefix are
176 : /// returned in `keys` ().
177 : ///
178 : /// `max_keys` controls the maximum number of keys that will be returned. If this is None, this function
179 : /// will iteratively call listobjects until it runs out of keys. Note that this is not safe to use on
180 : /// unlimted size buckets, as the full list of objects is allocated into a monolithic data structure.
181 : ///
182 : async fn list(
183 : &self,
184 : prefix: Option<&RemotePath>,
185 : _mode: ListingMode,
186 : max_keys: Option<NonZeroU32>,
187 : cancel: &CancellationToken,
188 : ) -> Result<Listing, DownloadError>;
189 :
190 : /// Streams the local file contents into remote into the remote storage entry.
191 : ///
192 : /// If the operation fails because of timeout or cancellation, the root cause of the error will be
193 : /// set to `TimeoutOrCancel`.
194 : async fn upload(
195 : &self,
196 : from: impl Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static,
197 : // S3 PUT request requires the content length to be specified,
198 : // otherwise it starts to fail with the concurrent connection count increasing.
199 : data_size_bytes: usize,
200 : to: &RemotePath,
201 : metadata: Option<StorageMetadata>,
202 : cancel: &CancellationToken,
203 : ) -> anyhow::Result<()>;
204 :
205 : /// Streams the remote storage entry contents.
206 : ///
207 : /// The returned download stream will obey initial timeout and cancellation signal by erroring
208 : /// on whichever happens first. Only one of the reasons will fail the stream, which is usually
209 : /// enough for `tokio::io::copy_buf` usage. If needed the error can be filtered out.
210 : ///
211 : /// Returns the metadata, if any was stored with the file previously.
212 : async fn download(
213 : &self,
214 : from: &RemotePath,
215 : cancel: &CancellationToken,
216 : ) -> Result<Download, DownloadError>;
217 :
218 : /// Streams a given byte range of the remote storage entry contents.
219 : ///
220 : /// The returned download stream will obey initial timeout and cancellation signal by erroring
221 : /// on whichever happens first. Only one of the reasons will fail the stream, which is usually
222 : /// enough for `tokio::io::copy_buf` usage. If needed the error can be filtered out.
223 : ///
224 : /// Returns the metadata, if any was stored with the file previously.
225 : async fn download_byte_range(
226 : &self,
227 : from: &RemotePath,
228 : start_inclusive: u64,
229 : end_exclusive: Option<u64>,
230 : cancel: &CancellationToken,
231 : ) -> Result<Download, DownloadError>;
232 :
233 : /// Delete a single path from remote storage.
234 : ///
235 : /// If the operation fails because of timeout or cancellation, the root cause of the error will be
236 : /// set to `TimeoutOrCancel`. In such situation it is unknown if the deletion went through.
237 : async fn delete(&self, path: &RemotePath, cancel: &CancellationToken) -> anyhow::Result<()>;
238 :
239 : /// Delete a multiple paths from remote storage.
240 : ///
241 : /// If the operation fails because of timeout or cancellation, the root cause of the error will be
242 : /// set to `TimeoutOrCancel`. In such situation it is unknown which deletions, if any, went
243 : /// through.
244 : async fn delete_objects<'a>(
245 : &self,
246 : paths: &'a [RemotePath],
247 : cancel: &CancellationToken,
248 : ) -> anyhow::Result<()>;
249 :
250 : /// Copy a remote object inside a bucket from one path to another.
251 : async fn copy(
252 : &self,
253 : from: &RemotePath,
254 : to: &RemotePath,
255 : cancel: &CancellationToken,
256 : ) -> anyhow::Result<()>;
257 :
258 : /// Resets the content of everything with the given prefix to the given state
259 : async fn time_travel_recover(
260 : &self,
261 : prefix: Option<&RemotePath>,
262 : timestamp: SystemTime,
263 : done_if_after: SystemTime,
264 : cancel: &CancellationToken,
265 : ) -> Result<(), TimeTravelError>;
266 : }
267 :
268 : /// DownloadStream is sensitive to the timeout and cancellation used with the original
269 : /// [`RemoteStorage::download`] request. The type yields `std::io::Result<Bytes>` to be compatible
270 : /// with `tokio::io::copy_buf`.
271 : // This has 'static because safekeepers do not use cancellation tokens (yet)
272 : pub type DownloadStream =
273 : Pin<Box<dyn Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static>>;
274 :
275 : pub struct Download {
276 : pub download_stream: DownloadStream,
277 : /// The last time the file was modified (`last-modified` HTTP header)
278 : pub last_modified: SystemTime,
279 : /// A way to identify this specific version of the resource (`etag` HTTP header)
280 : pub etag: Etag,
281 : /// Extra key-value data, associated with the current remote file.
282 : pub metadata: Option<StorageMetadata>,
283 : }
284 :
285 : impl Debug for Download {
286 0 : fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
287 0 : f.debug_struct("Download")
288 0 : .field("metadata", &self.metadata)
289 0 : .finish()
290 0 : }
291 : }
292 :
293 : /// Every storage, currently supported.
294 : /// Serves as a simple way to pass around the [`RemoteStorage`] without dealing with generics.
295 : #[derive(Clone)]
296 : // Require Clone for `Other` due to https://github.com/rust-lang/rust/issues/26925
297 : pub enum GenericRemoteStorage<Other: Clone = Arc<UnreliableWrapper>> {
298 : LocalFs(LocalFs),
299 : AwsS3(Arc<S3Bucket>),
300 : AzureBlob(Arc<AzureBlobStorage>),
301 : Unreliable(Other),
302 : }
303 :
304 : impl<Other: RemoteStorage> GenericRemoteStorage<Arc<Other>> {
305 194 : pub async fn list(
306 194 : &self,
307 194 : prefix: Option<&RemotePath>,
308 194 : mode: ListingMode,
309 194 : max_keys: Option<NonZeroU32>,
310 194 : cancel: &CancellationToken,
311 194 : ) -> anyhow::Result<Listing, DownloadError> {
312 194 : match self {
313 623 : Self::LocalFs(s) => s.list(prefix, mode, max_keys, cancel).await,
314 114 : Self::AwsS3(s) => s.list(prefix, mode, max_keys, cancel).await,
315 60 : Self::AzureBlob(s) => s.list(prefix, mode, max_keys, cancel).await,
316 0 : Self::Unreliable(s) => s.list(prefix, mode, max_keys, cancel).await,
317 : }
318 194 : }
319 :
320 : /// See [`RemoteStorage::upload`]
321 2929 : pub async fn upload(
322 2929 : &self,
323 2929 : from: impl Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static,
324 2929 : data_size_bytes: usize,
325 2929 : to: &RemotePath,
326 2929 : metadata: Option<StorageMetadata>,
327 2929 : cancel: &CancellationToken,
328 2929 : ) -> anyhow::Result<()> {
329 2929 : match self {
330 30952 : Self::LocalFs(s) => s.upload(from, data_size_bytes, to, metadata, cancel).await,
331 559 : Self::AwsS3(s) => s.upload(from, data_size_bytes, to, metadata, cancel).await,
332 282 : Self::AzureBlob(s) => s.upload(from, data_size_bytes, to, metadata, cancel).await,
333 76 : Self::Unreliable(s) => s.upload(from, data_size_bytes, to, metadata, cancel).await,
334 : }
335 2891 : }
336 :
337 58 : pub async fn download(
338 58 : &self,
339 58 : from: &RemotePath,
340 58 : cancel: &CancellationToken,
341 58 : ) -> Result<Download, DownloadError> {
342 58 : match self {
343 63 : Self::LocalFs(s) => s.download(from, cancel).await,
344 26 : Self::AwsS3(s) => s.download(from, cancel).await,
345 10 : Self::AzureBlob(s) => s.download(from, cancel).await,
346 0 : Self::Unreliable(s) => s.download(from, cancel).await,
347 : }
348 58 : }
349 :
350 15 : pub async fn download_byte_range(
351 15 : &self,
352 15 : from: &RemotePath,
353 15 : start_inclusive: u64,
354 15 : end_exclusive: Option<u64>,
355 15 : cancel: &CancellationToken,
356 15 : ) -> Result<Download, DownloadError> {
357 15 : match self {
358 0 : Self::LocalFs(s) => {
359 0 : s.download_byte_range(from, start_inclusive, end_exclusive, cancel)
360 0 : .await
361 : }
362 10 : Self::AwsS3(s) => {
363 10 : s.download_byte_range(from, start_inclusive, end_exclusive, cancel)
364 10 : .await
365 : }
366 5 : Self::AzureBlob(s) => {
367 5 : s.download_byte_range(from, start_inclusive, end_exclusive, cancel)
368 25 : .await
369 : }
370 0 : Self::Unreliable(s) => {
371 0 : s.download_byte_range(from, start_inclusive, end_exclusive, cancel)
372 0 : .await
373 : }
374 : }
375 15 : }
376 :
377 : /// See [`RemoteStorage::delete`]
378 136 : pub async fn delete(
379 136 : &self,
380 136 : path: &RemotePath,
381 136 : cancel: &CancellationToken,
382 136 : ) -> anyhow::Result<()> {
383 136 : match self {
384 2 : Self::LocalFs(s) => s.delete(path, cancel).await,
385 287 : Self::AwsS3(s) => s.delete(path, cancel).await,
386 223 : Self::AzureBlob(s) => s.delete(path, cancel).await,
387 0 : Self::Unreliable(s) => s.delete(path, cancel).await,
388 : }
389 136 : }
390 :
391 : /// See [`RemoteStorage::delete_objects`]
392 21 : pub async fn delete_objects(
393 21 : &self,
394 21 : paths: &[RemotePath],
395 21 : cancel: &CancellationToken,
396 21 : ) -> anyhow::Result<()> {
397 21 : match self {
398 6 : Self::LocalFs(s) => s.delete_objects(paths, cancel).await,
399 56 : Self::AwsS3(s) => s.delete_objects(paths, cancel).await,
400 25 : Self::AzureBlob(s) => s.delete_objects(paths, cancel).await,
401 0 : Self::Unreliable(s) => s.delete_objects(paths, cancel).await,
402 : }
403 21 : }
404 :
405 : /// See [`RemoteStorage::copy`]
406 3 : pub async fn copy_object(
407 3 : &self,
408 3 : from: &RemotePath,
409 3 : to: &RemotePath,
410 3 : cancel: &CancellationToken,
411 3 : ) -> anyhow::Result<()> {
412 3 : match self {
413 0 : Self::LocalFs(s) => s.copy(from, to, cancel).await,
414 4 : Self::AwsS3(s) => s.copy(from, to, cancel).await,
415 5 : Self::AzureBlob(s) => s.copy(from, to, cancel).await,
416 0 : Self::Unreliable(s) => s.copy(from, to, cancel).await,
417 : }
418 3 : }
419 :
420 : /// See [`RemoteStorage::time_travel_recover`].
421 6 : pub async fn time_travel_recover(
422 6 : &self,
423 6 : prefix: Option<&RemotePath>,
424 6 : timestamp: SystemTime,
425 6 : done_if_after: SystemTime,
426 6 : cancel: &CancellationToken,
427 6 : ) -> Result<(), TimeTravelError> {
428 6 : match self {
429 0 : Self::LocalFs(s) => {
430 0 : s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
431 0 : .await
432 : }
433 6 : Self::AwsS3(s) => {
434 6 : s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
435 68 : .await
436 : }
437 0 : Self::AzureBlob(s) => {
438 0 : s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
439 0 : .await
440 : }
441 0 : Self::Unreliable(s) => {
442 0 : s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
443 0 : .await
444 : }
445 : }
446 6 : }
447 : }
448 :
449 : impl GenericRemoteStorage {
450 198 : pub fn from_config(storage_config: &RemoteStorageConfig) -> anyhow::Result<Self> {
451 198 : let timeout = storage_config.timeout;
452 198 : Ok(match &storage_config.storage {
453 174 : RemoteStorageKind::LocalFs { local_path: path } => {
454 174 : info!("Using fs root '{path}' as a remote storage");
455 174 : Self::LocalFs(LocalFs::new(path.clone(), timeout)?)
456 : }
457 18 : RemoteStorageKind::AwsS3(s3_config) => {
458 18 : // The profile and access key id are only printed here for debugging purposes,
459 18 : // their values don't indicate the eventually taken choice for auth.
460 18 : let profile = std::env::var("AWS_PROFILE").unwrap_or_else(|_| "<none>".into());
461 18 : let access_key_id =
462 18 : std::env::var("AWS_ACCESS_KEY_ID").unwrap_or_else(|_| "<none>".into());
463 18 : info!("Using s3 bucket '{}' in region '{}' as a remote storage, prefix in bucket: '{:?}', bucket endpoint: '{:?}', profile: {profile}, access_key_id: {access_key_id}",
464 : s3_config.bucket_name, s3_config.bucket_region, s3_config.prefix_in_bucket, s3_config.endpoint);
465 18 : Self::AwsS3(Arc::new(S3Bucket::new(s3_config, timeout)?))
466 : }
467 6 : RemoteStorageKind::AzureContainer(azure_config) => {
468 6 : let storage_account = azure_config
469 6 : .storage_account
470 6 : .as_deref()
471 6 : .unwrap_or("<AZURE_STORAGE_ACCOUNT>");
472 6 : info!("Using azure container '{}' in account '{storage_account}' in region '{}' as a remote storage, prefix in container: '{:?}'",
473 : azure_config.container_name, azure_config.container_region, azure_config.prefix_in_container);
474 6 : Self::AzureBlob(Arc::new(AzureBlobStorage::new(azure_config, timeout)?))
475 : }
476 : })
477 198 : }
478 :
479 4 : pub fn unreliable_wrapper(s: Self, fail_first: u64) -> Self {
480 4 : Self::Unreliable(Arc::new(UnreliableWrapper::new(s, fail_first)))
481 4 : }
482 :
483 : /// See [`RemoteStorage::upload`], which this method calls with `None` as metadata.
484 1359 : pub async fn upload_storage_object(
485 1359 : &self,
486 1359 : from: impl Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static,
487 1359 : from_size_bytes: usize,
488 1359 : to: &RemotePath,
489 1359 : cancel: &CancellationToken,
490 1359 : ) -> anyhow::Result<()> {
491 1359 : self.upload(from, from_size_bytes, to, None, cancel)
492 3729 : .await
493 1350 : .with_context(|| {
494 0 : format!("Failed to upload data of length {from_size_bytes} to storage path {to:?}")
495 1350 : })
496 1350 : }
497 :
498 : /// Downloads the storage object into the `to_path` provided.
499 : /// `byte_range` could be specified to dowload only a part of the file, if needed.
500 0 : pub async fn download_storage_object(
501 0 : &self,
502 0 : byte_range: Option<(u64, Option<u64>)>,
503 0 : from: &RemotePath,
504 0 : cancel: &CancellationToken,
505 0 : ) -> Result<Download, DownloadError> {
506 0 : match byte_range {
507 0 : Some((start, end)) => self.download_byte_range(from, start, end, cancel).await,
508 0 : None => self.download(from, cancel).await,
509 : }
510 0 : }
511 : }
512 :
513 : /// Extra set of key-value pairs that contain arbitrary metadata about the storage entry.
514 : /// Immutable, cannot be changed once the file is created.
515 : #[derive(Debug, Clone, PartialEq, Eq)]
516 : pub struct StorageMetadata(HashMap<String, String>);
517 :
518 : impl<const N: usize> From<[(&str, &str); N]> for StorageMetadata {
519 0 : fn from(arr: [(&str, &str); N]) -> Self {
520 0 : let map: HashMap<String, String> = arr
521 0 : .iter()
522 0 : .map(|(k, v)| (k.to_string(), v.to_string()))
523 0 : .collect();
524 0 : Self(map)
525 0 : }
526 : }
527 :
528 : /// External backup storage configuration, enough for creating a client for that storage.
529 114 : #[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
530 : pub struct RemoteStorageConfig {
531 : /// The storage connection configuration.
532 : #[serde(flatten)]
533 : pub storage: RemoteStorageKind,
534 : /// A common timeout enforced for all requests after concurrency limiter permit has been
535 : /// acquired.
536 : #[serde(with = "humantime_serde", default = "default_timeout")]
537 : pub timeout: Duration,
538 : }
539 :
540 10 : fn default_timeout() -> Duration {
541 10 : RemoteStorageConfig::DEFAULT_TIMEOUT
542 10 : }
543 :
544 : /// A kind of a remote storage to connect to, with its connection configuration.
545 90 : #[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
546 : #[serde(untagged)]
547 : pub enum RemoteStorageKind {
548 : /// Storage based on local file system.
549 : /// Specify a root folder to place all stored files into.
550 : LocalFs { local_path: Utf8PathBuf },
551 : /// AWS S3 based storage, storing all files in the S3 bucket
552 : /// specified by the config
553 : AwsS3(S3Config),
554 : /// Azure Blob based storage, storing all files in the container
555 : /// specified by the config
556 : AzureContainer(AzureConfig),
557 : }
558 :
559 : /// AWS S3 bucket coordinates and access credentials to manage the bucket contents (read and write).
560 82 : #[derive(Clone, PartialEq, Eq, serde::Deserialize)]
561 : pub struct S3Config {
562 : /// Name of the bucket to connect to.
563 : pub bucket_name: String,
564 : /// The region where the bucket is located at.
565 : pub bucket_region: String,
566 : /// A "subfolder" in the bucket, to use the same bucket separately by multiple remote storage users at once.
567 : pub prefix_in_bucket: Option<String>,
568 : /// A base URL to send S3 requests to.
569 : /// By default, the endpoint is derived from a region name, assuming it's
570 : /// an AWS S3 region name, erroring on wrong region name.
571 : /// Endpoint provides a way to support other S3 flavors and their regions.
572 : ///
573 : /// Example: `http://127.0.0.1:5000`
574 : pub endpoint: Option<String>,
575 : /// AWS S3 has various limits on its API calls, we need not to exceed those.
576 : /// See [`DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT`] for more details.
577 : #[serde(default = "default_remote_storage_s3_concurrency_limit")]
578 : pub concurrency_limit: NonZeroUsize,
579 : #[serde(default = "default_max_keys_per_list_response")]
580 : pub max_keys_per_list_response: Option<i32>,
581 : #[serde(deserialize_with = "deserialize_storage_class", default)]
582 : pub upload_storage_class: Option<StorageClass>,
583 : }
584 :
585 10 : fn default_remote_storage_s3_concurrency_limit() -> NonZeroUsize {
586 10 : DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT
587 10 : .try_into()
588 10 : .unwrap()
589 10 : }
590 :
591 14 : fn default_max_keys_per_list_response() -> Option<i32> {
592 14 : DEFAULT_MAX_KEYS_PER_LIST_RESPONSE
593 14 : }
594 :
595 : impl Debug for S3Config {
596 0 : fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
597 0 : f.debug_struct("S3Config")
598 0 : .field("bucket_name", &self.bucket_name)
599 0 : .field("bucket_region", &self.bucket_region)
600 0 : .field("prefix_in_bucket", &self.prefix_in_bucket)
601 0 : .field("concurrency_limit", &self.concurrency_limit)
602 0 : .field(
603 0 : "max_keys_per_list_response",
604 0 : &self.max_keys_per_list_response,
605 0 : )
606 0 : .finish()
607 0 : }
608 : }
609 :
610 : /// Azure bucket coordinates and access credentials to manage the bucket contents (read and write).
611 16 : #[derive(Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
612 : pub struct AzureConfig {
613 : /// Name of the container to connect to.
614 : pub container_name: String,
615 : /// Name of the storage account the container is inside of
616 : pub storage_account: Option<String>,
617 : /// The region where the bucket is located at.
618 : pub container_region: String,
619 : /// A "subfolder" in the container, to use the same container separately by multiple remote storage users at once.
620 : pub prefix_in_container: Option<String>,
621 : /// Azure has various limits on its API calls, we need not to exceed those.
622 : /// See [`DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT`] for more details.
623 : #[serde(default = "default_remote_storage_azure_concurrency_limit")]
624 : pub concurrency_limit: NonZeroUsize,
625 : #[serde(default = "default_max_keys_per_list_response")]
626 : pub max_keys_per_list_response: Option<i32>,
627 : }
628 :
629 8 : fn default_remote_storage_azure_concurrency_limit() -> NonZeroUsize {
630 8 : NonZeroUsize::new(DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT).unwrap()
631 8 : }
632 :
633 : impl Debug for AzureConfig {
634 0 : fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
635 0 : f.debug_struct("AzureConfig")
636 0 : .field("bucket_name", &self.container_name)
637 0 : .field("storage_account", &self.storage_account)
638 0 : .field("bucket_region", &self.container_region)
639 0 : .field("prefix_in_container", &self.prefix_in_container)
640 0 : .field("concurrency_limit", &self.concurrency_limit)
641 0 : .field(
642 0 : "max_keys_per_list_response",
643 0 : &self.max_keys_per_list_response,
644 0 : )
645 0 : .finish()
646 0 : }
647 : }
648 :
649 8 : fn deserialize_storage_class<'de, D: serde::Deserializer<'de>>(
650 8 : deserializer: D,
651 8 : ) -> Result<Option<StorageClass>, D::Error> {
652 8 : Option::<String>::deserialize(deserializer).and_then(|s| {
653 8 : if let Some(s) = s {
654 : use serde::de::Error;
655 8 : let storage_class = StorageClass::from_str(&s).expect("infallible");
656 : #[allow(deprecated)]
657 8 : if matches!(storage_class, StorageClass::Unknown(_)) {
658 0 : return Err(D::Error::custom(format!(
659 0 : "Specified storage class unknown to SDK: '{s}'. Allowed values: {:?}",
660 0 : StorageClass::values()
661 0 : )));
662 8 : }
663 8 : Ok(Some(storage_class))
664 : } else {
665 0 : Ok(None)
666 : }
667 8 : })
668 8 : }
669 :
670 : impl RemoteStorageConfig {
671 : pub const DEFAULT_TIMEOUT: Duration = std::time::Duration::from_secs(120);
672 :
673 44 : pub fn from_toml(toml: &toml_edit::Item) -> anyhow::Result<Option<RemoteStorageConfig>> {
674 44 : let document: toml_edit::Document = match toml {
675 16 : toml_edit::Item::Table(toml) => toml.clone().into(),
676 28 : toml_edit::Item::Value(toml_edit::Value::InlineTable(toml)) => {
677 28 : toml.clone().into_table().into()
678 : }
679 0 : _ => bail!("toml not a table or inline table"),
680 : };
681 :
682 44 : if document.is_empty() {
683 22 : return Ok(None);
684 22 : }
685 22 :
686 22 : Ok(Some(toml_edit::de::from_document(document)?))
687 44 : }
688 : }
689 :
690 : struct ConcurrencyLimiter {
691 : // Every request to S3 can be throttled or cancelled, if a certain number of requests per second is exceeded.
692 : // Same goes to IAM, which is queried before every S3 request, if enabled. IAM has even lower RPS threshold.
693 : // The helps to ensure we don't exceed the thresholds.
694 : write: Arc<Semaphore>,
695 : read: Arc<Semaphore>,
696 : }
697 :
698 : impl ConcurrencyLimiter {
699 372 : fn for_kind(&self, kind: RequestKind) -> &Arc<Semaphore> {
700 372 : match kind {
701 31 : RequestKind::Get => &self.read,
702 153 : RequestKind::Put => &self.write,
703 30 : RequestKind::List => &self.read,
704 149 : RequestKind::Delete => &self.write,
705 3 : RequestKind::Copy => &self.write,
706 6 : RequestKind::TimeTravel => &self.write,
707 : }
708 372 : }
709 :
710 348 : async fn acquire(
711 348 : &self,
712 348 : kind: RequestKind,
713 348 : ) -> Result<tokio::sync::SemaphorePermit<'_>, tokio::sync::AcquireError> {
714 348 : self.for_kind(kind).acquire().await
715 348 : }
716 :
717 24 : async fn acquire_owned(
718 24 : &self,
719 24 : kind: RequestKind,
720 24 : ) -> Result<tokio::sync::OwnedSemaphorePermit, tokio::sync::AcquireError> {
721 24 : Arc::clone(self.for_kind(kind)).acquire_owned().await
722 24 : }
723 :
724 44 : fn new(limit: usize) -> ConcurrencyLimiter {
725 44 : Self {
726 44 : read: Arc::new(Semaphore::new(limit)),
727 44 : write: Arc::new(Semaphore::new(limit)),
728 44 : }
729 44 : }
730 : }
731 :
732 : #[cfg(test)]
733 : mod tests {
734 : use super::*;
735 :
736 12 : fn parse(input: &str) -> anyhow::Result<Option<RemoteStorageConfig>> {
737 12 : let toml = input.parse::<toml_edit::Document>().unwrap();
738 12 : RemoteStorageConfig::from_toml(toml.as_item())
739 12 : }
740 :
741 : #[test]
742 4 : fn test_object_name() {
743 4 : let k = RemotePath::new(Utf8Path::new("a/b/c")).unwrap();
744 4 : assert_eq!(k.object_name(), Some("c"));
745 :
746 4 : let k = RemotePath::new(Utf8Path::new("a/b/c/")).unwrap();
747 4 : assert_eq!(k.object_name(), Some("c"));
748 :
749 4 : let k = RemotePath::new(Utf8Path::new("a/")).unwrap();
750 4 : assert_eq!(k.object_name(), Some("a"));
751 :
752 : // XXX is it impossible to have an empty key?
753 4 : let k = RemotePath::new(Utf8Path::new("")).unwrap();
754 4 : assert_eq!(k.object_name(), None);
755 4 : }
756 :
757 : #[test]
758 4 : fn rempte_path_cannot_be_created_from_absolute_ones() {
759 4 : let err = RemotePath::new(Utf8Path::new("/")).expect_err("Should fail on absolute paths");
760 4 : assert_eq!(err.to_string(), "Path \"/\" is not relative");
761 4 : }
762 :
763 : #[test]
764 4 : fn parse_localfs_config_with_timeout() {
765 4 : let input = "local_path = '.'
766 4 : timeout = '5s'";
767 4 :
768 4 : let config = parse(input).unwrap().expect("it exists");
769 4 :
770 4 : assert_eq!(
771 4 : config,
772 4 : RemoteStorageConfig {
773 4 : storage: RemoteStorageKind::LocalFs {
774 4 : local_path: Utf8PathBuf::from(".")
775 4 : },
776 4 : timeout: Duration::from_secs(5)
777 4 : }
778 4 : );
779 4 : }
780 :
781 : #[test]
782 4 : fn test_s3_parsing() {
783 4 : let toml = "\
784 4 : bucket_name = 'foo-bar'
785 4 : bucket_region = 'eu-central-1'
786 4 : upload_storage_class = 'INTELLIGENT_TIERING'
787 4 : timeout = '7s'
788 4 : ";
789 4 :
790 4 : let config = parse(toml).unwrap().expect("it exists");
791 4 :
792 4 : assert_eq!(
793 4 : config,
794 4 : RemoteStorageConfig {
795 4 : storage: RemoteStorageKind::AwsS3(S3Config {
796 4 : bucket_name: "foo-bar".into(),
797 4 : bucket_region: "eu-central-1".into(),
798 4 : prefix_in_bucket: None,
799 4 : endpoint: None,
800 4 : concurrency_limit: default_remote_storage_s3_concurrency_limit(),
801 4 : max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
802 4 : upload_storage_class: Some(StorageClass::IntelligentTiering),
803 4 : }),
804 4 : timeout: Duration::from_secs(7)
805 4 : }
806 4 : );
807 4 : }
808 :
809 : #[test]
810 4 : fn test_azure_parsing() {
811 4 : let toml = "\
812 4 : container_name = 'foo-bar'
813 4 : container_region = 'westeurope'
814 4 : upload_storage_class = 'INTELLIGENT_TIERING'
815 4 : timeout = '7s'
816 4 : ";
817 4 :
818 4 : let config = parse(toml).unwrap().expect("it exists");
819 4 :
820 4 : assert_eq!(
821 4 : config,
822 4 : RemoteStorageConfig {
823 4 : storage: RemoteStorageKind::AzureContainer(AzureConfig {
824 4 : container_name: "foo-bar".into(),
825 4 : storage_account: None,
826 4 : container_region: "westeurope".into(),
827 4 : prefix_in_container: None,
828 4 : concurrency_limit: default_remote_storage_azure_concurrency_limit(),
829 4 : max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
830 4 : }),
831 4 : timeout: Duration::from_secs(7)
832 4 : }
833 4 : );
834 4 : }
835 : }
|
