yf_common/
rate_limit.rs

1//! Rate limiting for Yahoo Finance API requests
2
3use governor::{Quota, RateLimiter};
4use std::num::NonZeroU32;
5use std::sync::Arc;
6use tracing::debug;
7
8/// Type alias for the rate limiter
9pub type YfRateLimiter = RateLimiter<
10    governor::state::NotKeyed,
11    governor::state::InMemoryState,
12    governor::clock::DefaultClock,
13>;
14
15/// Default requests per minute
16pub const DEFAULT_REQUESTS_PER_MINUTE: u32 = 5;
17
18/// Configuration for rate limiting
19#[derive(Debug, Clone)]
20pub struct RateLimitConfig {
21    /// Requests allowed per minute
22    pub requests_per_minute: u32,
23}
24
25impl Default for RateLimitConfig {
26    fn default() -> Self {
27        Self {
28            requests_per_minute: DEFAULT_REQUESTS_PER_MINUTE,
29        }
30    }
31}
32
33impl RateLimitConfig {
34    /// Create a new rate limit configuration
35    pub fn new(requests_per_minute: u32) -> Self {
36        Self { requests_per_minute }
37    }
38
39    /// Build a rate limiter from this configuration
40    pub fn build_limiter(&self) -> Arc<YfRateLimiter> {
41        let quota = Quota::per_minute(
42            NonZeroU32::new(self.requests_per_minute).unwrap_or(NonZeroU32::new(1).unwrap())
43        );
44        Arc::new(RateLimiter::direct(quota))
45    }
46}
47
48/// Wait for a rate limit permit
49pub async fn wait_for_permit(limiter: &YfRateLimiter) {
50    limiter.until_ready().await;
51    debug!("Rate limit permit acquired");
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[test]
59    fn test_default_config() {
60        let config = RateLimitConfig::default();
61        assert_eq!(config.requests_per_minute, 5);
62    }
63
64    #[test]
65    fn test_custom_config() {
66        let config = RateLimitConfig::new(10);
67        assert_eq!(config.requests_per_minute, 10);
68    }
69
70    #[test]
71    fn test_build_limiter() {
72        let config = RateLimitConfig::new(5);
73        let limiter = config.build_limiter();
74        assert!(limiter.check().is_ok());
75    }
76
77    #[test]
78    fn test_rate_limiter_zero_fallback() {
79        // Zero should not panic, uses NonZeroU32 minimum
80        let config = RateLimitConfig::new(0);
81        let limiter = config.build_limiter();
82        // Should still work (defaults to 1)
83        assert!(limiter.check().is_ok());
84    }
85}
yf_common/rate_limit.rs

yf_common/
rate_limit.rs
