query_planning/error/context.rs
1//! Typed failure evidence carried by a [`super::QueryError`].
2//!
3//! Code that decides (retry, repartition, cleanup, attempt reports) reads these
4//! fields. [`super::QueryError::wire_facts`] renders them under the diagnostic
5//! fact keys that reports, Flight, and TSV already use.
6
7use std::borrow::Cow;
8
9use crate::plan::{JoinSide, NodeId};
10
11/// Typed evidence about why and where a query failed.
12///
13/// Every field is filled by the owner that observed it. Decoration keeps the
14/// first value: an outer layer never replaces the original failure location
15/// or allocator snapshot.
16#[derive(Debug, Clone, Default, PartialEq, Eq)]
17#[non_exhaustive]
18pub struct FailureContext {
19 /// Allocator snapshot of the budget scope the failing operation ran in.
20 pub memory: Option<MemoryFailure>,
21 /// Execution boundary that owned the failed unit of work.
22 pub boundary: Option<ResourceBoundary>,
23 /// Native allocator exception family that reported an allocation failure.
24 pub allocation_source: Option<AllocationSource>,
25 /// High-water mark of a completed resident-graph cache scope, present only
26 /// after that scope finished its fence and health proof. Not a wire fact.
27 pub graph_cache_scope_peak_bytes: Option<u64>,
28 /// An owner-reported cleanup wait reached its deadline.
29 pub cleanup_deadline_elapsed: bool,
30}
31
32/// The budget scope a failing operation ran in.
33///
34/// A snapshot shows the scope the error passed through; it does not prove a
35/// query-limit rejection. The error code carries the allocation attribution.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct MemoryFailure {
38 /// Stage boundary label of the budget request.
39 pub boundary: Cow<'static, str>,
40 /// Operation that ran under the budget scope.
41 pub operation: Cow<'static, str>,
42 /// Allocator counters of the scope's member, when sampling succeeded.
43 pub usage: Option<MemoryUsage>,
44}
45
46/// Allocator counters of one query-domain member.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub struct MemoryUsage {
49 /// Allocation limit the allocator enforced for the member.
50 pub limit_bytes: u64,
51 /// Highest requested bytes held concurrently by this member.
52 pub peak_bytes: u64,
53 /// Highest allocator-aligned bytes held concurrently by this member.
54 pub peak_aligned_bytes: u64,
55 /// Requested allocation bytes still held at the snapshot.
56 pub current_bytes: u64,
57 /// Aligned allocation bytes still held at the snapshot.
58 pub current_aligned_bytes: u64,
59 /// Number of rejected or failed allocations recorded for the member.
60 pub failed_allocation_count: u64,
61 /// Largest failed allocation request in bytes; zero when none was recorded.
62 pub largest_failed_allocation_bytes: u64,
63}
64
65/// Execution boundary that owned a failed unit of work.
66///
67/// Labels are diagnostics. Retry and repartition decisions read `side`
68/// together with the error code, never the labels.
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub struct ResourceBoundary {
71 /// Stable label of the memory-accounting boundary.
72 pub memory_boundary: &'static str,
73 /// Stage label for the failed operation.
74 pub stage: &'static str,
75 /// Plan node associated with the failure, or `None` when unavailable.
76 pub node_id: Option<NodeId>,
77 /// Semantic join input whose working set failed.
78 pub side: Option<JoinSide>,
79 /// Where execution can resume after a local retry.
80 pub resume: Option<&'static str>,
81}
82
83impl ResourceBoundary {
84 /// A boundary whose memory boundary and stage share one label.
85 pub const fn new(memory_boundary: &'static str, node_id: NodeId) -> Self {
86 Self { memory_boundary, stage: memory_boundary, node_id: Some(node_id), side: None, resume: None }
87 }
88
89 #[must_use]
90 /// Sets the diagnostic execution-stage label.
91 pub const fn with_stage(mut self, stage: &'static str) -> Self {
92 self.stage = stage;
93 self
94 }
95
96 #[must_use]
97 /// Records the semantic join input whose working set failed.
98 pub const fn with_side(mut self, side: JoinSide) -> Self {
99 self.side = Some(side);
100 self
101 }
102
103 #[must_use]
104 /// Records the execution position available to an operator-owned local retry.
105 pub const fn with_resume(mut self, resume: &'static str) -> Self {
106 self.resume = Some(resume);
107 self
108 }
109}
110
111/// Native allocator exception family behind an allocation failure.
112///
113/// `RmmOutOfMemory` is the callback-backed path through the query memory
114/// resource. The `bad_alloc` variants identify direct throwing paths; when
115/// they traverse a resident cache allocation scope that scope publishes an
116/// exact allocator event independently. `Unknown` stays fail-closed.
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118pub enum AllocationSource {
119 /// No allocator family was established; attribution remains unknown.
120 Unknown,
121 /// RMM reported an out-of-memory exception through the query resource path.
122 RmmOutOfMemory,
123 /// A direct RMM allocation threw `rmm::bad_alloc`.
124 RmmBadAlloc,
125 /// A native allocation threw `std::bad_alloc`.
126 StdBadAlloc,
127}
128
129impl AllocationSource {
130 /// Stable allocator-family label used in diagnostic facts.
131 pub const fn as_str(self) -> &'static str {
132 match self {
133 Self::Unknown => "unknown",
134 Self::RmmOutOfMemory => "rmm_out_of_memory",
135 Self::RmmBadAlloc => "rmm_bad_alloc",
136 Self::StdBadAlloc => "std_bad_alloc",
137 }
138 }
139}
140
141/// The only allocator cap mechanism; reported under `budget_enforcement`.
142const MEMORY_ENFORCEMENT: &str = "query_domain_member";
143/// Reported under `budget_status` for every error that passed a budget scope.
144const MEMORY_STATUS: &str = "exceeded_or_lower_layer_rejected";
145
146impl FailureContext {
147 /// Wire keys and values rendered from the typed fields.
148 pub(super) fn projection(&self) -> Vec<(&'static str, Cow<'_, str>)> {
149 let mut facts = Vec::new();
150 if let Some(memory) = &self.memory {
151 facts.push(("budget_boundary", Cow::Borrowed(memory.boundary.as_ref())));
152 facts.push(("budget_operation", Cow::Borrowed(memory.operation.as_ref())));
153 facts.push(("budget_enforcement", Cow::Borrowed(MEMORY_ENFORCEMENT)));
154 facts.push(("budget_status", Cow::Borrowed(MEMORY_STATUS)));
155 if let Some(usage) = memory.usage {
156 facts.extend(
157 [
158 ("budget_bytes", usage.limit_bytes),
159 ("budget_peak_bytes", usage.peak_bytes),
160 ("budget_peak_aligned_bytes", usage.peak_aligned_bytes),
161 ("budget_current_bytes", usage.current_bytes),
162 ("budget_current_aligned_bytes", usage.current_aligned_bytes),
163 ("budget_failed_allocation_count", usage.failed_allocation_count),
164 ("budget_largest_failed_allocation_bytes", usage.largest_failed_allocation_bytes),
165 ]
166 .map(|(key, value)| (key, Cow::Owned(value.to_string()))),
167 );
168 }
169 }
170 if let Some(boundary) = &self.boundary {
171 facts.push(("memory_boundary", Cow::Borrowed(boundary.memory_boundary)));
172 facts.push(("stage", Cow::Borrowed(boundary.stage)));
173 if let Some(node_id) = boundary.node_id {
174 facts.push(("node_id", Cow::Owned(node_id.to_string())));
175 }
176 if let Some(side) = boundary.side {
177 facts.push(("side", Cow::Owned(side.to_string())));
178 }
179 if let Some(resume) = boundary.resume {
180 facts.push(("resume_boundary", Cow::Borrowed(resume)));
181 }
182 }
183 if let Some(source) = self.allocation_source {
184 facts.push(("cugraph_allocation_source", Cow::Borrowed(source.as_str())));
185 }
186 facts
187 }
188
189 pub(super) fn projects(&self, key: &str) -> bool {
190 self.projection().iter().any(|(projected, _)| *projected == key)
191 }
192}