Skip to main content

datadog_agent_runtime/remote_agent/
status.rs

1//! Status details reported to the Datadog Agent.
2
3use std::collections::{hash_map::Entry, HashMap};
4
5use datadog_protos::agent::status::v1::{GetStatusDetailsResponse, StatusSection};
6
7/// A provider of status details for the Datadog Agent's status output.
8///
9/// Each time the Datadog Agent requests status details, the status service writes its built-in main section fields and
10/// then calls every registered provider in registration order. Providers write their fields into the main section or
11/// into named sections of their own.
12///
13/// The call happens while the status request is being answered, so it **SHOULD** only read state that is already at
14/// hand, such as internal metrics, rather than performing I/O.
15pub trait StatusSectionProvider: Send + Sync + 'static {
16    /// Writes this provider's status fields into the given builder.
17    fn write_status(&self, builder: &mut StatusBuilder);
18}
19
20/// Builder for the status details reported to the Datadog Agent.
21///
22/// Status details consist of a main section, which the Datadog Agent shows under the name of the remote agent, and
23/// any number of named sections.
24pub struct StatusBuilder {
25    main_section: StatusSection,
26    named_sections: HashMap<String, StatusSection>,
27}
28
29impl StatusBuilder {
30    /// Creates an empty `StatusBuilder`.
31    pub fn new() -> Self {
32        Self {
33            main_section: StatusSection { fields: HashMap::new() },
34            named_sections: HashMap::new(),
35        }
36    }
37
38    /// Returns a writer for the main section.
39    pub fn main_section(&mut self) -> StatusSectionWriter<'_> {
40        StatusSectionWriter {
41            section: &mut self.main_section,
42        }
43    }
44
45    /// Returns a writer for the named section with the given name, creating the section if it doesn't exist yet.
46    pub fn named_section<S: AsRef<str>>(&mut self, name: S) -> StatusSectionWriter<'_> {
47        match self.named_sections.entry(name.as_ref().to_string()) {
48            Entry::Occupied(entry) => StatusSectionWriter {
49                section: entry.into_mut(),
50            },
51            Entry::Vacant(entry) => {
52                let section = entry.insert(StatusSection { fields: HashMap::new() });
53                StatusSectionWriter { section }
54            }
55        }
56    }
57
58    /// Consumes the builder and returns the status details in their wire form.
59    pub fn into_response(self) -> GetStatusDetailsResponse {
60        GetStatusDetailsResponse {
61            main_section: Some(self.main_section),
62            named_sections: self.named_sections,
63        }
64    }
65}
66
67impl Default for StatusBuilder {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73/// Writer for the fields of a single status section.
74pub struct StatusSectionWriter<'a> {
75    section: &'a mut StatusSection,
76}
77
78impl StatusSectionWriter<'_> {
79    /// Sets the field with the given name to the given value, replacing any previous value.
80    pub fn set_field<S: AsRef<str>, V: AsRef<str>>(&mut self, name: S, value: V) -> &mut Self {
81        self.section
82            .fields
83            .insert(name.as_ref().to_string(), value.as_ref().to_string());
84        self
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn named_sections_accumulate_fields_across_writers() {
94        let mut builder = StatusBuilder::new();
95        builder.main_section().set_field("Version", "1.0.0");
96        builder.named_section("Example").set_field("First", "1");
97        builder
98            .named_section("Example")
99            .set_field("Second", "2")
100            .set_field("First", "one");
101
102        let response = builder.into_response();
103        let main = response.main_section.expect("main section is always present");
104        assert_eq!(main.fields.get("Version").map(String::as_str), Some("1.0.0"));
105
106        let example = &response.named_sections["Example"];
107        assert_eq!(example.fields.len(), 2);
108        assert_eq!(example.fields.get("First").map(String::as_str), Some("one"));
109        assert_eq!(example.fields.get("Second").map(String::as_str), Some("2"));
110    }
111}