diff --git a/libs/remote_storage/src/azure_blob.rs b/libs/remote_storage/src/azure_blob.rs index 24c1248304..220d4ef115 100644 --- a/libs/remote_storage/src/azure_blob.rs +++ b/libs/remote_storage/src/azure_blob.rs @@ -29,6 +29,7 @@ use http_types::{StatusCode, Url}; use tokio_util::sync::CancellationToken; use tracing::debug; +use crate::RemoteStorageActivity; use crate::{ error::Cancelled, s3_bucket::RequestKind, AzureConfig, ConcurrencyLimiter, Download, DownloadError, Listing, ListingMode, RemotePath, RemoteStorage, StorageMetadata, @@ -525,6 +526,10 @@ impl RemoteStorage for AzureBlobStorage { // https://learn.microsoft.com/en-us/azure/storage/blobs/point-in-time-restore-overview Err(TimeTravelError::Unimplemented) } + + fn activity(&self) -> RemoteStorageActivity { + self.concurrency_limiter.activity() + } } pin_project_lite::pin_project! { diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs index 708662f20f..f024021507 100644 --- a/libs/remote_storage/src/lib.rs +++ b/libs/remote_storage/src/lib.rs @@ -263,6 +263,17 @@ pub trait RemoteStorage: Send + Sync + 'static { done_if_after: SystemTime, cancel: &CancellationToken, ) -> Result<(), TimeTravelError>; + + /// Query how busy we currently are: may be used by callers which wish to politely + /// back off if there are already a lot of operations underway. + fn activity(&self) -> RemoteStorageActivity; +} + +pub struct RemoteStorageActivity { + pub read_available: usize, + pub read_total: usize, + pub write_available: usize, + pub write_total: usize, } /// DownloadStream is sensitive to the timeout and cancellation used with the original @@ -444,6 +455,15 @@ impl GenericRemoteStorage> { } } } + + pub fn activity(&self) -> RemoteStorageActivity { + match self { + Self::LocalFs(s) => s.activity(), + Self::AwsS3(s) => s.activity(), + Self::AzureBlob(s) => s.activity(), + Self::Unreliable(s) => s.activity(), + } + } } impl GenericRemoteStorage { @@ -774,6 +794,9 @@ struct ConcurrencyLimiter { // The helps to ensure we don't exceed the thresholds. write: Arc, read: Arc, + + write_total: usize, + read_total: usize, } impl ConcurrencyLimiter { @@ -802,10 +825,21 @@ impl ConcurrencyLimiter { Arc::clone(self.for_kind(kind)).acquire_owned().await } + fn activity(&self) -> RemoteStorageActivity { + RemoteStorageActivity { + read_available: self.read.available_permits(), + read_total: self.read_total, + write_available: self.write.available_permits(), + write_total: self.write_total, + } + } + fn new(limit: usize) -> ConcurrencyLimiter { Self { read: Arc::new(Semaphore::new(limit)), write: Arc::new(Semaphore::new(limit)), + read_total: limit, + write_total: limit, } } } diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index 1f7bcfc982..f12f6590a3 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -23,8 +23,8 @@ use tokio_util::{io::ReaderStream, sync::CancellationToken}; use utils::crashsafe::path_with_suffix_extension; use crate::{ - Download, DownloadError, Listing, ListingMode, RemotePath, TimeTravelError, TimeoutOrCancel, - REMOTE_STORAGE_PREFIX_SEPARATOR, + Download, DownloadError, Listing, ListingMode, RemotePath, RemoteStorageActivity, + TimeTravelError, TimeoutOrCancel, REMOTE_STORAGE_PREFIX_SEPARATOR, }; use super::{RemoteStorage, StorageMetadata}; @@ -605,6 +605,16 @@ impl RemoteStorage for LocalFs { ) -> Result<(), TimeTravelError> { Err(TimeTravelError::Unimplemented) } + + fn activity(&self) -> RemoteStorageActivity { + // LocalFS has no concurrency limiting: give callers the impression that plenty of units are available + RemoteStorageActivity { + read_available: 16, + read_total: 16, + write_available: 16, + write_total: 16, + } + } } fn storage_metadata_path(original_path: &Utf8Path) -> Utf8PathBuf { diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index c3d6c75e20..0f6772b274 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -47,8 +47,8 @@ use utils::backoff; use super::StorageMetadata; use crate::{ error::Cancelled, support::PermitCarrying, ConcurrencyLimiter, Download, DownloadError, - Listing, ListingMode, RemotePath, RemoteStorage, S3Config, TimeTravelError, TimeoutOrCancel, - MAX_KEYS_PER_DELETE, REMOTE_STORAGE_PREFIX_SEPARATOR, + Listing, ListingMode, RemotePath, RemoteStorage, RemoteStorageActivity, S3Config, + TimeTravelError, TimeoutOrCancel, MAX_KEYS_PER_DELETE, REMOTE_STORAGE_PREFIX_SEPARATOR, }; pub(super) mod metrics; @@ -975,6 +975,10 @@ impl RemoteStorage for S3Bucket { } Ok(()) } + + fn activity(&self) -> RemoteStorageActivity { + self.concurrency_limiter.activity() + } } /// On drop (cancellation) count towards [`metrics::BucketMetrics::cancelled_waits`]. diff --git a/libs/remote_storage/src/simulate_failures.rs b/libs/remote_storage/src/simulate_failures.rs index c467a2d196..66522e04ca 100644 --- a/libs/remote_storage/src/simulate_failures.rs +++ b/libs/remote_storage/src/simulate_failures.rs @@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken; use crate::{ Download, DownloadError, GenericRemoteStorage, Listing, ListingMode, RemotePath, RemoteStorage, - StorageMetadata, TimeTravelError, + RemoteStorageActivity, StorageMetadata, TimeTravelError, }; pub struct UnreliableWrapper { @@ -213,4 +213,8 @@ impl RemoteStorage for UnreliableWrapper { .time_travel_recover(prefix, timestamp, done_if_after, cancel) .await } + + fn activity(&self) -> RemoteStorageActivity { + self.inner.activity() + } }