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