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

inspect/
initiate.rs
