datadog_agent_runtime/remote_agent/
status.rs1use std::collections::{hash_map::Entry, HashMap};
4
5use datadog_protos::agent::status::v1::{GetStatusDetailsResponse, StatusSection};
6
7pub trait StatusSectionProvider: Send + Sync + 'static {
16 fn write_status(&self, builder: &mut StatusBuilder);
18}
19
20pub struct StatusBuilder {
25 main_section: StatusSection,
26 named_sections: HashMap<String, StatusSection>,
27}
28
29impl StatusBuilder {
30 pub fn new() -> Self {
32 Self {
33 main_section: StatusSection { fields: HashMap::new() },
34 named_sections: HashMap::new(),
35 }
36 }
37
38 pub fn main_section(&mut self) -> StatusSectionWriter<'_> {
40 StatusSectionWriter {
41 section: &mut self.main_section,
42 }
43 }
44
45 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 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
73pub struct StatusSectionWriter<'a> {
75 section: &'a mut StatusSection,
76}
77
78impl StatusSectionWriter<'_> {
79 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}