Refactor: simplify and allow scaling for rate limit definitions

zejiran · Oct 12, 2023 · 8c439b0 · 8c439b0
1 parent e9691ff
commit 8c439b0
Showing 1 changed file with 21 additions and 21 deletions.
diff --git a/rate_limited_notification_service/src/main.rs b/rate_limited_notification_service/src/main.rs
@@ -23,14 +23,6 @@ impl NotificationService {
         }
     }
 
-    fn create_rate_limit(max_requests: u32, per_duration: Duration) -> RateLimit {
-        RateLimit {
-            max_requests,
-            per_duration,
-            recipient_counters: HashMap::new(),
-        }
-    }
-
     fn send(
         &mut self,
         notification_type: &str,
@@ -40,7 +32,11 @@ impl NotificationService {
         let rate_limit = self
             .rate_limits
             .entry(notification_type.to_string())
-            .or_insert(Self::create_rate_limit(u32::MAX, Duration::from_secs(1)));
+            .or_insert(RateLimit {
+                max_requests: u32::MAX,
+                per_duration: Duration::from_secs(1),
+                recipient_counters: HashMap::new(),
+            });
 
         let recipient_counter = rate_limit
             .recipient_counters
@@ -80,18 +76,22 @@ fn main() {
     let mut service = NotificationService::new();
 
     // Define rate limits for different notification types
-    service.rate_limits.insert(
-        "status".to_string(),
-        NotificationService::create_rate_limit(2, Duration::from_secs(60)), // 2 per minute
-    );
-    service.rate_limits.insert(
-        "news".to_string(),
-        NotificationService::create_rate_limit(1, Duration::from_secs(24 * 60 * 60)), // 1 per day
-    );
-    service.rate_limits.insert(
-        "marketing".to_string(),
-        NotificationService::create_rate_limit(3, Duration::from_secs(60 * 60)), // 3 per hour
-    );
+    let rate_limits = vec![
+        ("status", 2, Duration::from_secs(60)),         // 2 per minute
+        ("news", 1, Duration::from_secs(24 * 60 * 60)), // 1 per day
+        ("marketing", 3, Duration::from_secs(60 * 60)), // 3 per hour
+    ];
+
+    for (notification_type, max_requests, per_duration) in rate_limits {
+        service.rate_limits.insert(
+            notification_type.to_string(),
+            RateLimit {
+                max_requests,
+                per_duration,
+                recipient_counters: HashMap::new(),
+            },
+        );
+    }
 
     // Example usage:
     for _ in 0..5 {