1use std::collections::BTreeMap;
4use std::time::Duration;
5
6use anyhow::{Context as _, Result, anyhow};
7use async_imap::types::{Quota, QuotaResource};
8
9use crate::chat::add_device_msg_with_importance;
10use crate::config::Config;
11use crate::context::Context;
12use crate::imap::get_watched_folders;
13use crate::imap::session::Session as ImapSession;
14use crate::log::warn;
15use crate::message::Message;
16use crate::tools::{self, time_elapsed};
17use crate::{EventType, stock_str};
18
19pub const QUOTA_WARN_THRESHOLD_PERCENTAGE: u64 = 80;
22
23pub const QUOTA_ERROR_THRESHOLD_PERCENTAGE: u64 = 95;
26
27pub const QUOTA_ALLCLEAR_PERCENTAGE: u64 = 75;
36
37#[derive(Debug)]
39pub struct QuotaInfo {
40 pub(crate) recent: Result<BTreeMap<String, Vec<QuotaResource>>>,
44
45 pub(crate) modified: tools::Time,
47}
48
49async fn get_unique_quota_roots_and_usage(
50 session: &mut ImapSession,
51 folders: Vec<String>,
52) -> Result<BTreeMap<String, Vec<QuotaResource>>> {
53 let mut unique_quota_roots: BTreeMap<String, Vec<QuotaResource>> = BTreeMap::new();
54 for folder in folders {
55 let (quota_roots, quotas) = &session.get_quota_root(&folder).await?;
56 for qr_entries in quota_roots {
58 for quota_root_name in &qr_entries.quota_root_names {
59 let quota: Quota = quotas
61 .iter()
62 .find(|q| &q.root_name == quota_root_name)
63 .cloned()
64 .context("quota_root should have a quota")?;
65 *unique_quota_roots
68 .entry(quota_root_name.clone())
69 .or_default() = quota.resources;
70 }
71 }
72 }
73 Ok(unique_quota_roots)
74}
75
76fn get_highest_usage<'t>(
77 unique_quota_roots: &'t BTreeMap<String, Vec<QuotaResource>>,
78) -> Result<(u64, &'t String, &'t QuotaResource)> {
79 let mut highest: Option<(u64, &'t String, &QuotaResource)> = None;
80 for (name, resources) in unique_quota_roots {
81 for r in resources {
82 let usage_percent = r.get_usage_percentage();
83 match highest {
84 None => {
85 highest = Some((usage_percent, name, r));
86 }
87 Some((up, ..)) => {
88 if up <= usage_percent {
89 highest = Some((usage_percent, name, r));
90 }
91 }
92 };
93 }
94 }
95
96 highest.context("no quota_resource found, this is unexpected")
97}
98
99pub fn needs_quota_warning(curr_percentage: u64, warned_at_percentage: u64) -> bool {
101 (curr_percentage >= QUOTA_WARN_THRESHOLD_PERCENTAGE
102 && warned_at_percentage < QUOTA_WARN_THRESHOLD_PERCENTAGE)
103 || (curr_percentage >= QUOTA_ERROR_THRESHOLD_PERCENTAGE
104 && warned_at_percentage < QUOTA_ERROR_THRESHOLD_PERCENTAGE)
105}
106
107impl Context {
108 pub(crate) async fn quota_needs_update(&self, transport_id: u32, ratelimit_secs: u64) -> bool {
111 let quota = self.quota.read().await;
112 quota.get(&transport_id).is_none_or(|quota| {
113 time_elapsed("a.modified) >= Duration::from_secs(ratelimit_secs)
114 })
115 }
116
117 pub(crate) async fn update_recent_quota(&self, session: &mut ImapSession) -> Result<()> {
126 let transport_id = session.transport_id();
127
128 info!(self, "Transport {transport_id}: Updating quota.");
129
130 let quota = if session.can_check_quota() {
131 let folders = get_watched_folders(self).await?;
132 get_unique_quota_roots_and_usage(session, folders).await
133 } else {
134 Err(anyhow!(stock_str::not_supported_by_provider(self)))
135 };
136
137 if let Ok(quota) = "a {
138 match get_highest_usage(quota) {
139 Ok((highest, _, _)) => {
140 if needs_quota_warning(
141 highest,
142 self.get_config_int(Config::QuotaExceeding).await? as u64,
143 ) {
144 self.set_config_internal(
145 Config::QuotaExceeding,
146 Some(&highest.to_string()),
147 )
148 .await?;
149 let mut msg = Message::new_text(stock_str::quota_exceeding(self, highest));
150 add_device_msg_with_importance(self, None, Some(&mut msg), true).await?;
151 } else if highest <= QUOTA_ALLCLEAR_PERCENTAGE {
152 self.set_config_internal(Config::QuotaExceeding, None)
153 .await?;
154 }
155 }
156 Err(err) => warn!(
157 self,
158 "Transport {transport_id}: Cannot get highest quota usage: {err:#}"
159 ),
160 }
161 }
162
163 self.quota.write().await.insert(
164 transport_id,
165 QuotaInfo {
166 recent: quota,
167 modified: tools::Time::now(),
168 },
169 );
170
171 info!(self, "Transport {transport_id}: Updated quota.");
172 self.emit_event(EventType::ConnectivityChanged);
173 Ok(())
174 }
175}
176
177#[cfg(test)]
178mod tests {
179 use super::*;
180 use crate::test_utils::TestContextManager;
181
182 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
183 async fn test_needs_quota_warning() -> Result<()> {
184 assert!(!needs_quota_warning(0, 0));
185 assert!(!needs_quota_warning(10, 0));
186 assert!(!needs_quota_warning(70, 0));
187 assert!(!needs_quota_warning(75, 0));
188 assert!(!needs_quota_warning(79, 0));
189 assert!(needs_quota_warning(80, 0));
190 assert!(needs_quota_warning(81, 0));
191 assert!(!needs_quota_warning(85, 80));
192 assert!(!needs_quota_warning(85, 81));
193 assert!(needs_quota_warning(95, 82));
194 assert!(!needs_quota_warning(97, 95));
195 assert!(!needs_quota_warning(97, 96));
196 assert!(!needs_quota_warning(1000, 96));
197 Ok(())
198 }
199
200 #[expect(clippy::assertions_on_constants)]
201 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
202 async fn test_quota_thresholds() -> anyhow::Result<()> {
203 assert!(QUOTA_ALLCLEAR_PERCENTAGE > 50);
204 assert!(QUOTA_ALLCLEAR_PERCENTAGE < QUOTA_WARN_THRESHOLD_PERCENTAGE);
205 assert!(QUOTA_WARN_THRESHOLD_PERCENTAGE < QUOTA_ERROR_THRESHOLD_PERCENTAGE);
206 assert!(QUOTA_ERROR_THRESHOLD_PERCENTAGE < 100);
207 Ok(())
208 }
209
210 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
211 async fn test_quota_needs_update() -> Result<()> {
212 let mut tcm = TestContextManager::new();
213 let t = &tcm.unconfigured().await;
214 const TIMEOUT: u64 = 60;
215 assert!(t.quota_needs_update(0, TIMEOUT).await);
216
217 *t.quota.write().await = {
218 let mut map = BTreeMap::new();
219 map.insert(
220 0,
221 QuotaInfo {
222 recent: Ok(Default::default()),
223 modified: tools::Time::now() - Duration::from_secs(TIMEOUT + 1),
224 },
225 );
226 map
227 };
228 assert!(t.quota_needs_update(0, TIMEOUT).await);
229
230 *t.quota.write().await = {
231 let mut map = BTreeMap::new();
232 map.insert(
233 0,
234 QuotaInfo {
235 recent: Ok(Default::default()),
236 modified: tools::Time::now(),
237 },
238 );
239 map
240 };
241 assert!(!t.quota_needs_update(0, TIMEOUT).await);
242
243 t.evtracker.clear_events();
244 t.set_primary_self_addr("new@addr").await?;
245 assert!(t.quota.read().await.is_empty());
246 t.evtracker
247 .get_matching(|evt| matches!(evt, EventType::ConnectivityChanged))
248 .await;
249 assert!(t.quota_needs_update(0, TIMEOUT).await);
250
251 Ok(())
252 }
253}