hiqlite 0.14.0

Hiqlite - highly-available, embeddable, raft-based SQLite + cache
Documentation
use crate::config::RateLimitConfig;
use crate::{Client, Error};
use std::cmp::max;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tokio::sync::oneshot;
use tokio::task;

impl Client {
    pub(crate) fn spawn_rate_limit_ticker(
        &self,
        config_cache: Option<RateLimitConfig>,
        config_db: Option<RateLimitConfig>,
        cache_await: crossbeam::channel::Receiver<oneshot::Sender<Result<(), Error>>>,
        db_await: crossbeam::channel::Receiver<oneshot::Sender<Result<(), Error>>>,
    ) {
        task::spawn(Self::run_ticker(
            self.clone(),
            config_cache,
            config_db,
            cache_await,
            db_await,
        ));
    }

    #[allow(unused_variables)]
    async fn run_ticker(
        slf: Client,
        config_cache: Option<RateLimitConfig>,
        config_db: Option<RateLimitConfig>,
        cache_await: crossbeam::channel::Receiver<oneshot::Sender<Result<(), Error>>>,
        db_await: crossbeam::channel::Receiver<oneshot::Sender<Result<(), Error>>>,
    ) {
        let mut interval = tokio::time::interval(Duration::from_secs(1));

        loop {
            interval.tick().await;

            #[cfg(feature = "cache")]
            if let Some(config) = &config_cache {
                let lim = slf.inner.rate_limit_cache.as_ref().unwrap();

                lim.try_update(Ordering::SeqCst, Ordering::SeqCst, |current| {
                    Some(max(current + config.rps, config.burst))
                })
                .ok();

                while let Ok(tx) = cache_await.try_recv() {
                    if tx.send(slf.try_rate_limit_cache()).is_err() {
                        lim.fetch_add(1, Ordering::Relaxed);
                    }
                }
            }

            #[cfg(feature = "sqlite")]
            if let Some(config) = &config_db {
                let lim = slf.inner.rate_limit_db.as_ref().unwrap();

                lim.try_update(Ordering::SeqCst, Ordering::SeqCst, |current| {
                    Some(max(current + config.rps, config.burst))
                })
                .ok();

                while let Ok(tx) = db_await.try_recv() {
                    if tx.send(slf.try_rate_limit_db()).is_err() {
                        lim.fetch_add(1, Ordering::Relaxed);
                    }
                }
            }
        }
    }

    #[cfg(feature = "cache")]
    pub(crate) fn try_rate_limit_cache(&self) -> Result<(), Error> {
        if let Some(v) = &self.inner.rate_limit_cache {
            v.try_update(Ordering::SeqCst, Ordering::SeqCst, |current| {
                if current > 0 { Some(current - 1) } else { None }
            })
            .map_err(|_| Error::RateLimit("Raft rate-limit hit for the Cache layer".into()))?;
        }

        Ok(())
    }

    #[cfg(feature = "cache")]
    pub(crate) async fn rate_limit_cache(&self) -> Result<(), Error> {
        match self.try_rate_limit_cache() {
            Ok(_) => Ok(()),
            Err(err) => {
                let (tx, rx) = oneshot::channel();
                if self.inner.rate_limit_cache_await.try_send(tx).is_ok() {
                    rx.await.unwrap()
                } else {
                    Err(err)
                }
            }
        }
    }

    #[cfg(feature = "sqlite")]
    pub(crate) fn try_rate_limit_db(&self) -> Result<(), Error> {
        if let Some(v) = &self.inner.rate_limit_db {
            v.try_update(Ordering::SeqCst, Ordering::SeqCst, |current| {
                if current > 0 { Some(current - 1) } else { None }
            })
            .map_err(|_| Error::RateLimit("Raft rate-limit hit for the DB layer".into()))?;
        }

        Ok(())
    }

    #[cfg(feature = "sqlite")]
    pub(crate) async fn rate_limit_db(&self) -> Result<(), Error> {
        match self.try_rate_limit_db() {
            Ok(_) => Ok(()),
            Err(err) => {
                let (tx, rx) = oneshot::channel();
                if self.inner.rate_limit_db_await.try_send(tx).is_ok() {
                    rx.await.unwrap()
                } else {
                    Err(err)
                }
            }
        }
    }
}