1use std::{collections::BTreeMap, path::Path};
11
12use serde::{Deserialize, Serialize};
13
14use crate::config::Backend;
15#[cfg(test)]
16mod backend_tests;
17
18pub const FILE: &str = "limits.json";
19const MAX_BYTES: u64 = 64 * 1024;
21const MAX_SLOTS: u32 = 100_000;
24pub const SESSIONS_MAX_BYTES: u64 = 1024 * 1024 * 1024;
26pub const SESSIONS_TARGET_BYTES: u64 = 768 * 1024 * 1024;
28const SESSION_TARGET_NUMERATOR: u64 = 3;
31const SESSION_TARGET_DENOMINATOR: u64 = 4;
32
33#[derive(Debug, Clone, Default, PartialEq, Eq)]
35pub struct Limits {
36 backends: BTreeMap<Backend, u32>,
37 agent_sessions: AgentSessions,
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
43#[serde(rename_all = "camelCase")]
44pub struct AgentSessions {
45 pub enabled: bool,
46 pub max_bytes: u64,
47 pub target_bytes: u64,
48}
49
50impl Default for AgentSessions {
51 fn default() -> Self {
52 Self {
53 enabled: true,
54 max_bytes: SESSIONS_MAX_BYTES,
55 target_bytes: SESSIONS_TARGET_BYTES,
56 }
57 }
58}
59
60#[derive(Deserialize)]
61#[serde(deny_unknown_fields)]
62struct Declared {
63 #[serde(default, deserialize_with = "backend_limits")]
64 backends: BTreeMap<Backend, u32>,
65 #[serde(default, rename = "agentSessions")]
66 agent_sessions: Option<DeclaredSessions>,
67}
68
69fn backend_limits<'de, D: serde::Deserializer<'de>>(
70 deserializer: D,
71) -> Result<BTreeMap<Backend, u32>, D::Error> {
72 let names = BTreeMap::<String, u32>::deserialize(deserializer)?;
75 names
76 .into_iter()
77 .map(|(name, limit)| {
78 let backend =
79 serde_json::from_value::<Backend>(serde_json::Value::String(name.clone()))
80 .map_err(|error| serde::de::Error::custom(format!("backends: {error}")))?;
81 if !(1..=MAX_SLOTS).contains(&limit) {
82 return Err(serde::de::Error::custom(format!(
83 "backends.{name} must be between 1 and {MAX_SLOTS}."
84 )));
85 }
86 Ok((backend, limit))
87 })
88 .collect()
89}
90
91#[derive(Deserialize)]
92#[serde(rename_all = "camelCase", deny_unknown_fields)]
93struct DeclaredSessions {
94 enabled: Option<bool>,
95 max_bytes: Option<u64>,
96 target_bytes: Option<u64>,
97}
98
99impl DeclaredSessions {
100 fn resolve(self) -> Result<AgentSessions, String> {
103 let max_bytes = self.max_bytes.unwrap_or(SESSIONS_MAX_BYTES);
104 if max_bytes == 0 {
105 return Err("agentSessions.maxBytes must be at least 1.".into());
106 }
107 let target_bytes = match self.target_bytes {
108 Some(target) if target >= max_bytes => {
109 return Err(format!(
110 "agentSessions.targetBytes ({target}) must be lower than maxBytes ({max_bytes})."
111 ));
112 }
113 Some(target) => target,
114 None if SESSIONS_TARGET_BYTES < max_bytes => SESSIONS_TARGET_BYTES,
115 None => max_bytes / SESSION_TARGET_DENOMINATOR * SESSION_TARGET_NUMERATOR,
116 };
117 Ok(AgentSessions {
118 enabled: self.enabled.unwrap_or(true),
119 max_bytes,
120 target_bytes,
121 })
122 }
123}
124
125impl Limits {
126 pub fn read(state: &Path) -> Result<Self, String> {
130 let path = state.join(FILE);
131 let invalid = |message: String| format!("{}: {message}", crate::platform::path_text(&path));
132 let too_large =
133 || invalid("limits.json must be a regular file no larger than 64 KiB.".into());
134 let file = match crate::platform::open_regular(&path) {
135 Ok(file) => file,
136 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
137 return Ok(Self::default());
138 }
139 Err(error) if error.kind() == std::io::ErrorKind::InvalidInput => {
140 return Err(too_large());
141 }
142 Err(error) => return Err(invalid(error.to_string())),
143 };
144 let mut text = String::new();
145 std::io::Read::read_to_string(&mut std::io::Read::take(&file, MAX_BYTES + 1), &mut text)
146 .map_err(|e| invalid(e.to_string()))?;
147 if text.len() as u64 > MAX_BYTES {
148 return Err(too_large());
149 }
150 let declared: Declared = serde_json::from_str(&text).map_err(|error| {
151 let message = error.to_string();
152 let message = if message.starts_with("backends:") || message.starts_with("backends.") {
155 message
156 .strip_suffix(&format!(
157 " at line {} column {}",
158 error.line(),
159 error.column()
160 ))
161 .unwrap_or(&message)
162 } else {
163 &message
164 };
165 invalid(message.to_owned())
166 })?;
167 let agent_sessions = declared
168 .agent_sessions
169 .map(DeclaredSessions::resolve)
170 .transpose()
171 .map_err(invalid)?
172 .unwrap_or_default();
173 Ok(Self {
174 backends: declared.backends,
175 agent_sessions,
176 })
177 }
178
179 pub fn limit(&self, backend: Backend) -> Option<u32> {
181 self.backends.get(&backend).copied()
182 }
183
184 pub fn backends(&self) -> &BTreeMap<Backend, u32> {
185 &self.backends
186 }
187
188 pub fn agent_sessions(&self) -> AgentSessions {
189 self.agent_sessions
190 }
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use std::fs;
197
198 fn read(text: &str) -> Result<Limits, String> {
199 let directory = crate::test_os::tempdir();
200 fs::write(directory.path().join(FILE), text).unwrap();
201 Limits::read(directory.path())
202 }
203
204 #[test]
205 fn agent_sessions_default_and_validate() {
206 assert_eq!(
207 read("{}").unwrap().agent_sessions(),
208 AgentSessions {
209 enabled: true,
210 max_bytes: 1_073_741_824,
211 target_bytes: 805_306_368,
212 }
213 );
214 assert_eq!(
215 read(r#"{"agentSessions":{"enabled":false}}"#)
216 .unwrap()
217 .agent_sessions(),
218 AgentSessions {
219 enabled: false,
220 max_bytes: 1_073_741_824,
221 target_bytes: 805_306_368,
222 }
223 );
224 assert_eq!(
225 read(r#"{"agentSessions":{"maxBytes":4000,"targetBytes":1000}}"#)
226 .unwrap()
227 .agent_sessions(),
228 AgentSessions {
229 enabled: true,
230 max_bytes: 4000,
231 target_bytes: 1000
232 }
233 );
234 assert_eq!(
236 read(r#"{"agentSessions":{"maxBytes":4000}}"#)
237 .unwrap()
238 .agent_sessions()
239 .target_bytes,
240 3000
241 );
242 for (max_bytes, target_bytes) in [(1, 0), (3, 0), (5, 3), (7, 3), (8, 6), (11, 6)] {
244 assert_eq!(
245 read(&format!(
246 r#"{{"agentSessions":{{"maxBytes":{max_bytes}}}}}"#
247 ))
248 .unwrap()
249 .agent_sessions()
250 .target_bytes,
251 target_bytes
252 );
253 }
254 assert_eq!(
255 read(r#"{"agentSessions":{"maxBytes":18446744073709551615}}"#)
256 .unwrap()
257 .agent_sessions()
258 .target_bytes,
259 805306368
260 );
261 for (text, error) in [
262 (
263 r#"{"agentSessions":{"maxBytes":100,"targetBytes":100}}"#,
264 "agentSessions.targetBytes (100) must be lower than maxBytes (100).",
265 ),
266 (
267 r#"{"agentSessions":{"targetBytes":2147483648}}"#,
268 "must be lower than maxBytes (1073741824)",
269 ),
270 (
271 r#"{"agentSessions":{"maxBytes":0}}"#,
272 "agentSessions.maxBytes must be at least 1.",
273 ),
274 (r#"{"agentSessions":{"maxbytes":10}}"#, "unknown field"),
275 (r#"{"agentSessions":{"maxBytes":-1}}"#, "invalid value"),
276 (r#"{"agentSessions":{"enabled":"no"}}"#, "invalid type"),
277 (r#"{"agentSessions":true}"#, "invalid type"),
278 ] {
279 let message = read(text).unwrap_err();
280 assert!(message.contains(error), "{text}: {message}");
281 assert!(message.contains(FILE), "{message}");
282 }
283 }
284}