bbb_api_wrapper/
throttler.rs1use 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
15pub 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
56pub 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
73fn 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
86pub fn unix_timestamp() -> u64 {
88 SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis().try_into().unwrap()
89}