inspect/
initiate.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Support for initiating inspection requests.
5
6mod natural_sort;
7
8use super::InspectMut;
9use super::InternalError;
10use super::InternalNode;
11use super::SensitivityLevel;
12use super::Value;
13use super::ValueKind;
14use crate::NumberFormat;
15use crate::RequestParams;
16use crate::RootParams;
17use alloc::borrow::ToOwned;
18use alloc::string::String;
19use alloc::string::ToString;
20use alloc::vec;
21use alloc::vec::Vec;
22use base64::display::Base64Display;
23use core::cmp::Ordering;
24use core::fmt;
25use core::fmt::Write;
26use core::future::Future;
27use core::future::poll_fn;
28use core::pin::Pin;
29use core::task::Context;
30use core::task::Poll;
31use core::time::Duration;
32use mesh::MeshPayload;
33use thiserror::Error;
34
35/// A node of an inspect result.
36#[derive(Debug, Clone, PartialEq, MeshPayload)]
37#[mesh(package = "inspect")]
38#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
39pub enum Node {
40    /// The node is known to exist in the inspect tree but was not evaluated.
41    #[mesh(1)]
42    Unevaluated,
43    /// Evaluation of this node did not complete successfully.
44    #[mesh(2)]
45    Failed(Error),
46    /// A value.
47    #[mesh(3)]
48    Value(Value),
49    /// An interior node, with zero or more children.
50    #[mesh(4)]
51    Dir(Vec<Entry>),
52}
53
54#[derive(Debug, Clone, PartialEq, MeshPayload)]
55#[mesh(package = "inspect")]
56#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
57/// A directory entry.
58pub struct Entry {
59    /// The name of the entry.
60    #[mesh(1)]
61    pub name: String,
62    /// The node at this entry.
63    #[mesh(2)]
64    pub node: Node,
65    /// The sensitivity level of this entry.
66    #[mesh(3)]
67    pub sensitivity: SensitivityLevel,
68}
69
70/// A node resolution error.
71#[derive(Debug, Clone, PartialEq, Eq, Error, MeshPayload)]
72#[mesh(package = "inspect")]
73#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
74pub enum Error {
75    /// Deferred request never resolved.
76    #[error("unresolved")]
77    #[mesh(1)]
78    Unresolved,
79    /// Mesh channel error.
80    #[error("channel error: {0}")]
81    #[mesh(2)]
82    Mesh(String),
83    /// The node is immutable.
84    #[error("immutable node")]
85    #[mesh(3)]
86    Immutable,
87    /// The update value could not be applied.
88    #[error("update error: {0}")]
89    #[mesh(4)]
90    Update(String),
91    /// A requested node is not a directory.
92    #[error("not a directory")]
93    #[mesh(5)]
94    NotADirectory,
95    /// A requested node was not found.
96    #[error("not found")]
97    #[mesh(6)]
98    NotFound,
99    /// An internal error occurred.
100    #[error("internal error")]
101    #[mesh(7)]
102    Internal,
103}
104
105impl From<InternalError> for Error {
106    fn from(value: InternalError) -> Self {
107        match value {
108            InternalError::Immutable => Self::Immutable,
109            InternalError::Update(v) => Self::Update(v),
110            InternalError::NotADirectory => Self::NotADirectory,
111            InternalError::Unresolved => Self::Unresolved,
112            InternalError::Mesh(v) => Self::Mesh(v),
113        }
114    }
115}
116
117/// Implement Debug by calling Display.
118struct DebugFromDisplay<T>(T);
119
120impl<T: fmt::Display> fmt::Debug for DebugFromDisplay<T> {
121    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122        self.0.fmt(f)
123    }
124}
125
126impl fmt::Display for Node {
127    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128        match self {
129            Node::Unevaluated => f.pad("_"),
130            Node::Failed(err) => write!(f, "error ({err})"),
131            Node::Value(v) => fmt::Display::fmt(v, f),
132            Node::Dir(children) => {
133                let mut map = f.debug_map();
134                for entry in children {
135                    map.entry(
136                        &DebugFromDisplay(&entry.name),
137                        &DebugFromDisplay(&entry.node),
138                    );
139                }
140                map.finish()
141            }
142        }
143    }
144}
145
146impl Node {
147    fn merge_list(children: &mut Vec<Entry>) {
148        // Sort the new list of children.
149        children.sort_by(|a, b| natural_sort::compare(&a.name, &b.name));
150
151        // Merge duplicates.
152        {
153            let mut last: Option<&mut Entry> = None;
154            for entry in children.iter_mut() {
155                if entry.name.is_empty() {
156                    continue;
157                }
158                match &mut last {
159                    Some(last_entry) if *last_entry.name == entry.name => {
160                        last_entry
161                            .node
162                            .merge(core::mem::replace(&mut entry.node, Node::Unevaluated));
163                        entry.name.clear();
164                    }
165                    _ => {
166                        last = Some(entry);
167                    }
168                }
169            }
170        }
171
172        // Remove nameless children.
173        children.retain(|entry| !entry.name.is_empty());
174    }
175
176    fn merge(&mut self, other: Node) {
177        if matches!(other, Node::Unevaluated) {
178            return;
179        }
180        match self {
181            Node::Unevaluated | Node::Failed(_) | Node::Value(_) => {
182                // Cannot merge, so the later node takes precedence.
183                *self = other;
184            }
185
186            Node::Dir(children) => match other {
187                Node::Unevaluated | Node::Failed(_) | Node::Value(_) => {
188                    // Cannot merge, so the directory takes precedence.
189                }
190                Node::Dir(mut other_children) => {
191                    children.append(&mut other_children);
192                    Self::merge_list(children);
193                }
194            },
195        }
196    }
197
198    fn skip(mut self, mut n: usize) -> Node {
199        while n > 0 {
200            self = match self {
201                Node::Dir(d) => {
202                    if d.len() == 1 {
203                        d.into_iter().next().unwrap().node
204                    } else if d.is_empty() {
205                        return Node::Failed(Error::NotFound);
206                    } else {
207                        // The walk should not have produced a multi-child
208                        // directory node here.
209                        return Node::Failed(Error::Internal);
210                    }
211                }
212                Node::Failed(_) | Node::Unevaluated => return self,
213                Node::Value(_) => {
214                    // The walk should not have produced a value here.
215                    return Node::Failed(Error::Internal);
216                }
217            };
218            n -= 1;
219        }
220        self
221    }
222
223    fn compute_since(&self, last: &Node, t: f64) -> Node {
224        match (self, last) {
225            (Node::Value(value), Node::Value(last)) if value.flags.count() => {
226                let kind = match (&value.kind, &last.kind) {
227                    (ValueKind::Unsigned(x), ValueKind::Unsigned(y)) => {
228                        ValueKind::Double((x - y) as f64 / t)
229                    }
230                    (ValueKind::Signed(x), ValueKind::Signed(y)) => {
231                        ValueKind::Double((x - y) as f64 / t)
232                    }
233                    (ValueKind::Float(x), ValueKind::Float(y)) => {
234                        ValueKind::Double((x - y) as f64 / t)
235                    }
236                    (ValueKind::Double(x), ValueKind::Double(y)) => ValueKind::Double((x - y) / t),
237                    (kind, _) => kind.clone(),
238                };
239                Node::Value(Value {
240                    kind,
241                    flags: value.flags,
242                })
243            }
244            (Node::Dir(this), Node::Dir(last)) => {
245                let mut children = Vec::new();
246                let mut this = this.iter().peekable();
247                let mut last = last.iter().peekable();
248                while let (Some(&this_entry), Some(&last_entry)) = (this.peek(), last.peek()) {
249                    match this_entry.name.cmp(&last_entry.name) {
250                        Ordering::Less => {
251                            children.push(this_entry.clone());
252                            this.next();
253                        }
254                        Ordering::Equal => {
255                            children.push(Entry {
256                                node: this_entry.node.compute_since(&last_entry.node, t),
257                                ..this_entry.clone()
258                            });
259                            this.next();
260                            last.next();
261                        }
262                        Ordering::Greater => {
263                            last.next();
264                        }
265                    }
266                }
267                children.extend(this.cloned());
268                Node::Dir(children)
269            }
270            (node, _) => node.clone(),
271        }
272    }
273
274    /// Computes the differences in this node from a previous snapshot of the
275    /// same node.
276    ///
277    /// For ordinary values, the result will just have the new values.
278    ///
279    /// For values marked as counters, the result will be the difference since
280    /// the last node, divided by `duration` (in seconds). The value type will
281    /// be changed to floating point to capture the non-integral portion of the
282    /// result.
283    pub fn since(&self, last: &Node, duration: Duration) -> Self {
284        self.compute_since(last, duration.as_secs_f64())
285    }
286
287    /// Returns an object that implements [`Display`](core::fmt::Display) to output JSON.
288    pub fn json(&self) -> impl '_ + fmt::Display {
289        JsonDisplay(self)
290    }
291}
292
293struct JsonDisplay<'a>(&'a Node);
294
295impl fmt::Display for JsonDisplay<'_> {
296    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
297        match self.0 {
298            Node::Unevaluated => f.write_str("null"),
299            // Errors are serialized using a key that should not occur in valid output, since those
300            // are normally valid Rust identifiers.
301            Node::Failed(error) => write!(f, r#"{{"$error":{:?}}}"#, error.to_string()),
302            Node::Value(value) => match &value.kind {
303                ValueKind::Signed(v) => {
304                    // For human readability, values that should be hex or binary are serialized as
305                    // strings.
306                    if value.flags.hex() {
307                        write!(f, r#""{v:#x}""#)
308                    } else if value.flags.binary() {
309                        write!(f, r#""{v:#b}""#)
310                    } else {
311                        write!(f, "{v}")
312                    }
313                }
314                ValueKind::Unsigned(v) => {
315                    if value.flags.hex() {
316                        write!(f, r#""{v:#x}""#)
317                    } else if value.flags.binary() {
318                        write!(f, r#""{v:#b}""#)
319                    } else {
320                        write!(f, "{v}")
321                    }
322                }
323                ValueKind::Float(v) => write!(f, "{v}"),
324                ValueKind::Double(v) => write!(f, "{v}"),
325                ValueKind::Bool(v) => write!(f, "{v}"),
326                ValueKind::String(v) => write!(f, "{v:?}"),
327                ValueKind::Bytes(b) => {
328                    // Use base64 encoding to match typical JSON conventions.
329                    write!(
330                        f,
331                        r#""{}""#,
332                        Base64Display::new(b, &base64::engine::general_purpose::STANDARD_NO_PAD)
333                    )
334                }
335            },
336            Node::Dir(children) => {
337                f.write_char('{')?;
338                let mut comma = "";
339                for entry in children {
340                    let child = JsonDisplay(&entry.node);
341                    let name = &entry.name;
342                    write!(f, "{comma}{name:?}:{child}")?;
343                    comma = ",";
344                }
345                f.write_char('}')?;
346                Ok(())
347            }
348        }
349    }
350}
351
352fn path_node_count(path: &str) -> usize {
353    path.split('/').filter(|x| !x.is_empty()).count()
354}
355
356/// A builder for an inspection request.
357pub struct InspectionBuilder<'a> {
358    path: &'a str,
359    depth: Option<usize>,
360    sensitivity: Option<SensitivityLevel>,
361}
362
363impl<'a> InspectionBuilder<'a> {
364    /// Creates a new builder for an inspection request.
365    pub fn new(path: &'a str) -> Self {
366        Self {
367            path,
368            depth: None,
369            sensitivity: None,
370        }
371    }
372
373    /// Sets the maximum depth of the inspection request.
374    pub fn depth(mut self, depth: Option<usize>) -> Self {
375        self.depth = depth;
376        self
377    }
378
379    /// Sets the [`SensitivityLevel`] of the inspection request.
380    pub fn sensitivity(mut self, sensitivity: Option<SensitivityLevel>) -> Self {
381        self.sensitivity = sensitivity;
382        self
383    }
384
385    /// Inspects `obj` for state at the initially given `path`.
386    pub fn inspect(self, obj: impl InspectMut) -> Inspection {
387        let root = self.root();
388        let (params, skip) = self.build(&root);
389        Inspection {
390            node: params.inspect(obj),
391            skip,
392        }
393    }
394
395    /// Updates a value in `obj` at the initially given `path` to value `value`.
396    pub fn update(self, value: &'a str, obj: impl InspectMut) -> Update {
397        let mut root = self.root();
398        root.value = Some(value);
399        let (params, skip) = self.build(&root);
400        Update {
401            node: params.inspect(obj),
402            skip,
403        }
404    }
405
406    fn root(&self) -> RootParams<'_> {
407        RootParams {
408            full_path: self.path,
409            value: None,
410            sensitivity: self.sensitivity.unwrap_or(SensitivityLevel::Sensitive),
411        }
412    }
413
414    fn build<'b>(&self, root: &'b RootParams<'_>) -> (RequestParams<'b>, usize) {
415        // Account for the root node by bumping depth.
416        // Also enforce a maximum depth of 4096, anything deeper than that is
417        // most likely a bug, and we don't want to cause an infinite loop.
418        const MAX_INSPECT_DEPTH: usize = 4096;
419        let depth_with_root = if let Some(depth) = self.depth {
420            depth.saturating_add(1).min(MAX_INSPECT_DEPTH)
421        } else {
422            MAX_INSPECT_DEPTH
423        };
424        let params = RequestParams {
425            root,
426            path_start: 0,
427            depth: depth_with_root,
428            number_format: NumberFormat::default(),
429        };
430        (params, path_node_count(root.full_path))
431    }
432}
433
434/// Inspects `obj` for state at `path`.
435///
436/// ```rust
437/// # use inspect::{Inspect, Request, Node, inspect, Value, ValueKind};
438/// # use futures::executor::block_on;
439/// # use core::time::Duration;
440/// struct Obj;
441/// impl Inspect for Obj {
442///     fn inspect(&self, req: Request) {
443///         req.respond().field("field", 3);
444///     }
445/// }
446/// let mut inspection = inspect("field", &Obj);
447/// block_on(inspection.resolve());
448/// let node = inspection.results();
449/// assert!(matches!(node, Node::Value(Value { kind: ValueKind::Signed(3), .. })));
450/// ```
451pub fn inspect(path: &str, obj: impl InspectMut) -> Inspection {
452    InspectionBuilder::new(path).inspect(obj)
453}
454
455/// An active inspection, returned by [`inspect()`] or [`InspectionBuilder::inspect()`].
456#[derive(Debug)]
457pub struct Inspection {
458    node: InternalNode,
459    skip: usize,
460}
461
462impl Inspection {
463    /// Resolves any deferred inspection nodes, waiting indefinitely until they
464    /// are responded to.
465    ///
466    /// This future may be dropped (e.g. after a timeout) to stop collecting
467    /// inspection results without losing results that have already been
468    /// collected.
469    pub async fn resolve(&mut self) {
470        self.node.resolve().await
471    }
472
473    /// Returns the current results of the inspection.
474    ///
475    /// This may have unresolved nodes if [`Self::resolve`] was not called or
476    /// was cancelled before it completed.
477    pub fn results(self) -> Node {
478        self.node.into_node().skip(self.skip)
479    }
480}
481
482/// Updates a value in `obj` at `path` to value `value`.
483pub fn update(path: &str, value: &str, obj: impl InspectMut) -> Update {
484    InspectionBuilder::new(path).update(value, obj)
485}
486
487/// An active update operation, returned by [`update()`] or [`InspectionBuilder::update()`].
488pub struct Update {
489    node: InternalNode,
490    skip: usize,
491}
492
493impl Future for Update {
494    type Output = Result<Value, Error>;
495
496    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
497        let this = self.get_mut();
498        core::task::ready!(this.node.poll_resolve(cx));
499        Poll::Ready(
500            match std::mem::replace(&mut this.node, InternalNode::Unevaluated)
501                .into_node()
502                .skip(this.skip)
503            {
504                Node::Unevaluated => Err(Error::Unresolved),
505                Node::Failed(err) => Err(err),
506                Node::Value(v) => Ok(v),
507                Node::Dir(_) => Err(Error::Unresolved),
508            },
509        )
510    }
511}
512
513impl InternalNode {
514    fn poll_resolve(&mut self, cx: &mut Context<'_>) -> Poll<()> {
515        loop {
516            match self {
517                InternalNode::Dir(children) => {
518                    // Poll each child, even if a previous one returned pending,
519                    // in order to collect as many results as possible. This is
520                    // important in case the resolve operation is timed out.
521                    if !children
522                        .iter_mut()
523                        .all(|entry| entry.node.poll_resolve(cx).is_ready())
524                    {
525                        break Poll::Pending;
526                    }
527                    // Remember that this node is resolved to avoid recursing
528                    // again.
529                    *self = InternalNode::DirResolved(core::mem::take(children));
530                }
531                InternalNode::Deferred(recv) => match Pin::new(recv).poll(cx) {
532                    Poll::Ready(node) => {
533                        *self = match node {
534                            Ok(node) => {
535                                // N.B. This could be another deferred node, or
536                                // a directory with deferred child nodes.
537                                node
538                            }
539                            Err(err) => InternalNode::Failed(match err {
540                                mesh::RecvError::Closed => InternalError::Unresolved,
541                                mesh::RecvError::Error(err) => InternalError::Mesh(err.to_string()),
542                            }),
543                        };
544                    }
545                    _ => break Poll::Pending,
546                },
547                _ => break Poll::Ready(()),
548            }
549        }
550    }
551
552    async fn resolve(&mut self) {
553        poll_fn(|cx| self.poll_resolve(cx)).await
554    }
555
556    fn into_node(self) -> Node {
557        match self {
558            InternalNode::Dir(children) | InternalNode::DirResolved(children) => {
559                // Convert child nodes and merge any nameless children.
560                let mut child_nodes = Vec::new();
561                for entry in children {
562                    if matches!(entry.node, InternalNode::Ignored) {
563                        continue;
564                    }
565                    let mut child_node = entry.node.into_node();
566
567                    if entry.name.is_empty() {
568                        if let Node::Dir(grandchildren) = child_node {
569                            // No name for the node--merge the grandchildren in.
570                            child_nodes.extend(grandchildren);
571                        }
572                    } else {
573                        let mut name = entry.name;
574                        let root_len = {
575                            // Handle multi-level names like foo/bar/baz.
576                            let mut names = name.split('/');
577                            let root_name = names.next().unwrap();
578                            for interior_name in names.rev() {
579                                child_node = Node::Dir(vec![Entry {
580                                    name: interior_name.to_owned(),
581                                    node: child_node,
582                                    sensitivity: entry.sensitivity,
583                                }]);
584                            }
585                            root_name.len()
586                        };
587                        name.truncate(root_len);
588                        child_nodes.push(Entry {
589                            name,
590                            node: child_node,
591                            sensitivity: entry.sensitivity,
592                        });
593                    }
594                }
595
596                Node::merge_list(&mut child_nodes);
597                Node::Dir(child_nodes)
598            }
599            InternalNode::Value(v) => Node::Value(v),
600            InternalNode::Failed(err) => Node::Failed(err.into()),
601            InternalNode::Deferred(_) => Node::Failed(Error::Unresolved),
602            InternalNode::DepthExhausted | InternalNode::Unevaluated | InternalNode::Ignored => {
603                Node::Unevaluated
604            }
605        }
606    }
607}
inspect/initiate.rs

inspect/
initiate.rs
