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 | Node::Failed(_) => f.write_str("null"),
299            Node::Value(value) => match &value.kind {
300                ValueKind::Signed(v) => write!(f, "{}", v),
301                ValueKind::Unsigned(v) => write!(f, "{}", v),
302                ValueKind::Float(v) => write!(f, "{}", v),
303                ValueKind::Double(v) => write!(f, "{}", v),
304                ValueKind::Bool(v) => write!(f, "{}", v),
305                ValueKind::String(v) => write!(f, "{:?}", v),
306                ValueKind::Bytes(b) => {
307                    // Use base64 encoding to match typical JSON conventions.
308                    write!(
309                        f,
310                        r#""{}""#,
311                        Base64Display::new(b, &base64::engine::general_purpose::STANDARD_NO_PAD)
312                    )
313                }
314            },
315            Node::Dir(children) => {
316                f.write_char('{')?;
317                let mut comma = "";
318                for entry in children {
319                    let child = JsonDisplay(&entry.node);
320                    let name = &entry.name;
321                    write!(f, "{comma}{name:?}:{child}")?;
322                    comma = ",";
323                }
324                f.write_char('}')?;
325                Ok(())
326            }
327        }
328    }
329}
330
331fn path_node_count(path: &str) -> usize {
332    path.split('/').filter(|x| !x.is_empty()).count()
333}
334
335/// A builder for an inspection request.
336pub struct InspectionBuilder<'a> {
337    path: &'a str,
338    depth: Option<usize>,
339    sensitivity: Option<SensitivityLevel>,
340}
341
342impl<'a> InspectionBuilder<'a> {
343    /// Creates a new builder for an inspection request.
344    pub fn new(path: &'a str) -> Self {
345        Self {
346            path,
347            depth: None,
348            sensitivity: None,
349        }
350    }
351
352    /// Sets the maximum depth of the inspection request.
353    pub fn depth(mut self, depth: Option<usize>) -> Self {
354        self.depth = depth;
355        self
356    }
357
358    /// Sets the [`SensitivityLevel`] of the inspection request.
359    pub fn sensitivity(mut self, sensitivity: Option<SensitivityLevel>) -> Self {
360        self.sensitivity = sensitivity;
361        self
362    }
363
364    /// Inspects `obj` for state at the initially given `path`.
365    pub fn inspect(self, obj: impl InspectMut) -> Inspection {
366        let root = self.root();
367        let (params, skip) = self.build(&root);
368        Inspection {
369            node: params.inspect(obj),
370            skip,
371        }
372    }
373
374    /// Updates a value in `obj` at the initially given `path` to value `value`.
375    pub fn update(self, value: &'a str, obj: impl InspectMut) -> Update {
376        let mut root = self.root();
377        root.value = Some(value);
378        let (params, skip) = self.build(&root);
379        Update {
380            node: params.inspect(obj),
381            skip,
382        }
383    }
384
385    fn root(&self) -> RootParams<'_> {
386        RootParams {
387            full_path: self.path,
388            value: None,
389            sensitivity: self.sensitivity.unwrap_or(SensitivityLevel::Sensitive),
390        }
391    }
392
393    fn build<'b>(&self, root: &'b RootParams<'_>) -> (RequestParams<'b>, usize) {
394        // Account for the root node by bumping depth.
395        // Also enforce a maximum depth of 4096, anything deeper than that is
396        // most likely a bug, and we don't want to cause an infinite loop.
397        const MAX_INSPECT_DEPTH: usize = 4096;
398        let depth_with_root = if let Some(depth) = self.depth {
399            depth.saturating_add(1).min(MAX_INSPECT_DEPTH)
400        } else {
401            MAX_INSPECT_DEPTH
402        };
403        let params = RequestParams {
404            root,
405            path_start: 0,
406            depth: depth_with_root,
407            number_format: NumberFormat::default(),
408        };
409        (params, path_node_count(root.full_path))
410    }
411}
412
413/// Inspects `obj` for state at `path`.
414///
415/// ```rust
416/// # use inspect::{Inspect, Request, Node, inspect, Value, ValueKind};
417/// # use futures::executor::block_on;
418/// # use core::time::Duration;
419/// struct Obj;
420/// impl Inspect for Obj {
421///     fn inspect(&self, req: Request) {
422///         req.respond().field("field", 3);
423///     }
424/// }
425/// let mut inspection = inspect("field", &Obj);
426/// block_on(inspection.resolve());
427/// let node = inspection.results();
428/// assert!(matches!(node, Node::Value(Value { kind: ValueKind::Signed(3), .. })));
429/// ```
430pub fn inspect(path: &str, obj: impl InspectMut) -> Inspection {
431    InspectionBuilder::new(path).inspect(obj)
432}
433
434/// An active inspection, returned by [`inspect()`] or [`InspectionBuilder::inspect()`].
435#[derive(Debug)]
436pub struct Inspection {
437    node: InternalNode,
438    skip: usize,
439}
440
441impl Inspection {
442    /// Resolves any deferred inspection nodes, waiting indefinitely until they
443    /// are responded to.
444    ///
445    /// This future may be dropped (e.g. after a timeout) to stop collecting
446    /// inspection results without losing results that have already been
447    /// collected.
448    pub async fn resolve(&mut self) {
449        self.node.resolve().await
450    }
451
452    /// Returns the current results of the inspection.
453    ///
454    /// This may have unresolved nodes if [`Self::resolve`] was not called or
455    /// was cancelled before it completed.
456    pub fn results(self) -> Node {
457        self.node.into_node().skip(self.skip)
458    }
459}
460
461/// Updates a value in `obj` at `path` to value `value`.
462pub fn update(path: &str, value: &str, obj: impl InspectMut) -> Update {
463    InspectionBuilder::new(path).update(value, obj)
464}
465
466/// An active update operation, returned by [`update()`] or [`InspectionBuilder::update()`].
467pub struct Update {
468    node: InternalNode,
469    skip: usize,
470}
471
472impl Future for Update {
473    type Output = Result<Value, Error>;
474
475    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
476        let this = self.get_mut();
477        core::task::ready!(this.node.poll_resolve(cx));
478        Poll::Ready(
479            match std::mem::replace(&mut this.node, InternalNode::Unevaluated)
480                .into_node()
481                .skip(this.skip)
482            {
483                Node::Unevaluated => Err(Error::Unresolved),
484                Node::Failed(err) => Err(err),
485                Node::Value(v) => Ok(v),
486                Node::Dir(_) => Err(Error::Unresolved),
487            },
488        )
489    }
490}
491
492impl InternalNode {
493    fn poll_resolve(&mut self, cx: &mut Context<'_>) -> Poll<()> {
494        loop {
495            match self {
496                InternalNode::Dir(children) => {
497                    // Poll each child, even if a previous one returned pending,
498                    // in order to collect as many results as possible. This is
499                    // important in case the resolve operation is timed out.
500                    if !children
501                        .iter_mut()
502                        .all(|entry| entry.node.poll_resolve(cx).is_ready())
503                    {
504                        break Poll::Pending;
505                    }
506                    // Remember that this node is resolved to avoid recursing
507                    // again.
508                    *self = InternalNode::DirResolved(core::mem::take(children));
509                }
510                InternalNode::Deferred(recv) => match Pin::new(recv).poll(cx) {
511                    Poll::Ready(node) => {
512                        *self = match node {
513                            Ok(node) => {
514                                // N.B. This could be another deferred node, or
515                                // a directory with deferred child nodes.
516                                node
517                            }
518                            Err(err) => InternalNode::Failed(match err {
519                                mesh::RecvError::Closed => InternalError::Unresolved,
520                                mesh::RecvError::Error(err) => InternalError::Mesh(err.to_string()),
521                            }),
522                        };
523                    }
524                    _ => break Poll::Pending,
525                },
526                _ => break Poll::Ready(()),
527            }
528        }
529    }
530
531    async fn resolve(&mut self) {
532        poll_fn(|cx| self.poll_resolve(cx)).await
533    }
534
535    fn into_node(self) -> Node {
536        match self {
537            InternalNode::Dir(children) | InternalNode::DirResolved(children) => {
538                // Convert child nodes and merge any nameless children.
539                let mut child_nodes = Vec::new();
540                for entry in children {
541                    if matches!(entry.node, InternalNode::Ignored) {
542                        continue;
543                    }
544                    let mut child_node = entry.node.into_node();
545
546                    if entry.name.is_empty() {
547                        if let Node::Dir(grandchildren) = child_node {
548                            // No name for the node--merge the grandchildren in.
549                            child_nodes.extend(grandchildren);
550                        }
551                    } else {
552                        let mut name = entry.name;
553                        let root_len = {
554                            // Handle multi-level names like foo/bar/baz.
555                            let mut names = name.split('/');
556                            let root_name = names.next().unwrap();
557                            for interior_name in names.rev() {
558                                child_node = Node::Dir(vec![Entry {
559                                    name: interior_name.to_owned(),
560                                    node: child_node,
561                                    sensitivity: entry.sensitivity,
562                                }]);
563                            }
564                            root_name.len()
565                        };
566                        name.truncate(root_len);
567                        child_nodes.push(Entry {
568                            name,
569                            node: child_node,
570                            sensitivity: entry.sensitivity,
571                        });
572                    }
573                }
574
575                Node::merge_list(&mut child_nodes);
576                Node::Dir(child_nodes)
577            }
578            InternalNode::Value(v) => Node::Value(v),
579            InternalNode::Failed(err) => Node::Failed(err.into()),
580            InternalNode::Deferred(_) => Node::Failed(Error::Unresolved),
581            InternalNode::DepthExhausted | InternalNode::Unevaluated | InternalNode::Ignored => {
582                Node::Unevaluated
583            }
584        }
585    }
586}
inspect/initiate.rs

inspect/
initiate.rs
