chie_shared/types/
quota.rs

1//! Quota management types for CHIE Protocol.
2//!
3//! This module provides types for tracking and managing user quotas,
4//! including storage, bandwidth, and rate limits.
5
6#[cfg(feature = "schema")]
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10use super::core::Bytes;
11
12/// Storage quota information for a user or node.
13///
14/// # Examples
15///
16/// ```
17/// use chie_shared::StorageQuota;
18///
19/// // Create 10 GB quota with 5 GB used and 1 GB reserved
20/// let quota = StorageQuota::new(
21///     10 * 1024 * 1024 * 1024,  // 10 GB total
22///     5 * 1024 * 1024 * 1024,   // 5 GB used
23///     1 * 1024 * 1024 * 1024,   // 1 GB reserved
24/// );
25///
26/// // Check available space
27/// assert_eq!(quota.available_bytes(), 4 * 1024 * 1024 * 1024);
28///
29/// // Check utilization
30/// assert_eq!(quota.utilization(), 0.5);
31///
32/// // Check if can allocate more
33/// assert!(quota.can_allocate(3 * 1024 * 1024 * 1024));
34/// assert!(!quota.is_nearly_full());
35/// ```
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
37#[cfg_attr(feature = "schema", derive(JsonSchema))]
38pub struct StorageQuota {
39    /// Total allocated storage in bytes.
40    pub total_bytes: Bytes,
41    /// Currently used storage in bytes.
42    pub used_bytes: Bytes,
43    /// Reserved storage in bytes (pending operations).
44    pub reserved_bytes: Bytes,
45}
46
47impl StorageQuota {
48    /// Create a new storage quota.
49    pub fn new(total_bytes: Bytes, used_bytes: Bytes, reserved_bytes: Bytes) -> Self {
50        Self {
51            total_bytes,
52            used_bytes,
53            reserved_bytes,
54        }
55    }
56
57    /// Get available storage in bytes.
58    pub fn available_bytes(&self) -> Bytes {
59        self.total_bytes
60            .saturating_sub(self.used_bytes)
61            .saturating_sub(self.reserved_bytes)
62    }
63
64    /// Get utilization percentage (0.0 to 1.0).
65    pub fn utilization(&self) -> f64 {
66        if self.total_bytes == 0 {
67            0.0
68        } else {
69            self.used_bytes as f64 / self.total_bytes as f64
70        }
71    }
72
73    /// Check if there's enough space for additional bytes.
74    pub fn can_allocate(&self, bytes: Bytes) -> bool {
75        self.available_bytes() >= bytes
76    }
77
78    /// Check if quota is nearly full (>90% used).
79    pub fn is_nearly_full(&self) -> bool {
80        self.utilization() > 0.9
81    }
82
83    /// Get total allocated bytes (used + reserved).
84    pub fn allocated_bytes(&self) -> Bytes {
85        self.used_bytes.saturating_add(self.reserved_bytes)
86    }
87}
88
89impl Default for StorageQuota {
90    fn default() -> Self {
91        Self::new(0, 0, 0)
92    }
93}
94
95/// Bandwidth quota for a time period.
96///
97/// # Examples
98///
99/// ```
100/// use chie_shared::BandwidthQuota;
101///
102/// let now = chrono::Utc::now().timestamp_millis() as u64;
103///
104/// // Create 100 GB monthly bandwidth quota
105/// let quota = BandwidthQuota::new(
106///     100 * 1024 * 1024 * 1024,  // 100 GB total
107///     50 * 1024 * 1024 * 1024,   // 50 GB used
108///     30 * 24 * 60 * 60,         // 30 days in seconds
109///     now,
110/// );
111///
112/// // Check remaining bandwidth
113/// assert_eq!(quota.remaining_bytes(), 50 * 1024 * 1024 * 1024);
114///
115/// // Check utilization
116/// assert_eq!(quota.utilization(), 0.5);
117///
118/// // Check if can consume more
119/// assert!(quota.can_consume(30 * 1024 * 1024 * 1024));
120/// assert!(!quota.is_exceeded());
121/// ```
122#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
123#[cfg_attr(feature = "schema", derive(JsonSchema))]
124pub struct BandwidthQuota {
125    /// Total bandwidth quota in bytes.
126    pub total_bytes: Bytes,
127    /// Bandwidth used in current period (bytes).
128    pub used_bytes: Bytes,
129    /// Period duration in seconds.
130    pub period_seconds: u64,
131    /// Period start timestamp (Unix milliseconds).
132    pub period_start_ms: u64,
133}
134
135impl BandwidthQuota {
136    /// Create a new bandwidth quota.
137    pub fn new(
138        total_bytes: Bytes,
139        used_bytes: Bytes,
140        period_seconds: u64,
141        period_start_ms: u64,
142    ) -> Self {
143        Self {
144            total_bytes,
145            used_bytes,
146            period_seconds,
147            period_start_ms,
148        }
149    }
150
151    /// Get remaining bandwidth in bytes.
152    pub fn remaining_bytes(&self) -> Bytes {
153        self.total_bytes.saturating_sub(self.used_bytes)
154    }
155
156    /// Get utilization percentage (0.0 to 1.0).
157    pub fn utilization(&self) -> f64 {
158        if self.total_bytes == 0 {
159            0.0
160        } else {
161            self.used_bytes as f64 / self.total_bytes as f64
162        }
163    }
164
165    /// Check if there's enough bandwidth remaining.
166    pub fn can_consume(&self, bytes: Bytes) -> bool {
167        self.remaining_bytes() >= bytes
168    }
169
170    /// Check if quota is exceeded.
171    pub fn is_exceeded(&self) -> bool {
172        self.used_bytes >= self.total_bytes
173    }
174
175    /// Get average bytes per second used.
176    pub fn avg_bytes_per_second(&self) -> f64 {
177        if self.period_seconds == 0 {
178            0.0
179        } else {
180            self.used_bytes as f64 / self.period_seconds as f64
181        }
182    }
183
184    /// Check if period has expired (given current time).
185    pub fn is_period_expired(&self, current_time_ms: u64) -> bool {
186        let elapsed_ms = current_time_ms.saturating_sub(self.period_start_ms);
187        let period_ms = self.period_seconds * 1000;
188        elapsed_ms >= period_ms
189    }
190}
191
192impl Default for BandwidthQuota {
193    fn default() -> Self {
194        Self::new(0, 0, 0, 0)
195    }
196}
197
198/// Rate limit quota for API requests.
199#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
200#[cfg_attr(feature = "schema", derive(JsonSchema))]
201pub struct RateLimitQuota {
202    /// Maximum requests allowed in period.
203    pub max_requests: u32,
204    /// Current request count in period.
205    pub current_requests: u32,
206    /// Period duration in seconds.
207    pub period_seconds: u64,
208    /// Period start timestamp (Unix milliseconds).
209    pub period_start_ms: u64,
210}
211
212impl RateLimitQuota {
213    /// Create a new rate limit quota.
214    pub fn new(
215        max_requests: u32,
216        current_requests: u32,
217        period_seconds: u64,
218        period_start_ms: u64,
219    ) -> Self {
220        Self {
221            max_requests,
222            current_requests,
223            period_seconds,
224            period_start_ms,
225        }
226    }
227
228    /// Get remaining requests allowed.
229    pub fn remaining_requests(&self) -> u32 {
230        self.max_requests.saturating_sub(self.current_requests)
231    }
232
233    /// Check if another request is allowed.
234    pub fn is_allowed(&self) -> bool {
235        self.current_requests < self.max_requests
236    }
237
238    /// Check if rate limit is exceeded.
239    pub fn is_exceeded(&self) -> bool {
240        self.current_requests >= self.max_requests
241    }
242
243    /// Get utilization percentage (0.0 to 1.0).
244    pub fn utilization(&self) -> f64 {
245        if self.max_requests == 0 {
246            0.0
247        } else {
248            self.current_requests as f64 / self.max_requests as f64
249        }
250    }
251
252    /// Check if period has expired (given current time).
253    pub fn is_period_expired(&self, current_time_ms: u64) -> bool {
254        let elapsed_ms = current_time_ms.saturating_sub(self.period_start_ms);
255        let period_ms = self.period_seconds * 1000;
256        elapsed_ms >= period_ms
257    }
258}
259
260impl Default for RateLimitQuota {
261    fn default() -> Self {
262        Self::new(0, 0, 0, 0)
263    }
264}
265
266/// Combined quota information for a user.
267#[derive(Debug, Clone, Serialize, Deserialize)]
268#[cfg_attr(feature = "schema", derive(JsonSchema))]
269pub struct UserQuota {
270    /// Storage quota.
271    pub storage: StorageQuota,
272    /// Bandwidth quota.
273    pub bandwidth: BandwidthQuota,
274    /// API rate limit quota.
275    pub rate_limit: RateLimitQuota,
276}
277
278impl UserQuota {
279    /// Create new user quota.
280    pub fn new(
281        storage: StorageQuota,
282        bandwidth: BandwidthQuota,
283        rate_limit: RateLimitQuota,
284    ) -> Self {
285        Self {
286            storage,
287            bandwidth,
288            rate_limit,
289        }
290    }
291
292    /// Check if user can perform an operation requiring both storage and bandwidth.
293    pub fn can_perform_operation(&self, storage_bytes: Bytes, bandwidth_bytes: Bytes) -> bool {
294        self.storage.can_allocate(storage_bytes) && self.bandwidth.can_consume(bandwidth_bytes)
295    }
296
297    /// Check if any quota is nearly exhausted (>90%).
298    pub fn has_warning_threshold(&self) -> bool {
299        self.storage.is_nearly_full()
300            || self.bandwidth.utilization() > 0.9
301            || self.rate_limit.utilization() > 0.9
302    }
303}
304
305impl Default for UserQuota {
306    fn default() -> Self {
307        Self::new(
308            StorageQuota::default(),
309            BandwidthQuota::default(),
310            RateLimitQuota::default(),
311        )
312    }
313}
314
315/// Builder for StorageQuota with fluent API.
316#[derive(Debug, Default)]
317pub struct StorageQuotaBuilder {
318    total_bytes: Option<Bytes>,
319    used_bytes: Option<Bytes>,
320    reserved_bytes: Option<Bytes>,
321}
322
323impl StorageQuotaBuilder {
324    /// Create a new builder.
325    pub fn new() -> Self {
326        Self::default()
327    }
328
329    /// Set total bytes.
330    pub fn total_bytes(mut self, bytes: Bytes) -> Self {
331        self.total_bytes = Some(bytes);
332        self
333    }
334
335    /// Set used bytes.
336    pub fn used_bytes(mut self, bytes: Bytes) -> Self {
337        self.used_bytes = Some(bytes);
338        self
339    }
340
341    /// Set reserved bytes.
342    pub fn reserved_bytes(mut self, bytes: Bytes) -> Self {
343        self.reserved_bytes = Some(bytes);
344        self
345    }
346
347    /// Build the StorageQuota.
348    pub fn build(self) -> StorageQuota {
349        StorageQuota::new(
350            self.total_bytes.unwrap_or(0),
351            self.used_bytes.unwrap_or(0),
352            self.reserved_bytes.unwrap_or(0),
353        )
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    #[test]
362    fn test_storage_quota() {
363        let quota = StorageQuota::new(1000, 400, 100);
364        assert_eq!(quota.available_bytes(), 500);
365        assert_eq!(quota.utilization(), 0.4);
366        assert!(quota.can_allocate(400));
367        assert!(!quota.can_allocate(600));
368        assert!(!quota.is_nearly_full());
369        assert_eq!(quota.allocated_bytes(), 500);
370    }
371
372    #[test]
373    fn test_storage_quota_nearly_full() {
374        let quota = StorageQuota::new(1000, 950, 0);
375        assert!(quota.is_nearly_full());
376    }
377
378    #[test]
379    fn test_storage_quota_saturating() {
380        let quota = StorageQuota::new(100, 80, 50);
381        assert_eq!(quota.available_bytes(), 0); // Saturating sub
382    }
383
384    #[test]
385    fn test_bandwidth_quota() {
386        let quota = BandwidthQuota::new(1000, 400, 3600, 0);
387        assert_eq!(quota.remaining_bytes(), 600);
388        assert_eq!(quota.utilization(), 0.4);
389        assert!(quota.can_consume(500));
390        assert!(!quota.can_consume(700));
391        assert!(!quota.is_exceeded());
392    }
393
394    #[test]
395    fn test_bandwidth_quota_exceeded() {
396        let quota = BandwidthQuota::new(1000, 1200, 3600, 0);
397        assert!(quota.is_exceeded());
398        assert_eq!(quota.remaining_bytes(), 0);
399    }
400
401    #[test]
402    fn test_bandwidth_quota_period_expired() {
403        let quota = BandwidthQuota::new(1000, 400, 3600, 0);
404        assert!(!quota.is_period_expired(1000 * 1000)); // 1000 seconds
405        assert!(quota.is_period_expired(3600 * 1000 + 1)); // Just past period
406    }
407
408    #[test]
409    fn test_bandwidth_quota_avg_bps() {
410        let quota = BandwidthQuota::new(1000, 500, 10, 0);
411        assert_eq!(quota.avg_bytes_per_second(), 50.0);
412    }
413
414    #[test]
415    fn test_rate_limit_quota() {
416        let quota = RateLimitQuota::new(100, 40, 60, 0);
417        assert_eq!(quota.remaining_requests(), 60);
418        assert!(quota.is_allowed());
419        assert!(!quota.is_exceeded());
420        assert_eq!(quota.utilization(), 0.4);
421    }
422
423    #[test]
424    fn test_rate_limit_quota_exceeded() {
425        let quota = RateLimitQuota::new(100, 120, 60, 0);
426        assert!(!quota.is_allowed());
427        assert!(quota.is_exceeded());
428        assert_eq!(quota.remaining_requests(), 0);
429    }
430
431    #[test]
432    fn test_rate_limit_period_expired() {
433        let quota = RateLimitQuota::new(100, 40, 60, 0);
434        assert!(!quota.is_period_expired(30 * 1000)); // 30 seconds
435        assert!(quota.is_period_expired(60 * 1000 + 1)); // Just past period
436    }
437
438    #[test]
439    fn test_user_quota() {
440        let storage = StorageQuota::new(1000, 400, 0);
441        let bandwidth = BandwidthQuota::new(2000, 800, 3600, 0);
442        let rate_limit = RateLimitQuota::new(100, 40, 60, 0);
443
444        let user_quota = UserQuota::new(storage, bandwidth, rate_limit);
445
446        assert!(user_quota.can_perform_operation(500, 1000));
447        assert!(!user_quota.can_perform_operation(700, 1000)); // Storage exceeded
448        assert!(!user_quota.can_perform_operation(500, 1500)); // Bandwidth exceeded
449        assert!(!user_quota.has_warning_threshold());
450    }
451
452    #[test]
453    fn test_user_quota_warning() {
454        let storage = StorageQuota::new(1000, 950, 0);
455        let bandwidth = BandwidthQuota::new(2000, 800, 3600, 0);
456        let rate_limit = RateLimitQuota::new(100, 40, 60, 0);
457
458        let user_quota = UserQuota::new(storage, bandwidth, rate_limit);
459        assert!(user_quota.has_warning_threshold());
460    }
461
462    #[test]
463    fn test_storage_quota_serialization() {
464        let quota = StorageQuota::new(1000, 400, 100);
465        let json = serde_json::to_string(&quota).unwrap();
466        let deserialized: StorageQuota = serde_json::from_str(&json).unwrap();
467        assert_eq!(quota, deserialized);
468    }
469
470    #[test]
471    fn test_storage_quota_builder() {
472        let quota = StorageQuotaBuilder::new()
473            .total_bytes(1000)
474            .used_bytes(400)
475            .reserved_bytes(100)
476            .build();
477
478        assert_eq!(quota.total_bytes, 1000);
479        assert_eq!(quota.used_bytes, 400);
480        assert_eq!(quota.reserved_bytes, 100);
481    }
482
483    #[test]
484    fn test_storage_quota_builder_partial() {
485        let quota = StorageQuotaBuilder::new().total_bytes(1000).build();
486
487        assert_eq!(quota.total_bytes, 1000);
488        assert_eq!(quota.used_bytes, 0);
489        assert_eq!(quota.reserved_bytes, 0);
490    }
491}
chie_shared/types/quota.rs

chie_shared/types/
quota.rs
