making sure there can be only a single client task claiming more bandwidth
This commit is contained in:
@@ -2,21 +2,37 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
use si_scale::helpers::bibytes2;
|
||||
use std::sync::atomic::{AtomicI64, Ordering};
|
||||
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use time::OffsetDateTime;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct BandwidthClaimGuard {
|
||||
inner: Arc<ClientBandwidthInner>,
|
||||
}
|
||||
|
||||
impl Drop for BandwidthClaimGuard {
|
||||
fn drop(&mut self) {
|
||||
let old = self.inner.claiming_more.swap(false, Ordering::SeqCst);
|
||||
assert!(
|
||||
old,
|
||||
"critical failure: there were multiple BandwidthClaimGuard existing"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ClientBandwidth {
|
||||
inner: Arc<ClientBandwidthInner>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ClientBandwidthInner {
|
||||
/// the actual bandwidth amount (in bytes) available
|
||||
available: AtomicI64,
|
||||
|
||||
/// flag to indicate whether this client is currently in the process of claiming additional bandwidth
|
||||
claiming_more: AtomicBool,
|
||||
|
||||
/// defines the timestamp when the bandwidth information has been logged to the logs stream
|
||||
last_logged_ts: AtomicI64,
|
||||
|
||||
@@ -29,11 +45,28 @@ impl ClientBandwidth {
|
||||
ClientBandwidth {
|
||||
inner: Arc::new(ClientBandwidthInner {
|
||||
available: AtomicI64::new(0),
|
||||
claiming_more: AtomicBool::new(false),
|
||||
last_logged_ts: AtomicI64::new(0),
|
||||
last_updated_ts: AtomicI64::new(0),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn begin_bandwidth_claim(&self) -> Option<BandwidthClaimGuard> {
|
||||
if self
|
||||
.inner
|
||||
.claiming_more
|
||||
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_ok()
|
||||
{
|
||||
Some(BandwidthClaimGuard {
|
||||
inner: self.inner.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn remaining(&self) -> i64 {
|
||||
self.inner.available.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
@@ -598,6 +598,11 @@ impl<C, St> GatewayClient<C, St> {
|
||||
return Err(GatewayClientError::NoBandwidthControllerAvailable);
|
||||
}
|
||||
|
||||
let Some(_claim_guard) = self.bandwidth.begin_bandwidth_claim() else {
|
||||
debug!("there's already an existing bandwidth claim ongoing");
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
warn!("Not enough bandwidth. Trying to get more bandwidth, this might take a while");
|
||||
if !self.cfg.bandwidth.require_tickets {
|
||||
info!("The client is running in disabled credentials mode - attempting to claim bandwidth without a credential");
|
||||
|
||||
Reference in New Issue
Block a user