Skip to main content

millipede_core/
config.rs

1//! Crawler configuration and environment-variable resolution.
2
3use crate::events::EventBus;
4use std::{
5    fmt,
6    path::{Path, PathBuf},
7    str::FromStr,
8    sync::Arc,
9    time::Duration,
10};
11
12/// Logging verbosity for crawler diagnostics.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum LogLevel {
15    /// Disables logging.
16    Off,
17    /// Emits errors only.
18    Error,
19    /// Emits warnings and errors.
20    Warn,
21    /// Emits informational messages and above.
22    Info,
23    /// Emits debug messages and above.
24    Debug,
25    /// Emits all trace messages.
26    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/// Errors produced while resolving crawler configuration.
59#[derive(Debug, thiserror::Error)]
60#[non_exhaustive]
61pub enum ConfigError {
62    /// An environment variable contains an invalid value.
63    #[error("invalid value {value:?} for {name}: {message}")]
64    InvalidEnvVar {
65        /// The environment variable name.
66        name: &'static str,
67        /// The rejected value.
68        value: String,
69        /// A description of the required format.
70        message: String,
71    },
72    /// A string is not a supported logging level.
73    #[error("invalid log level {0:?}")]
74    InvalidLogLevel(String),
75    /// The state persistence interval is zero.
76    #[error("persist_state_interval must be greater than zero")]
77    ZeroPersistStateInterval,
78}
79
80/// Builds a resolved [`Configuration`].
81#[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    /// Sets the default dataset identifier.
124    pub fn default_dataset_id(mut self, value: impl Into<String>) -> Self {
125        self.default_dataset_id = Some(value.into());
126        self
127    }
128    /// Sets the default key-value store identifier.
129    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    /// Sets the default request queue identifier.
134    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    /// Sets the storage directory.
139    pub fn storage_dir(mut self, value: impl Into<PathBuf>) -> Self {
140        self.storage_dir = Some(value.into());
141        self
142    }
143    /// Sets the maximum used CPU ratio.
144    pub fn max_used_cpu_ratio(mut self, value: f32) -> Self {
145        self.max_used_cpu_ratio = Some(value);
146        self
147    }
148    /// Sets the available memory ratio.
149    pub fn available_memory_ratio(mut self, value: f32) -> Self {
150        self.available_memory_ratio = Some(value);
151        self
152    }
153    /// Sets the memory limit in bytes.
154    pub fn memory_bytes(mut self, value: u64) -> Self {
155        self.memory_bytes = Some(value);
156        self
157    }
158    /// Sets the interval between state persistence events.
159    pub fn persist_state_interval(mut self, value: Duration) -> Self {
160        self.persist_state_interval = Some(value);
161        self
162    }
163    /// Sets whether storage is purged at startup.
164    pub fn purge_on_start(mut self, value: bool) -> Self {
165        self.purge_on_start = Some(value);
166        self
167    }
168    /// Sets the logging verbosity.
169    pub fn log_level(mut self, value: LogLevel) -> Self {
170        self.log_level = Some(value);
171        self
172    }
173
174    /// Sets the storage backend client.
175    pub fn storage_client(mut self, value: Arc<dyn crate::storage::StorageClient>) -> Self {
176        self.storage_client = Some(value);
177        self
178    }
179
180    /// Resolves this builder using process environment overrides.
181    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
316/// Fully resolved crawler configuration.
317///
318/// Unlike the interface's non-optional storage getter, core exposes an optional client because it
319/// cannot depend on `millipede-storage-memory` without a dependency cycle. The crawler builder
320/// takes the client from `CrawlerBuilder::storage_client` or `Configuration::storage_client` and
321/// returns `CrawlerBuildError::MissingStorage` when neither is set; default in-memory wiring lives
322/// with the umbrella crate's examples.
323pub 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    /// Creates an empty configuration builder.
366    pub fn builder() -> ConfigurationBuilder {
367        ConfigurationBuilder::default()
368    }
369    /// Returns the crawler event bus.
370    pub fn events(&self) -> &EventBus {
371        &self.events
372    }
373    /// Returns the default dataset identifier.
374    pub fn default_dataset_id(&self) -> &str {
375        &self.default_dataset_id
376    }
377    /// Returns the default key-value store identifier.
378    pub fn default_key_value_store_id(&self) -> &str {
379        &self.default_key_value_store_id
380    }
381    /// Returns the default request queue identifier.
382    pub fn default_request_queue_id(&self) -> &str {
383        &self.default_request_queue_id
384    }
385    /// Returns the storage directory.
386    pub fn storage_dir(&self) -> &Path {
387        &self.storage_dir
388    }
389    /// Returns the maximum used CPU ratio, when configured.
390    pub fn max_used_cpu_ratio(&self) -> Option<f32> {
391        self.max_used_cpu_ratio
392    }
393    /// Returns the available memory ratio, when configured.
394    pub fn available_memory_ratio(&self) -> Option<f32> {
395        self.available_memory_ratio
396    }
397    /// Returns the memory limit in bytes, when configured.
398    pub fn memory_bytes(&self) -> Option<u64> {
399        self.memory_bytes
400    }
401    /// Returns the interval between state persistence events.
402    pub fn persist_state_interval(&self) -> Duration {
403        self.persist_state_interval
404    }
405    /// Returns whether storage is purged at startup.
406    pub fn purge_on_start(&self) -> bool {
407        self.purge_on_start
408    }
409    /// Returns the logging verbosity.
410    pub fn log_level(&self) -> LogLevel {
411        self.log_level
412    }
413    /// Returns the configured storage client, when one was injected.
414    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}