skp_ratelimit/storage/
entry.rs

1//! Storage entry type for rate limiting state.
2
3use serde::{Deserialize, Serialize};
4
5/// Entry stored in the storage backend.
6///
7/// This struct contains all state needed by rate limiting algorithms,
8/// designed to be flexible enough for any algorithm type.
9#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
10pub struct StorageEntry {
11    /// Request count (for window-based algorithms).
12    pub count: u64,
13
14    /// Window start timestamp (Unix milliseconds).
15    pub window_start: u64,
16
17    /// Theoretical Arrival Time for GCRA (Unix milliseconds).
18    pub tat: Option<u64>,
19
20    /// Available tokens (for token bucket algorithm).
21    pub tokens: Option<f64>,
22
23    /// Last update timestamp (Unix milliseconds).
24    pub last_update: u64,
25
26    /// Previous window count (for sliding window).
27    pub prev_count: Option<u64>,
28
29    /// Request timestamps (for sliding log algorithm).
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub timestamps: Option<Vec<u64>>,
32
33    /// Optional metadata (algorithm-specific).
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub metadata: Option<Vec<u8>>,
36}
37
38impl StorageEntry {
39    /// Create a new storage entry for window-based algorithms.
40    pub fn new(count: u64, window_start: u64) -> Self {
41        Self {
42            count,
43            window_start,
44            tat: None,
45            tokens: None,
46            last_update: window_start,
47            prev_count: None,
48            timestamps: None,
49            metadata: None,
50        }
51    }
52
53    /// Create a storage entry for GCRA algorithm.
54    pub fn with_tat(tat: u64) -> Self {
55        Self {
56            count: 0,
57            window_start: tat,
58            tat: Some(tat),
59            tokens: None,
60            last_update: tat,
61            prev_count: None,
62            timestamps: None,
63            metadata: None,
64        }
65    }
66
67    /// Create a storage entry for token bucket.
68    pub fn with_tokens(tokens: f64, last_update: u64) -> Self {
69        Self {
70            count: 0,
71            window_start: last_update,
72            tat: None,
73            tokens: Some(tokens),
74            last_update,
75            prev_count: None,
76            timestamps: None,
77            metadata: None,
78        }
79    }
80
81    /// Create a storage entry for sliding log.
82    pub fn with_timestamps(timestamps: Vec<u64>) -> Self {
83        let now = timestamps.last().copied().unwrap_or(0);
84        Self {
85            count: timestamps.len() as u64,
86            window_start: now,
87            tat: None,
88            tokens: None,
89            last_update: now,
90            prev_count: None,
91            timestamps: Some(timestamps),
92            metadata: None,
93        }
94    }
95
96    /// Set the TAT value.
97    pub fn set_tat(mut self, tat: u64) -> Self {
98        self.tat = Some(tat);
99        self
100    }
101
102    /// Set the token count.
103    pub fn set_tokens(mut self, tokens: f64) -> Self {
104        self.tokens = Some(tokens);
105        self
106    }
107
108    /// Set the last update timestamp.
109    pub fn set_last_update(mut self, last_update: u64) -> Self {
110        self.last_update = last_update;
111        self
112    }
113
114    /// Set previous window count.
115    pub fn set_prev_count(mut self, count: u64) -> Self {
116        self.prev_count = Some(count);
117        self
118    }
119
120    /// Set metadata.
121    pub fn set_metadata(mut self, metadata: Vec<u8>) -> Self {
122        self.metadata = Some(metadata);
123        self
124    }
125
126    /// Get tokens, defaulting to 0.0 if not set.
127    pub fn tokens_or_default(&self) -> f64 {
128        self.tokens.unwrap_or(0.0)
129    }
130
131    /// Get TAT, defaulting to 0 if not set.
132    pub fn tat_or_default(&self) -> u64 {
133        self.tat.unwrap_or(0)
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn test_entry_new() {
143        let entry = StorageEntry::new(5, 1000);
144        assert_eq!(entry.count, 5);
145        assert_eq!(entry.window_start, 1000);
146        assert!(entry.tat.is_none());
147        assert!(entry.tokens.is_none());
148    }
149
150    #[test]
151    fn test_entry_with_tat() {
152        let entry = StorageEntry::with_tat(5000);
153        assert_eq!(entry.tat, Some(5000));
154        assert_eq!(entry.tat_or_default(), 5000);
155    }
156
157    #[test]
158    fn test_entry_with_tokens() {
159        let entry = StorageEntry::with_tokens(10.5, 2000);
160        assert_eq!(entry.tokens, Some(10.5));
161        assert_eq!(entry.tokens_or_default(), 10.5);
162        assert_eq!(entry.last_update, 2000);
163    }
164
165    #[test]
166    fn test_entry_with_timestamps() {
167        let timestamps = vec![1000, 2000, 3000];
168        let entry = StorageEntry::with_timestamps(timestamps.clone());
169        assert_eq!(entry.timestamps, Some(timestamps));
170        assert_eq!(entry.count, 3);
171    }
172
173    #[test]
174    fn test_entry_serialization() {
175        let entry = StorageEntry::new(10, 1000).set_tokens(5.5).set_tat(2000);
176        let json = serde_json::to_string(&entry).unwrap();
177        let deserialized: StorageEntry = serde_json::from_str(&json).unwrap();
178        assert_eq!(entry, deserialized);
179    }
180}
skp_ratelimit/storage/entry.rs

skp_ratelimit/storage/
entry.rs
