api_key_pool/
lib.rs

1//! Pool of API keys to circumvent rate-limits.
2//!
3//! This package provides an easy way to use multiple API keys to bypass draconian rate-limit policies.
4//!
5//! # Example
6//!
7//! ```
8//! use chrono::Duration;
9//! use tokio::time;
10//!
11//! use api_key_pool::*;
12//!
13//! #[tokio::main]
14//! async fn main() {
15//!     // Create a RateLimitPolicy to be applied to all API keys.
16//!     // Note: An APIPool can have APIKeys with different RateLimitPolicies.
17//!     //       For the sake of simplicity, this example assumes identical policies.
18//!     let pol = RateLimitPolicy::new(1, Duration::seconds(2));
19//!
20//!     // Create the APIKeys.
21//!     let api1 = APIKey::new("1", pol);
22//!     let api2 = APIKey::new("2", pol);
23//!     let api3 = APIKey::new("3", pol);
24//!
25//!     // Create the APIKeyPool.
26//!     let mut pool = APIKeyPool::new();
27//!     pool.add_key(api1).await;
28//!     pool.add_key(api2).await;
29//!     pool.add_key(api3).await;
30//!
31//!     // Simulate 20 requests.
32//!     let mut ctr = 0;
33//!     while ctr < 20 {
34//!         // Use the APIKey if available (according to its respective RateLimitPolicy) or sleep.
35//!         if let Some(key) = pool.poll_for_key().await {
36//!             println!("{}", key);
37//!             ctr += 1;
38//!         } else {
39//!             println!("Have to sleep.");
40//!             time::sleep(time::Duration::from_millis(500)).await;
41//!         }
42//!     }
43//! }
44//! ```
45
46
47use std::cmp::Reverse;
48use std::collections::BinaryHeap;
49use std::sync::Arc;
50
51use chrono::{DateTime, Utc};
52use tokio::sync::Mutex;
53
54/// A pool of API keys.
55#[derive(Default)]
56pub struct APIKeyPool {
57    /// Collection holding the API keys.
58    api_keys: Arc<Mutex<Vec<APIKey>>>,
59}
60
61impl APIKeyPool {
62    /// Returns an empty API key pool.
63    pub fn new() -> Self {
64        Self {
65            api_keys: Arc::new(Mutex::new(Vec::new())),
66        }
67    }
68
69    /// Adds an API key to an API key pool.
70    ///
71    /// # Arguments
72    ///
73    /// * `key` - the API key to be added.
74    pub async fn add_key(&mut self, key: APIKey) {
75        self.api_keys.lock().await.push(key);
76    }
77
78    /// Checks the API key pool for any available API keys, and returns the API key if available.
79    pub async fn poll_for_key(&mut self) -> Option<String> {
80        // TODO: Performance can be improved by keeping track of index of last used key.
81        for key in &mut self.api_keys.lock().await.iter_mut() {
82            if key.is_ready().await {
83                return Some(key.use_key().await);
84            }
85        }
86        None
87    }
88}
89
90/// An API key, with its associated RateLimitPolicy
91pub struct APIKey {
92    /// The API key code.
93    key: String,
94    /// The rate limit policy that governs this API key.
95    policy: RateLimitPolicy,
96    /// Min-heap used to calculate if the key is available.
97    times: Arc<Mutex<BinaryHeap<Reverse<DateTime<Utc>>>>>,
98}
99
100impl APIKey {
101    /// Returns an API key with the given policy and code.
102    ///
103    /// # Arguments
104    ///
105    /// * `key` - the API key code.
106    /// * `policy` - the rate limit policy governing the API key.
107    pub fn new(key: &str, policy: RateLimitPolicy) -> Self {
108        let mut _times = BinaryHeap::new();
109        _times.reserve(policy.count);
110        let times = Arc::new(Mutex::new(_times));
111        Self {
112            key: String::from(key),
113            policy,
114            times,
115        }
116    }
117
118    /// Returns the code of an API key.
119    fn get_key(&self) -> String {
120        self.key.clone()
121    }
122
123    /// Checks to see if the API key is available for use.
124    async fn is_ready(&self) -> bool {
125        // If we have used the API key less than N times, we can use it again.
126        if self.times.lock().await.len() < self.policy.count {
127            return true;
128        }
129        if let Some(oldest) = self.times.lock().await.peek() {
130            // If the oldest time used is at least D duration ago.
131            if oldest.0 < Utc::now() - self.policy.per {
132                return true;
133            }
134        }
135        false
136    }
137
138    /// Uses the key.
139    async fn use_key(&mut self) -> String {
140        if self.times.lock().await.len() >= self.policy.count {
141            self.times.lock().await.pop();
142        }
143        self.times.lock().await.push(Reverse(Utc::now()));
144        self.get_key().clone()
145    }
146}
147
148/// A policy for rate-limiting an API key.
149#[derive(Clone, Copy)]
150pub struct RateLimitPolicy {
151    /// The number of times an API key can be used in the specified duration.
152    pub count: usize,
153    /// The duration.
154    pub per: chrono::Duration,
155}
156
157impl RateLimitPolicy {
158    /// Returns a rate-limit policy with the parameters.
159    ///
160    /// # Arguments
161    ///
162    /// * `count` - N times
163    /// * `per` - per D duration
164    pub fn new(count: usize, per: chrono::Duration) -> Self {
165        Self { count, per }
166    }
167}
api_key_pool/lib.rs

api_key_pool/
lib.rs
