1use crate::events::EventBus;
4use std::{
5 fmt,
6 path::{Path, PathBuf},
7 str::FromStr,
8 sync::Arc,
9 time::Duration,
10};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum LogLevel {
15 Off,
17 Error,
19 Warn,
21 Info,
23 Debug,
25 Trace,
27}
28
29impl FromStr for LogLevel {
30 type Err = ConfigError;
31
32 fn from_str(value: &str) -> Result<Self, Self::Err> {
33 match value.to_ascii_lowercase().as_str() {
34 "off" => Ok(Self::Off),
35 "error" => Ok(Self::Error),
36 "warn" => Ok(Self::Warn),
37 "info" => Ok(Self::Info),
38 "debug" => Ok(Self::Debug),
39 "trace" => Ok(Self::Trace),
40 _ => Err(ConfigError::InvalidLogLevel(value.to_owned())),
41 }
42 }
43}
44
45impl fmt::Display for LogLevel {
46 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
47 formatter.write_str(match self {
48 Self::Off => "off",
49 Self::Error => "error",
50 Self::Warn => "warn",
51 Self::Info => "info",
52 Self::Debug => "debug",
53 Self::Trace => "trace",
54 })
55 }
56}
57
58#[derive(Debug, thiserror::Error)]
60#[non_exhaustive]
61pub enum ConfigError {
62 #[error("invalid value {value:?} for {name}: {message}")]
64 InvalidEnvVar {
65 name: &'static str,
67 value: String,
69 message: String,
71 },
72 #[error("invalid log level {0:?}")]
74 InvalidLogLevel(String),
75 #[error("persist_state_interval must be greater than zero")]
77 ZeroPersistStateInterval,
78}
79
80#[derive(Default, Clone)]
82#[must_use = "builders do nothing unless consumed by build"]
83pub struct ConfigurationBuilder {
84 default_dataset_id: Option<String>,
85 default_key_value_store_id: Option<String>,
86 default_request_queue_id: Option<String>,
87 storage_dir: Option<PathBuf>,
88 max_used_cpu_ratio: Option<f32>,
89 available_memory_ratio: Option<f32>,
90 memory_bytes: Option<u64>,
91 persist_state_interval: Option<Duration>,
92 purge_on_start: Option<bool>,
93 log_level: Option<LogLevel>,
94 storage_client: Option<Arc<dyn crate::storage::StorageClient>>,
95}
96
97impl fmt::Debug for ConfigurationBuilder {
98 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
99 formatter
100 .debug_struct("ConfigurationBuilder")
101 .field("default_dataset_id", &self.default_dataset_id)
102 .field(
103 "default_key_value_store_id",
104 &self.default_key_value_store_id,
105 )
106 .field("default_request_queue_id", &self.default_request_queue_id)
107 .field("storage_dir", &self.storage_dir)
108 .field("max_used_cpu_ratio", &self.max_used_cpu_ratio)
109 .field("available_memory_ratio", &self.available_memory_ratio)
110 .field("memory_bytes", &self.memory_bytes)
111 .field("persist_state_interval", &self.persist_state_interval)
112 .field("purge_on_start", &self.purge_on_start)
113 .field("log_level", &self.log_level)
114 .field(
115 "storage_client",
116 &self.storage_client.as_ref().map(|_| "<dyn StorageClient>"),
117 )
118 .finish()
119 }
120}
121
122impl ConfigurationBuilder {
123 pub fn default_dataset_id(mut self, value: impl Into<String>) -> Self {
125 self.default_dataset_id = Some(value.into());
126 self
127 }
128 pub fn default_key_value_store_id(mut self, value: impl Into<String>) -> Self {
130 self.default_key_value_store_id = Some(value.into());
131 self
132 }
133 pub fn default_request_queue_id(mut self, value: impl Into<String>) -> Self {
135 self.default_request_queue_id = Some(value.into());
136 self
137 }
138 pub fn storage_dir(mut self, value: impl Into<PathBuf>) -> Self {
140 self.storage_dir = Some(value.into());
141 self
142 }
143 pub fn max_used_cpu_ratio(mut self, value: f32) -> Self {
145 self.max_used_cpu_ratio = Some(value);
146 self
147 }
148 pub fn available_memory_ratio(mut self, value: f32) -> Self {
150 self.available_memory_ratio = Some(value);
151 self
152 }
153 pub fn memory_bytes(mut self, value: u64) -> Self {
155 self.memory_bytes = Some(value);
156 self
157 }
158 pub fn persist_state_interval(mut self, value: Duration) -> Self {
160 self.persist_state_interval = Some(value);
161 self
162 }
163 pub fn purge_on_start(mut self, value: bool) -> Self {
165 self.purge_on_start = Some(value);
166 self
167 }
168 pub fn log_level(mut self, value: LogLevel) -> Self {
170 self.log_level = Some(value);
171 self
172 }
173
174 pub fn storage_client(mut self, value: Arc<dyn crate::storage::StorageClient>) -> Self {
176 self.storage_client = Some(value);
177 self
178 }
179
180 pub fn build(self) -> Result<Configuration, ConfigError> {
182 self.build_with_env(|name| std::env::var(name).ok())
183 }
184
185 pub(crate) fn build_with_env(
186 self,
187 lookup: impl Fn(&str) -> Option<String>,
188 ) -> Result<Configuration, ConfigError> {
189 let default_dataset_id = self
190 .default_dataset_id
191 .or_else(|| lookup("CRAWLEE_DEFAULT_DATASET_ID"))
192 .unwrap_or_else(|| "default".into());
193 let default_key_value_store_id = self
194 .default_key_value_store_id
195 .or_else(|| lookup("CRAWLEE_DEFAULT_KEY_VALUE_STORE_ID"))
196 .unwrap_or_else(|| "default".into());
197 let default_request_queue_id = self
198 .default_request_queue_id
199 .or_else(|| lookup("CRAWLEE_DEFAULT_REQUEST_QUEUE_ID"))
200 .unwrap_or_else(|| "default".into());
201 let storage_dir = self
202 .storage_dir
203 .or_else(|| lookup("CRAWLEE_STORAGE_DIR").map(PathBuf::from))
204 .unwrap_or_else(|| PathBuf::from("./storage"));
205 let max_used_cpu_ratio = match self.max_used_cpu_ratio {
206 Some(value) => Some(value),
207 None => optional_parse(
208 &lookup,
209 "CRAWLEE_MAX_USED_CPU_RATIO",
210 "expected a floating-point number",
211 )?,
212 };
213 let available_memory_ratio = match self.available_memory_ratio {
214 Some(value) => Some(value),
215 None => optional_parse(
216 &lookup,
217 "CRAWLEE_AVAILABLE_MEMORY_RATIO",
218 "expected a floating-point number",
219 )?,
220 };
221 let memory_bytes = match self.memory_bytes {
222 Some(value) => Some(value),
223 None => match optional_parse::<u64>(
224 &lookup,
225 "CRAWLEE_MEMORY_MBYTES",
226 "expected an unsigned integer",
227 )? {
228 Some(megabytes) => Some(megabytes.checked_mul(1024 * 1024).ok_or_else(|| {
229 ConfigError::InvalidEnvVar {
230 name: "CRAWLEE_MEMORY_MBYTES",
231 value: megabytes.to_string(),
232 message: "value is too large to convert to bytes".into(),
233 }
234 })?),
235 None => None,
236 },
237 };
238 let persist_state_interval = match self.persist_state_interval {
239 Some(value) => value,
240 None => Duration::from_millis(
241 optional_parse(
242 &lookup,
243 "CRAWLEE_PERSIST_STATE_INTERVAL_MILLIS",
244 "expected an unsigned integer",
245 )?
246 .unwrap_or(60_000),
247 ),
248 };
249 if persist_state_interval.is_zero() {
250 return Err(ConfigError::ZeroPersistStateInterval);
251 }
252 let purge_on_start = match self.purge_on_start {
253 Some(value) => value,
254 None => match lookup("CRAWLEE_PURGE_ON_START") {
255 Some(value) => parse_bool("CRAWLEE_PURGE_ON_START", value)?,
256 None => true,
257 },
258 };
259 let log_level = match self.log_level {
260 Some(value) => value,
261 None => match lookup("CRAWLEE_LOG_LEVEL") {
262 Some(value) => value.parse().map_err(|_| ConfigError::InvalidEnvVar {
263 name: "CRAWLEE_LOG_LEVEL",
264 value,
265 message: "expected off, error, warn, info, debug, or trace".into(),
266 })?,
267 None => LogLevel::Info,
268 },
269 };
270
271 Ok(Configuration {
272 events: EventBus::default(),
273 default_dataset_id,
274 default_key_value_store_id,
275 default_request_queue_id,
276 storage_dir,
277 max_used_cpu_ratio,
278 available_memory_ratio,
279 memory_bytes,
280 persist_state_interval,
281 purge_on_start,
282 log_level,
283 storage_client: self.storage_client,
284 })
285 }
286}
287
288fn optional_parse<T: FromStr>(
289 lookup: &impl Fn(&str) -> Option<String>,
290 name: &'static str,
291 message: &str,
292) -> Result<Option<T>, ConfigError> {
293 lookup(name)
294 .map(|value| {
295 value.parse().map_err(|_| ConfigError::InvalidEnvVar {
296 name,
297 value,
298 message: message.into(),
299 })
300 })
301 .transpose()
302}
303
304fn parse_bool(name: &'static str, value: String) -> Result<bool, ConfigError> {
305 match value.to_ascii_lowercase().as_str() {
306 "1" | "true" => Ok(true),
307 "0" | "false" => Ok(false),
308 _ => Err(ConfigError::InvalidEnvVar {
309 name,
310 value,
311 message: "expected 1, true, 0, or false".into(),
312 }),
313 }
314}
315
316pub struct Configuration {
324 events: EventBus,
325 default_dataset_id: String,
326 default_key_value_store_id: String,
327 default_request_queue_id: String,
328 storage_dir: PathBuf,
329 max_used_cpu_ratio: Option<f32>,
330 available_memory_ratio: Option<f32>,
331 memory_bytes: Option<u64>,
332 persist_state_interval: Duration,
333 purge_on_start: bool,
334 log_level: LogLevel,
335 storage_client: Option<Arc<dyn crate::storage::StorageClient>>,
336}
337
338impl fmt::Debug for Configuration {
339 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
340 formatter
341 .debug_struct("Configuration")
342 .field("events", &self.events)
343 .field("default_dataset_id", &self.default_dataset_id)
344 .field(
345 "default_key_value_store_id",
346 &self.default_key_value_store_id,
347 )
348 .field("default_request_queue_id", &self.default_request_queue_id)
349 .field("storage_dir", &self.storage_dir)
350 .field("max_used_cpu_ratio", &self.max_used_cpu_ratio)
351 .field("available_memory_ratio", &self.available_memory_ratio)
352 .field("memory_bytes", &self.memory_bytes)
353 .field("persist_state_interval", &self.persist_state_interval)
354 .field("purge_on_start", &self.purge_on_start)
355 .field("log_level", &self.log_level)
356 .field(
357 "storage_client",
358 &self.storage_client.as_ref().map(|_| "<dyn StorageClient>"),
359 )
360 .finish()
361 }
362}
363
364impl Configuration {
365 pub fn builder() -> ConfigurationBuilder {
367 ConfigurationBuilder::default()
368 }
369 pub fn events(&self) -> &EventBus {
371 &self.events
372 }
373 pub fn default_dataset_id(&self) -> &str {
375 &self.default_dataset_id
376 }
377 pub fn default_key_value_store_id(&self) -> &str {
379 &self.default_key_value_store_id
380 }
381 pub fn default_request_queue_id(&self) -> &str {
383 &self.default_request_queue_id
384 }
385 pub fn storage_dir(&self) -> &Path {
387 &self.storage_dir
388 }
389 pub fn max_used_cpu_ratio(&self) -> Option<f32> {
391 self.max_used_cpu_ratio
392 }
393 pub fn available_memory_ratio(&self) -> Option<f32> {
395 self.available_memory_ratio
396 }
397 pub fn memory_bytes(&self) -> Option<u64> {
399 self.memory_bytes
400 }
401 pub fn persist_state_interval(&self) -> Duration {
403 self.persist_state_interval
404 }
405 pub fn purge_on_start(&self) -> bool {
407 self.purge_on_start
408 }
409 pub fn log_level(&self) -> LogLevel {
411 self.log_level
412 }
413 pub fn storage_client(&self) -> Option<&Arc<dyn crate::storage::StorageClient>> {
415 self.storage_client.as_ref()
416 }
417}
418
419impl Default for Configuration {
420 fn default() -> Self {
421 ConfigurationBuilder::default()
422 .build_with_env(|_| None)
423 .expect("built-in configuration defaults are valid")
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 use super::*;
430 use std::collections::HashMap;
431
432 fn with_env(
433 builder: ConfigurationBuilder,
434 values: &[(&str, &str)],
435 ) -> Result<Configuration, ConfigError> {
436 let values: HashMap<&str, &str> = values.iter().copied().collect();
437 builder.build_with_env(|name| values.get(name).map(|value| (*value).to_owned()))
438 }
439
440 #[test]
441 fn defaults_are_resolved_without_environment() {
442 let config = with_env(Configuration::builder(), &[]).unwrap();
443 assert_eq!(config.default_dataset_id(), "default");
444 assert_eq!(config.default_key_value_store_id(), "default");
445 assert_eq!(config.default_request_queue_id(), "default");
446 assert_eq!(config.storage_dir(), Path::new("./storage"));
447 assert_eq!(config.persist_state_interval(), Duration::from_secs(60));
448 assert!(config.purge_on_start());
449 assert_eq!(config.log_level(), LogLevel::Info);
450 assert_eq!(config.available_memory_ratio(), None);
451 assert_eq!(config.memory_bytes(), None);
452 }
453
454 #[test]
455 fn environment_overrides_purge_on_start() {
456 let config = with_env(
457 Configuration::builder(),
458 &[("CRAWLEE_PURGE_ON_START", "false")],
459 )
460 .unwrap();
461 assert!(!config.purge_on_start());
462 }
463
464 #[test]
465 fn environment_memory_megabytes_are_converted_to_bytes() {
466 let config = with_env(
467 Configuration::builder(),
468 &[("CRAWLEE_MEMORY_MBYTES", "512")],
469 )
470 .unwrap();
471 assert_eq!(config.memory_bytes(), Some(512 * 1024 * 1024));
472 }
473
474 #[test]
475 fn environment_max_used_cpu_ratio_is_resolved() {
476 let config = with_env(
477 Configuration::builder(),
478 &[("CRAWLEE_MAX_USED_CPU_RATIO", "0.75")],
479 )
480 .unwrap();
481 assert_eq!(config.max_used_cpu_ratio(), Some(0.75));
482 }
483
484 #[test]
485 fn environment_persist_interval_is_parsed_as_milliseconds() {
486 let config = with_env(
487 Configuration::builder(),
488 &[("CRAWLEE_PERSIST_STATE_INTERVAL_MILLIS", "5000")],
489 )
490 .unwrap();
491 assert_eq!(config.persist_state_interval(), Duration::from_secs(5));
492 }
493
494 #[test]
495 fn zero_persist_interval_is_rejected() {
496 let builder_error = with_env(
497 Configuration::builder().persist_state_interval(Duration::ZERO),
498 &[],
499 )
500 .unwrap_err();
501 assert!(matches!(
502 builder_error,
503 ConfigError::ZeroPersistStateInterval
504 ));
505
506 let environment_error = with_env(
507 Configuration::builder(),
508 &[("CRAWLEE_PERSIST_STATE_INTERVAL_MILLIS", "0")],
509 )
510 .unwrap_err();
511 assert!(matches!(
512 environment_error,
513 ConfigError::ZeroPersistStateInterval
514 ));
515 }
516
517 #[test]
518 fn environment_log_level_is_case_insensitive() {
519 let config = with_env(Configuration::builder(), &[("CRAWLEE_LOG_LEVEL", "debug")]).unwrap();
520 assert_eq!(config.log_level(), LogLevel::Debug);
521 }
522
523 #[test]
524 fn builder_value_beats_environment() {
525 let config = with_env(
526 Configuration::builder().purge_on_start(true),
527 &[("CRAWLEE_PURGE_ON_START", "false")],
528 )
529 .unwrap();
530 assert!(config.purge_on_start());
531 }
532
533 #[test]
534 fn invalid_float_identifies_environment_variable() {
535 let error = with_env(
536 Configuration::builder(),
537 &[("CRAWLEE_AVAILABLE_MEMORY_RATIO", "nope")],
538 )
539 .unwrap_err();
540 assert!(matches!(
541 error,
542 ConfigError::InvalidEnvVar {
543 name: "CRAWLEE_AVAILABLE_MEMORY_RATIO",
544 ..
545 }
546 ));
547 }
548}