resymo_agent/
manager.rs

1use crate::collector::{self, Collector, Error};
2use crate::command::{self, Command};
3use crate::config::Commands;
4use crate::{
5    collector::{disk_free, load_avg, memory, swap},
6    config::Collectors,
7};
8use std::collections::{BTreeMap, HashMap};
9
10#[derive(Default)]
11pub struct Manager {
12    pub collectors: HashMap<String, Box<dyn Collector>>,
13    pub commands: HashMap<String, Box<dyn Command>>,
14}
15
16impl Manager {
17    pub fn new() -> Self {
18        Self {
19            collectors: Default::default(),
20            commands: Default::default(),
21        }
22    }
23
24    pub fn register_collector<N: Into<String>, C: Collector + 'static>(
25        &mut self,
26        name: N,
27        collector: C,
28    ) {
29        self.collectors.insert(name.into(), Box::new(collector));
30    }
31
32    pub fn extend_collectors(
33        &mut self,
34        collectors: impl IntoIterator<Item = (impl Into<String>, impl Collector + 'static)>,
35    ) {
36        self.collectors.extend(
37            collectors
38                .into_iter()
39                .map(|(k, v)| (k.into(), Box::new(v) as Box<dyn Collector + 'static>)),
40        );
41    }
42
43    pub fn register_command<N: Into<String>, C: Command + 'static>(&mut self, name: N, command: C) {
44        self.commands.insert(name.into(), Box::new(command));
45    }
46
47    pub fn extend_commands(
48        &mut self,
49        commands: impl IntoIterator<Item = (impl Into<String>, impl Command + 'static)>,
50    ) {
51        self.commands.extend(
52            commands
53                .into_iter()
54                .map(|(k, v)| (k.into(), Box::new(v) as Box<dyn Command + 'static>)),
55        );
56    }
57
58    pub async fn collect_one(&self, name: &str) -> Result<Option<serde_json::Value>, Error> {
59        Ok(match self.collectors.get(name) {
60            Some(collector) => Some(
61                collector
62                    .collect()
63                    .await
64                    .map_err(|err| Error::Collector(err.to_string()))?,
65            ),
66            None => None,
67        })
68    }
69
70    pub async fn collect_all(&self) -> Result<BTreeMap<String, serde_json::Value>, Error> {
71        let mut result = BTreeMap::new();
72        for (name, collector) in &self.collectors {
73            result.insert(
74                name.to_string(),
75                collector
76                    .collect()
77                    .await
78                    .map_err(|err| Error::Collector(err.to_string()))?,
79            );
80        }
81        Ok(result)
82    }
83}
84
85impl TryFrom<(Collectors, Commands)> for Manager {
86    type Error = anyhow::Error;
87
88    fn try_from((collectors, commands): (Collectors, Commands)) -> anyhow::Result<Manager> {
89        let mut manager = Manager::new();
90
91        // collectors
92
93        if !collectors.memory.disabled {
94            manager.register_collector("memory", memory::Collector);
95        }
96        if !collectors.swap.disabled {
97            manager.register_collector("swap", swap::Collector);
98        }
99        if !collectors.disk_free.disabled {
100            manager.register_collector("disk_free", disk_free::Collector);
101        }
102        if !collectors.load_avg.disabled {
103            manager.register_collector("load_avg", load_avg::Collector);
104        }
105        if !collectors.exec.disabled {
106            manager.extend_collectors(collector::exec::Collector::new(collectors.exec));
107        }
108
109        // commands
110
111        if !commands.exec.disabled {
112            manager.extend_commands(command::exec::Command::new(commands.exec));
113        }
114
115        // return
116
117        Ok(manager)
118    }
119}
resymo_agent/manager.rs

resymo_agent/
manager.rs
