|
| 1 | +/* |
| 2 | + * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | + * |
| 4 | + * This source code is licensed under the MIT license found in the |
| 5 | + * LICENSE file in the root directory of this source tree. |
| 6 | + */ |
| 7 | + |
| 8 | +//! Demand tree collection for debugging and testing laziness. |
| 9 | +//! |
| 10 | +//! A [`DemandCollector`] records cross-module demand calls into a tree. A |
| 11 | +//! collector is usually owned by a `Transaction`, scoping collection to a |
| 12 | +//! single check run — so parallel checks don't interfere with one another. |
| 13 | +//! Parent/child nesting is tracked via a per-thread scratch stack. |
| 14 | +//! |
| 15 | +//! The tree is machine-readable via serde. `#[serde(default)]` on optional |
| 16 | +//! fields lets the schema grow without breaking existing consumers. |
| 17 | +
|
| 18 | +use std::cell::RefCell; |
| 19 | +use std::fmt; |
| 20 | +use std::sync::Arc; |
| 21 | +use std::sync::Mutex; |
| 22 | + |
| 23 | +use serde::Deserialize; |
| 24 | +use serde::Serialize; |
| 25 | + |
| 26 | +/// What kind of cross-module demand a node represents. |
| 27 | +#[derive(Debug, Clone, Serialize, Deserialize)] |
| 28 | +#[serde(tag = "type", rename_all = "snake_case")] |
| 29 | +pub enum DemandKind { |
| 30 | + /// A leaf event for an Exports-level demand (e.g. `module_exists`, |
| 31 | + /// `export_exists`). Carries the reason string identifying which |
| 32 | + /// `LookupExport` method triggered the demand. |
| 33 | + Exports { reason: String }, |
| 34 | + /// A cross-module Answer lookup span. May have children if the |
| 35 | + /// computation recursively demanded data from other modules. |
| 36 | + /// `key` is the `Debug`-formatted lookup key. |
| 37 | + Answer { key: String }, |
| 38 | +} |
| 39 | + |
| 40 | +/// A node in the demand tree. |
| 41 | +#[derive(Debug, Clone, Serialize, Deserialize)] |
| 42 | +pub struct DemandNode { |
| 43 | + /// The module that made the demand. |
| 44 | + pub from: String, |
| 45 | + /// The module the demand was made against. |
| 46 | + pub target: String, |
| 47 | + /// What kind of demand this was. |
| 48 | + pub kind: DemandKind, |
| 49 | + /// Nested demands made while computing this one. |
| 50 | + #[serde(default, skip_serializing_if = "Vec::is_empty")] |
| 51 | + pub children: Vec<DemandNode>, |
| 52 | +} |
| 53 | + |
| 54 | +thread_local! { |
| 55 | + /// Per-thread stack of in-flight parent spans, used to nest closed demand |
| 56 | + /// nodes under their caller. This is inherently per-thread scratch space: |
| 57 | + /// enter and exit always happen on the same thread, and the stack empties |
| 58 | + /// itself as long as every enter has a matching exit. |
| 59 | + static STACK: RefCell<Vec<DemandNode>> = const { RefCell::new(Vec::new()) }; |
| 60 | +} |
| 61 | + |
| 62 | +/// A demand-tree collection session. Cloning produces another handle to the |
| 63 | +/// same underlying roots (the collector is reference-counted internally). |
| 64 | +#[derive(Clone, Default)] |
| 65 | +pub struct DemandCollector { |
| 66 | + roots: Arc<Mutex<Vec<DemandNode>>>, |
| 67 | +} |
| 68 | + |
| 69 | +impl DemandCollector { |
| 70 | + pub fn new() -> Self { |
| 71 | + Self::default() |
| 72 | + } |
| 73 | + |
| 74 | + /// Open a demand span for a cross-module Answer lookup. Hold the returned |
| 75 | + /// guard for the duration of the demand — dropping the guard (including |
| 76 | + /// on unwind) closes the span and attaches it to its parent, keeping |
| 77 | + /// enter/exit balanced even across panics. `key` is the lookup key, |
| 78 | + /// formatted via `Debug`. |
| 79 | + #[inline] |
| 80 | + pub fn enter( |
| 81 | + &self, |
| 82 | + from: impl fmt::Display, |
| 83 | + target: impl fmt::Display, |
| 84 | + key: impl fmt::Debug, |
| 85 | + ) -> DemandSpan<'_> { |
| 86 | + STACK.with(|stack| { |
| 87 | + stack.borrow_mut().push(DemandNode { |
| 88 | + from: from.to_string(), |
| 89 | + target: target.to_string(), |
| 90 | + kind: DemandKind::Answer { |
| 91 | + key: format!("{key:?}"), |
| 92 | + }, |
| 93 | + children: Vec::new(), |
| 94 | + }); |
| 95 | + }); |
| 96 | + DemandSpan { collector: self } |
| 97 | + } |
| 98 | + |
| 99 | + /// Record a leaf event for an Exports-level demand. `reason` identifies |
| 100 | + /// which `LookupExport` method triggered the demand. |
| 101 | + #[inline] |
| 102 | + pub fn exports_event(&self, from: impl fmt::Display, target: impl fmt::Display, reason: &str) { |
| 103 | + self.attach(DemandNode { |
| 104 | + from: from.to_string(), |
| 105 | + target: target.to_string(), |
| 106 | + kind: DemandKind::Exports { |
| 107 | + reason: reason.to_owned(), |
| 108 | + }, |
| 109 | + children: Vec::new(), |
| 110 | + }); |
| 111 | + } |
| 112 | + |
| 113 | + /// Attach a completed node either to the current parent on the stack or, |
| 114 | + /// if no parent is in flight, to the collector's shared root list. |
| 115 | + fn attach(&self, node: DemandNode) { |
| 116 | + let leftover = STACK.with(|stack| { |
| 117 | + let mut stack = stack.borrow_mut(); |
| 118 | + match stack.last_mut() { |
| 119 | + Some(parent) => { |
| 120 | + parent.children.push(node); |
| 121 | + None |
| 122 | + } |
| 123 | + None => Some(node), |
| 124 | + } |
| 125 | + }); |
| 126 | + if let Some(node) = leftover { |
| 127 | + self.roots.lock().unwrap().push(node); |
| 128 | + } |
| 129 | + } |
| 130 | + |
| 131 | + /// Take the collected tree roots, leaving the collector empty. |
| 132 | + pub fn take_roots(&self) -> Vec<DemandNode> { |
| 133 | + std::mem::take(&mut *self.roots.lock().unwrap()) |
| 134 | + } |
| 135 | +} |
| 136 | + |
| 137 | +/// RAII guard for an in-flight Answer demand span. Dropping the guard — |
| 138 | +/// whether via normal control flow or unwind — pops the span off the |
| 139 | +/// per-thread stack and attaches it to its parent (or to the collector's |
| 140 | +/// roots if no parent is in flight). |
| 141 | +#[must_use = "demand span is closed when the guard is dropped; hold it for the duration of the demand"] |
| 142 | +pub struct DemandSpan<'a> { |
| 143 | + collector: &'a DemandCollector, |
| 144 | +} |
| 145 | + |
| 146 | +impl Drop for DemandSpan<'_> { |
| 147 | + fn drop(&mut self) { |
| 148 | + if let Some(node) = STACK.with(|stack| stack.borrow_mut().pop()) { |
| 149 | + self.collector.attach(node); |
| 150 | + } |
| 151 | + } |
| 152 | +} |
![meta-codesync[bot]](https://avatars.githubusercontent.com/in/1373969?v=4&size=40){kind=avatar}
0 commit comments