1use std::{
2    borrow::Cow,
3    collections::hash_map::{self, HashMap},
4};
5
6use crate::{
7    encoder::EncodeMetric,
8    raw::{Metadata, MetricType, Unit},
9    registry::{is_lowercase, is_snake_case, RegistryError},
10};
11
12pub struct RegistrySystem {
50    namespace: Cow<'static, str>,
52    const_labels: Vec<(Cow<'static, str>, Cow<'static, str>)>,
53
54    pub(crate) metrics: HashMap<Metadata, Box<dyn EncodeMetric + 'static>>,
55    pub(crate) subsystems: HashMap<Cow<'static, str>, RegistrySystem>,
56}
57
58pub struct RegistrySystemBuilder {
60    prefix: Option<Cow<'static, str>>,
61    pub(crate) system_name: Cow<'static, str>,
62    const_labels: Vec<(Cow<'static, str>, Cow<'static, str>)>,
63}
64
65impl RegistrySystemBuilder {
67    pub fn with_const_labels<N, V>(mut self, labels: impl IntoIterator<Item = (N, V)>) -> Self
73    where
74        N: Into<Cow<'static, str>>,
75        V: Into<Cow<'static, str>>,
76    {
77        self.const_labels =
78            labels.into_iter().map(|(name, value)| (name.into(), value.into())).collect();
79        self
80    }
81}
82
83impl RegistrySystemBuilder {
85    fn new(system_name: impl Into<Cow<'static, str>>) -> Self {
86        Self { prefix: None, system_name: system_name.into(), const_labels: vec![] }
87    }
88
89    pub(crate) fn with_prefix(mut self, prefix: Option<impl Into<Cow<'static, str>>>) -> Self {
90        self.prefix = prefix.map(|prefix| prefix.into());
91        self
92    }
93
94    pub(crate) fn with_inherited_const_labels<N, V>(
95        mut self,
96        inherited_labels: impl IntoIterator<Item = (N, V)>,
97    ) -> Self
98    where
99        N: Into<Cow<'static, str>>,
100        V: Into<Cow<'static, str>>,
101    {
102        let labels = inherited_labels
103            .into_iter()
104            .map(|(name, value)| (name.into(), value.into()))
105            .chain(self.const_labels)
106            .collect::<Vec<_>>();
107        self.const_labels = labels;
108        self
109    }
110
111    pub(crate) fn build(self) -> RegistrySystem {
112        let namespace = match self.prefix {
113            Some(prefix) => Cow::Owned(format!("{}_{}", prefix, self.system_name)),
114            None => self.system_name,
115        };
116        RegistrySystem {
117            namespace,
118            const_labels: self.const_labels,
119            metrics: HashMap::default(),
120            subsystems: HashMap::default(),
121        }
122    }
123}
124
125impl RegistrySystem {
126    pub fn builder(system_name: impl Into<Cow<'static, str>>) -> RegistrySystemBuilder {
128        let system_name = system_name.into();
129        assert!(is_lowercase(&system_name), "invalid subsystem name, must be lowercase");
130        RegistrySystemBuilder::new(system_name)
131    }
132
133    pub fn namespace(&self) -> &str {
135        &self.namespace
136    }
137
138    pub fn constant_labels(&self) -> &[(Cow<'static, str>, Cow<'static, str>)] {
140        &self.const_labels
141    }
142}
143
144impl RegistrySystem {
146    pub fn register(
150        &mut self,
151        name: impl Into<Cow<'static, str>>,
152        help: impl Into<Cow<'static, str>>,
153        metric: impl EncodeMetric + 'static,
154    ) -> Result<&mut Self, RegistryError> {
155        self.do_register(name, help, None, metric)
156    }
157
158    pub fn register_with_unit(
163        &mut self,
164        name: impl Into<Cow<'static, str>>,
165        help: impl Into<Cow<'static, str>>,
166        unit: Unit,
167        metric: impl EncodeMetric + 'static,
168    ) -> Result<&mut Self, RegistryError> {
169        match metric.metric_type() {
170            MetricType::StateSet | MetricType::Info | MetricType::Unknown => {
171                return Err(RegistryError::MustHaveAnEmptyUnitString)
172            },
173            _ => {},
174        }
175        self.do_register(name, help, Some(unit), metric)
176    }
177
178    fn do_register(
179        &mut self,
180        name: impl Into<Cow<'static, str>>,
181        help: impl Into<Cow<'static, str>>,
182        unit: Option<Unit>,
183        metric: impl EncodeMetric + 'static,
184    ) -> Result<&mut Self, RegistryError> {
185        let name = name.into();
186        if !is_snake_case(&name) {
187            return Err(RegistryError::InvalidNameFormat);
188        }
189
190        let metadata = Metadata::new(name, help, metric.metric_type(), unit);
191        match self.metrics.entry(metadata) {
192            hash_map::Entry::Vacant(entry) => {
193                entry.insert(Box::new(metric));
194                Ok(self)
195            },
196            hash_map::Entry::Occupied(_) => Err(RegistryError::AlreadyExists),
197        }
198    }
199}
200
201impl RegistrySystem {
203    pub fn subsystem(&mut self, name: impl Into<Cow<'static, str>>) -> &mut Self {
210        let name = name.into();
211        self.subsystems.entry(name).or_insert_with_key(|name| {
212            RegistrySystem::builder(name.clone())
213                .with_prefix(Some(self.namespace.clone()))
215                .with_inherited_const_labels(self.const_labels.clone())
217                .build()
218        })
219    }
220
221    pub fn attach_subsystem(&mut self, builder: RegistrySystemBuilder) -> &mut RegistrySystem {
228        let name = builder.system_name.clone();
229        self.subsystems.entry(name).or_insert_with(|| {
230            builder
231                .with_prefix(Some(self.namespace.clone()))
233                .with_inherited_const_labels(self.const_labels.clone())
235                .build()
236        })
237    }
238}