From 2fa8da811784f17b7c8632e7f61684935d1db30f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C4=99drzej=20Stuczy=C5=84ski?=
 <jedrzej.stuczynski@gmail.com>
Date: Tue, 17 Sep 2024 18:39:49 +0100
Subject: [PATCH] making sure there can be only a single client task claiming
 more bandwidth

---
 .../gateway-client/src/bandwidth.rs           | 39 +++++++++++++++++--
 .../gateway-client/src/client/mod.rs          |  5 +++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/common/client-libs/gateway-client/src/bandwidth.rs b/common/client-libs/gateway-client/src/bandwidth.rs
index f232c004ec..1f45fdc8ea 100644
--- a/common/client-libs/gateway-client/src/bandwidth.rs
+++ b/common/client-libs/gateway-client/src/bandwidth.rs
@@ -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)
     }
diff --git a/common/client-libs/gateway-client/src/client/mod.rs b/common/client-libs/gateway-client/src/client/mod.rs
index d899ec2928..ad30cc2c8e 100644
--- a/common/client-libs/gateway-client/src/client/mod.rs
+++ b/common/client-libs/gateway-client/src/client/mod.rs
@@ -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");