1 // Copyright 2024, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 use std::collections::VecDeque;
16 use std::time::Duration;
17 use std::time::Instant;
18 
19 // 24 hours.
20 const HISTORY_EXPIRY: Duration = Duration::from_secs(60 * 60 * 24);
21 
22 /// Stores the log of zram writeback size to calculate daily limit.
23 pub struct ZramWritebackHistory {
24     history: VecDeque<(u64, Instant)>,
25 }
26 
27 impl ZramWritebackHistory {
28     /// Creates a new [ZramWritebackHistory].
new() -> Self29     pub fn new() -> Self {
30         Self { history: VecDeque::new() }
31     }
32 
33     /// Records a new log of zram writeback.
record(&mut self, pages: u64, now: Instant)34     pub fn record(&mut self, pages: u64, now: Instant) {
35         self.history.push_back((pages, now));
36     }
37 
38     /// Evicts expired records.
cleanup(&mut self, now: Instant)39     pub fn cleanup(&mut self, now: Instant) {
40         while !self.history.is_empty() && now - self.history.front().unwrap().1 > HISTORY_EXPIRY {
41             self.history.pop_front();
42         }
43     }
44 
45     /// Calculates the daily limit of zram writeback left.
calculate_daily_limit(&self, max_pages_per_day: u64, now: Instant) -> u6446     pub fn calculate_daily_limit(&self, max_pages_per_day: u64, now: Instant) -> u64 {
47         let pages_written = self
48             .history
49             .iter()
50             .filter(|(_, t)| now.saturating_duration_since(*t) < HISTORY_EXPIRY)
51             .map(|(p, _)| p)
52             .sum::<u64>();
53         if pages_written >= max_pages_per_day {
54             return 0;
55         }
56         max_pages_per_day - pages_written
57     }
58 }
59 
60 #[cfg(test)]
61 mod tests {
62     use super::*;
63 
64     #[test]
test_calculate_daily_limit()65     fn test_calculate_daily_limit() {
66         let mut history = ZramWritebackHistory::new();
67         let base_time = Instant::now();
68 
69         // records 1 day before is ignored.
70         history.record(1, base_time);
71         history.record(1, base_time);
72         history.record(2, base_time + Duration::from_secs(1));
73         history.record(3, base_time + HISTORY_EXPIRY);
74         assert_eq!(history.calculate_daily_limit(100, base_time + HISTORY_EXPIRY), 95);
75     }
76 
77     #[test]
test_calculate_daily_limit_empty()78     fn test_calculate_daily_limit_empty() {
79         let history = ZramWritebackHistory::new();
80         assert_eq!(history.calculate_daily_limit(100, Instant::now()), 100);
81     }
82 
83     #[test]
test_calculate_daily_limit_exceeds_max()84     fn test_calculate_daily_limit_exceeds_max() {
85         let mut history = ZramWritebackHistory::new();
86         let base_time = Instant::now();
87         // records 1 day before is ignored.
88         history.record(1, base_time);
89         history.record(2, base_time + Duration::from_secs(1));
90         history.record(3, base_time + HISTORY_EXPIRY);
91 
92         assert_eq!(history.calculate_daily_limit(1, base_time + HISTORY_EXPIRY), 0);
93         assert_eq!(history.calculate_daily_limit(2, base_time + HISTORY_EXPIRY), 0);
94         assert_eq!(history.calculate_daily_limit(3, base_time + HISTORY_EXPIRY), 0);
95         assert_eq!(history.calculate_daily_limit(4, base_time + HISTORY_EXPIRY), 0);
96         assert_eq!(history.calculate_daily_limit(5, base_time + HISTORY_EXPIRY), 0);
97         assert_eq!(history.calculate_daily_limit(6, base_time + HISTORY_EXPIRY), 1);
98     }
99 
100     #[test]
test_calculate_daily_limit_after_cleanup()101     fn test_calculate_daily_limit_after_cleanup() {
102         let mut history = ZramWritebackHistory::new();
103         let base_time = Instant::now();
104         // records 1 day before will be cleaned up.
105         history.record(1, base_time);
106         history.record(1, base_time);
107         history.record(2, base_time + Duration::from_secs(1));
108         history.record(3, base_time + HISTORY_EXPIRY);
109 
110         history.cleanup(base_time + HISTORY_EXPIRY);
111 
112         // The same result as test_calculate_daily_limit
113         assert_eq!(history.calculate_daily_limit(100, base_time + HISTORY_EXPIRY), 95);
114     }
115 }
116