nym/common/client-libs/gateway-client/src/bandwidth.rs

// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0

use si_scale::helpers::bibytes2;
use std::fmt::{Display, Formatter};
use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use time::OffsetDateTime;

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>,
}

// simple helper for logging purposes to accommodate 'unknown' case
pub(crate) enum UpgradeModeEnabledWrapper {
    True,
    False,
    Unknown,
}

impl From<Option<bool>> for UpgradeModeEnabledWrapper {
    fn from(value: Option<bool>) -> Self {
        match value {
            Some(true) => UpgradeModeEnabledWrapper::True,
            Some(false) => UpgradeModeEnabledWrapper::False,
            None => UpgradeModeEnabledWrapper::Unknown,
        }
    }
}

impl From<bool> for UpgradeModeEnabledWrapper {
    fn from(value: bool) -> Self {
        Some(value).into()
    }
}

impl Display for UpgradeModeEnabledWrapper {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            UpgradeModeEnabledWrapper::True => write!(f, "true"),
            UpgradeModeEnabledWrapper::False => write!(f, "false"),
            UpgradeModeEnabledWrapper::Unknown => write!(f, "unknown"),
        }
    }
}

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,

    /// defines the timestamp when the bandwidth value was last updated
    last_updated_ts: AtomicI64,
}

impl ClientBandwidth {
    pub(crate) fn new_empty() -> Self {
        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)
    }

    pub(crate) fn maybe_log_bandwidth(
        &self,
        now: Option<OffsetDateTime>,
        upgrade_mode: impl Into<UpgradeModeEnabledWrapper>,
    ) {
        let last = self.last_logged();
        let now = now.unwrap_or_else(OffsetDateTime::now_utc);
        if last + Duration::from_secs(10) < now {
            self.log_bandwidth(Some(now), upgrade_mode)
        }
    }

    pub(crate) fn log_bandwidth(
        &self,
        now: Option<OffsetDateTime>,
        upgrade_mode: impl Into<UpgradeModeEnabledWrapper>,
    ) {
        let now = now.unwrap_or_else(OffsetDateTime::now_utc);
        let upgrade_mode = upgrade_mode.into();

        let remaining = self.remaining();
        let remaining_bi2 = bibytes2(remaining as f64);

        if remaining < 0 {
            tracing::warn!(
                "OUT OF BANDWIDTH. remaining: {remaining_bi2}. in 'upgrade mode': {upgrade_mode}"
            );
        } else if remaining < 1_000_000 {
            tracing::info!(
                "remaining bandwidth: {remaining_bi2}. in 'upgrade mode': {upgrade_mode}"
            );
        } else {
            tracing::trace!(
                "remaining bandwidth: {remaining_bi2}. in 'upgrade mode': {upgrade_mode}"
            );
        }

        self.inner
            .last_logged_ts
            .store(now.unix_timestamp(), Ordering::Relaxed)
    }

    pub(crate) fn update_and_maybe_log(
        &self,
        remaining: i64,
        upgrade_mode: impl Into<UpgradeModeEnabledWrapper>,
    ) {
        let now = OffsetDateTime::now_utc();
        self.inner.available.store(remaining, Ordering::Release);
        self.inner
            .last_updated_ts
            .store(now.unix_timestamp(), Ordering::Relaxed);
        self.maybe_log_bandwidth(Some(now), upgrade_mode)
    }

    pub(crate) fn update_and_log(
        &self,
        remaining: i64,
        upgrade_mode: impl Into<UpgradeModeEnabledWrapper>,
    ) {
        let now = OffsetDateTime::now_utc();
        self.inner.available.store(remaining, Ordering::Release);
        self.inner
            .last_updated_ts
            .store(now.unix_timestamp(), Ordering::Relaxed);
        self.log_bandwidth(Some(now), upgrade_mode)
    }

    fn last_logged(&self) -> OffsetDateTime {
        // SAFETY: this value is always populated with valid timestamps
        #[allow(clippy::unwrap_used)]
        OffsetDateTime::from_unix_timestamp(self.inner.last_logged_ts.load(Ordering::Relaxed))
            .unwrap()
    }
}