bbb_api_wrapper/
throttler.rs

1// Copyright (c) 2021 Harry [Majored] [hello@majored.pw]
2// MIT License (https://github.com/Majored/rs-bbb-api-wrapper/blob/main/LICENSE)
3
4//! Holds key types for tracking our compliance with the API's rate limits.
5
6use std::convert::TryInto;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::{SystemTime, UNIX_EPOCH};
9
10pub enum RequestType {
11    READ,
12    WRITE,
13}
14
15/// A strucutre for storing the relevant atomic values in order to track our compliance with the API's rate limits.
16pub struct RateLimitStore {
17    pub read_last_retry: AtomicU64,
18    pub read_last_request: AtomicU64,
19
20    pub write_last_retry: AtomicU64,
21    pub write_last_request: AtomicU64,
22}
23
24impl RateLimitStore {
25    pub fn new() -> Self {
26        RateLimitStore {
27            read_last_retry: AtomicU64::new(0),
28            read_last_request: AtomicU64::new(unix_timestamp()),
29
30            write_last_retry: AtomicU64::new(0),
31            write_last_request: AtomicU64::new(unix_timestamp()),
32        }
33    }
34
35    pub fn store_read(&self, retry: u64) {
36        self.read_last_retry.store(retry, Ordering::Release);
37        self.read_last_request.store(unix_timestamp(), Ordering::Release);
38    }
39
40    pub fn store_write(&self, retry: u64) {
41        self.write_last_retry.store(retry, Ordering::Release);
42        self.write_last_request.store(unix_timestamp(), Ordering::Release);
43    }
44
45    pub fn reset_read(&self) {
46        self.read_last_retry.store(0, Ordering::Release);
47        self.read_last_request.store(unix_timestamp(), Ordering::Release);
48    }
49
50    pub fn reset_write(&self) {
51        self.write_last_retry.store(0, Ordering::Release);
52        self.write_last_request.store(unix_timestamp(), Ordering::Release);
53    }
54}
55
56/// Compute how long, if at all, we should stall the next request in order to be compliant with rate limiting.
57///
58/// Returned value is in milliseconds. A value of 0 indiciates that there's no need to stall the calling request.
59pub fn stall_for(store: &RateLimitStore, request_type: RequestType) -> u64 {
60    let time = unix_timestamp();
61    let mut stall_for = 0;
62
63    if let RequestType::READ = request_type {
64        stall_for = stall_for_helper(&store.read_last_retry, &store.read_last_request, time);
65    }
66    if let RequestType::WRITE = request_type {
67        stall_for = stall_for_helper(&store.write_last_retry, &store.write_last_request, time);
68    }
69
70    stall_for
71}
72
73/// A helper function for `stall_for` which computes over a generic set of rate limiting parameters.
74fn stall_for_helper(a_last_retry: &AtomicU64, a_last_request: &AtomicU64, time: u64) -> u64 {
75    let mut stall_for = 0;
76    let last_retry = a_last_retry.load(Ordering::Acquire);
77    let last_request = a_last_request.load(Ordering::Acquire);
78
79    if last_retry > 0 && (time - last_request) < last_retry {
80        stall_for = last_retry - (time - last_request);
81    }
82
83    stall_for
84}
85
86/// Return the current time as a UNIX millisecond timestamp.
87pub fn unix_timestamp() -> u64 {
88    SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis().try_into().unwrap()
89}