junobuild_shared/rate/
utils.rs

1use crate::rate::types::{RateConfig, RateTokenStore, RateTokens};
2use ic_cdk::api::time;
3use std::cmp::min;
4
5pub fn increment_and_assert_rate_store(
6    key: &String,
7    config: &Option<RateConfig>,
8    rate_tokens: &mut RateTokenStore,
9) -> Result<(), String> {
10    let config = match config {
11        Some(config) => config,
12        None => return Ok(()),
13    };
14
15    if let Some(tokens) = rate_tokens.get_mut(key) {
16        increment_and_assert_rate(config, tokens)?;
17    } else {
18        rate_tokens.insert(key.clone(), RateTokens::default());
19    }
20
21    Ok(())
22}
23
24pub fn increment_and_assert_rate(
25    config: &RateConfig,
26    tokens: &mut RateTokens,
27) -> Result<(), String> {
28    let new_tokens = (time() - tokens.updated_at) / config.time_per_token_ns;
29    if new_tokens > 0 {
30        // The number of tokens is capped otherwise tokens might accumulate
31        tokens.tokens = min(config.max_tokens, tokens.tokens + new_tokens);
32        tokens.updated_at += config.time_per_token_ns * new_tokens;
33    }
34
35    // deduct a token for the current call
36    if tokens.tokens > 0 {
37        tokens.tokens -= 1;
38        Ok(())
39    } else {
40        Err("Rate limit reached, try again later.".to_string())
41    }
42}
junobuild_shared/rate/utils.rs

junobuild_shared/rate/
utils.rs
