1use std::collections::VecDeque;
6use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
7use std::sync::{Arc, OnceLock, PoisonError, RwLock};
8use std::time::Duration;
9
10use serde::{Deserialize, Serialize};
11
12use crate::BackendKind;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
16#[non_exhaustive]
17pub enum MemoryTier {
18 Device,
20 Host,
22 Disk,
24}
25
26impl MemoryTier {
27 const fn index(self) -> usize {
28 match self {
29 MemoryTier::Device => 0,
30 MemoryTier::Host => 1,
31 MemoryTier::Disk => 2,
32 }
33 }
34}
35
36#[derive(Debug)]
38pub struct OperatorStats {
39 id: usize,
40 name: String,
41 backend: BackendKind,
42 rows_in: AtomicU64,
43 rows_out: AtomicU64,
44 batches: AtomicU64,
45 elapsed_ns: AtomicU64,
46 bytes_h2d: AtomicU64,
47 bytes_d2h: AtomicU64,
48 memory_bytes: AtomicU64,
49 fallback_batches: AtomicU64,
50}
51
52impl OperatorStats {
53 fn new(id: usize, name: String, backend: BackendKind) -> Self {
54 Self {
55 id,
56 name,
57 backend,
58 rows_in: AtomicU64::new(0),
59 rows_out: AtomicU64::new(0),
60 batches: AtomicU64::new(0),
61 elapsed_ns: AtomicU64::new(0),
62 bytes_h2d: AtomicU64::new(0),
63 bytes_d2h: AtomicU64::new(0),
64 memory_bytes: AtomicU64::new(0),
65 fallback_batches: AtomicU64::new(0),
66 }
67 }
68
69 pub const fn id(&self) -> usize {
71 self.id
72 }
73
74 pub fn name(&self) -> &str {
76 &self.name
77 }
78
79 pub const fn backend(&self) -> BackendKind {
81 self.backend
82 }
83
84 pub fn record_batch(&self, rows_in: u64, rows_out: u64, elapsed: Duration) {
86 self.rows_in.fetch_add(rows_in, Ordering::Relaxed);
87 self.rows_out.fetch_add(rows_out, Ordering::Relaxed);
88 self.batches.fetch_add(1, Ordering::Relaxed);
89 let ns = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
90 self.elapsed_ns.fetch_add(ns, Ordering::Relaxed);
91 }
92
93 pub fn record_transfer(&self, h2d: u64, d2h: u64) {
95 self.bytes_h2d.fetch_add(h2d, Ordering::Relaxed);
96 self.bytes_d2h.fetch_add(d2h, Ordering::Relaxed);
97 }
98
99 pub fn set_memory_bytes(&self, bytes: u64) {
101 self.memory_bytes.store(bytes, Ordering::Relaxed);
102 }
103
104 pub fn record_fallback(&self) {
112 self.fallback_batches.fetch_add(1, Ordering::Relaxed);
113 }
114
115 pub fn fallback_batches(&self) -> u64 {
117 self.fallback_batches.load(Ordering::Relaxed)
118 }
119
120 pub fn snapshot(&self) -> OperatorSnapshot {
122 OperatorSnapshot {
123 id: self.id,
124 name: self.name.clone(),
125 backend: self.backend,
126 rows_in: self.rows_in.load(Ordering::Relaxed),
127 rows_out: self.rows_out.load(Ordering::Relaxed),
128 batches: self.batches.load(Ordering::Relaxed),
129 elapsed_ns: self.elapsed_ns.load(Ordering::Relaxed),
130 bytes_h2d: self.bytes_h2d.load(Ordering::Relaxed),
131 bytes_d2h: self.bytes_d2h.load(Ordering::Relaxed),
132 memory_bytes: self.memory_bytes.load(Ordering::Relaxed),
133 fallback_batches: self.fallback_batches.load(Ordering::Relaxed),
134 }
135 }
136}
137
138#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
140pub struct OperatorSnapshot {
141 pub id: usize,
143 pub name: String,
145 pub backend: BackendKind,
147 pub rows_in: u64,
149 pub rows_out: u64,
151 pub batches: u64,
153 pub elapsed_ns: u64,
155 pub bytes_h2d: u64,
157 pub bytes_d2h: u64,
159 pub memory_bytes: u64,
161 pub fallback_batches: u64,
165}
166
167impl OperatorSnapshot {
168 pub fn fell_back_entirely(&self) -> bool {
172 self.batches > 0 && self.fallback_batches >= self.batches
173 }
174
175 pub fn mean_batch_latency_ms(&self) -> f64 {
177 if self.batches == 0 {
178 0.0
179 } else {
180 self.elapsed_ns as f64 / self.batches as f64 / 1_000_000.0
181 }
182 }
183}
184
185#[derive(Debug, Default)]
187pub struct TierGauges {
188 bytes: [AtomicU64; 3],
189}
190
191impl TierGauges {
192 pub fn set(&self, tier: MemoryTier, bytes: u64) {
194 self.bytes[tier.index()].store(bytes, Ordering::Relaxed);
195 }
196
197 pub fn add(&self, tier: MemoryTier, bytes: u64) {
199 self.bytes[tier.index()].fetch_add(bytes, Ordering::Relaxed);
200 }
201
202 pub fn sub(&self, tier: MemoryTier, bytes: u64) {
204 let _ = self.bytes[tier.index()].fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
205 Some(v.saturating_sub(bytes))
206 });
207 }
208
209 pub fn get(&self, tier: MemoryTier) -> u64 {
211 self.bytes[tier.index()].load(Ordering::Relaxed)
212 }
213
214 pub fn snapshot(&self) -> TierSnapshot {
216 TierSnapshot {
217 device_bytes: self.get(MemoryTier::Device),
218 host_bytes: self.get(MemoryTier::Host),
219 disk_bytes: self.get(MemoryTier::Disk),
220 }
221 }
222}
223
224#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
226pub struct TierSnapshot {
227 pub device_bytes: u64,
229 pub host_bytes: u64,
231 pub disk_bytes: u64,
233}
234
235#[derive(Debug, Default)]
237pub struct SpillCounters {
238 demotions: AtomicU64,
239 promotions: AtomicU64,
240 spilled_bytes: AtomicU64,
241 reloaded_bytes: AtomicU64,
242}
243
244impl SpillCounters {
245 pub fn record_demotion(&self, bytes: u64) {
247 self.demotions.fetch_add(1, Ordering::Relaxed);
248 self.spilled_bytes.fetch_add(bytes, Ordering::Relaxed);
249 }
250
251 pub fn record_promotion(&self, bytes: u64) {
253 self.promotions.fetch_add(1, Ordering::Relaxed);
254 self.reloaded_bytes.fetch_add(bytes, Ordering::Relaxed);
255 }
256
257 pub fn snapshot(&self) -> SpillSnapshot {
259 SpillSnapshot {
260 demotions: self.demotions.load(Ordering::Relaxed),
261 promotions: self.promotions.load(Ordering::Relaxed),
262 spilled_bytes: self.spilled_bytes.load(Ordering::Relaxed),
263 reloaded_bytes: self.reloaded_bytes.load(Ordering::Relaxed),
264 }
265 }
266}
267
268#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
276pub struct TierCapacity {
277 pub device_bytes: Option<u64>,
279 pub host_bytes: Option<u64>,
281 pub spill_on_query_path: bool,
287}
288
289#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
291pub struct SpillSnapshot {
292 pub demotions: u64,
294 pub promotions: u64,
296 pub spilled_bytes: u64,
298 pub reloaded_bytes: u64,
300}
301
302#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
304pub struct PlanNodeSummary {
305 pub depth: usize,
307 pub name: String,
309 pub backend: BackendKind,
311 pub detail: String,
313}
314
315#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
323pub struct PlacementSkip {
324 pub node: String,
326 pub reason: String,
328}
329
330pub const MAX_PLACEMENT_SKIPS: usize = 256;
334
335pub const MAX_OPERATORS: usize = 4096;
341
342#[derive(Debug, Default)]
344pub struct TelemetryHub {
345 operators: RwLock<VecDeque<Arc<OperatorStats>>>,
346 next_operator_id: AtomicUsize,
347 tiers: TierGauges,
348 spill: SpillCounters,
349 plan: RwLock<Vec<PlanNodeSummary>>,
350 skips: RwLock<Vec<PlacementSkip>>,
351 capacity: RwLock<TierCapacity>,
352}
353
354static GLOBAL: OnceLock<Arc<TelemetryHub>> = OnceLock::new();
356
357impl TelemetryHub {
358 pub fn new() -> Arc<Self> {
360 Arc::new(Self::default())
361 }
362
363 pub fn global() -> &'static Arc<TelemetryHub> {
372 GLOBAL.get_or_init(TelemetryHub::new)
373 }
374
375 pub fn register_operator(
379 &self,
380 name: impl Into<String>,
381 backend: BackendKind,
382 ) -> Arc<OperatorStats> {
383 let id = self.next_operator_id.fetch_add(1, Ordering::Relaxed);
384 let stats = Arc::new(OperatorStats::new(id, name.into(), backend));
385 let mut ops = self
386 .operators
387 .write()
388 .unwrap_or_else(PoisonError::into_inner);
389 if ops.len() >= MAX_OPERATORS {
390 ops.pop_front();
391 }
392 ops.push_back(Arc::clone(&stats));
393 stats
394 }
395
396 pub fn clear_operators(&self) {
399 self.operators
400 .write()
401 .unwrap_or_else(PoisonError::into_inner)
402 .clear();
403 }
404
405 pub fn record_skip(&self, node: impl Into<String>, reason: impl Into<String>) {
407 let mut skips = self.skips.write().unwrap_or_else(PoisonError::into_inner);
408 if skips.len() >= MAX_PLACEMENT_SKIPS {
409 return;
410 }
411 skips.push(PlacementSkip {
412 node: node.into(),
413 reason: reason.into(),
414 });
415 }
416
417 pub fn skips(&self) -> Vec<PlacementSkip> {
419 self.skips
420 .read()
421 .unwrap_or_else(PoisonError::into_inner)
422 .clone()
423 }
424
425 pub fn clear_skips(&self) {
428 self.skips
429 .write()
430 .unwrap_or_else(PoisonError::into_inner)
431 .clear();
432 }
433
434 pub fn set_plan(&self, nodes: Vec<PlanNodeSummary>) {
436 *self.plan.write().unwrap_or_else(PoisonError::into_inner) = nodes;
437 }
438
439 pub fn set_capacity(&self, capacity: TierCapacity) {
442 *self
443 .capacity
444 .write()
445 .unwrap_or_else(PoisonError::into_inner) = capacity;
446 }
447
448 pub fn capacity(&self) -> TierCapacity {
450 *self.capacity.read().unwrap_or_else(PoisonError::into_inner)
451 }
452
453 pub const fn tiers(&self) -> &TierGauges {
455 &self.tiers
456 }
457
458 pub const fn spill(&self) -> &SpillCounters {
460 &self.spill
461 }
462
463 pub fn snapshot(&self) -> TelemetrySnapshot {
465 let operators = self
466 .operators
467 .read()
468 .unwrap_or_else(PoisonError::into_inner)
469 .iter()
470 .map(|o| o.snapshot())
471 .collect();
472 let plan = self
473 .plan
474 .read()
475 .unwrap_or_else(PoisonError::into_inner)
476 .clone();
477 TelemetrySnapshot {
478 operators,
479 tiers: self.tiers.snapshot(),
480 spill: self.spill.snapshot(),
481 plan,
482 capacity: self.capacity(),
483 }
484 }
485}
486
487#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
489pub struct TelemetrySnapshot {
490 pub operators: Vec<OperatorSnapshot>,
492 pub tiers: TierSnapshot,
494 pub spill: SpillSnapshot,
496 pub plan: Vec<PlanNodeSummary>,
498 pub capacity: TierCapacity,
500}
501
502impl TelemetrySnapshot {
503 pub fn to_prometheus(&self) -> String {
515 let mut out = String::new();
516 let counters: [OperatorCounter; 7] = [
517 ("oxide_operator_rows_in_total", "Input rows.", |o| o.rows_in),
518 ("oxide_operator_rows_out_total", "Output rows.", |o| {
519 o.rows_out
520 }),
521 ("oxide_operator_batches_total", "Batches processed.", |o| {
522 o.batches
523 }),
524 (
525 "oxide_operator_fallback_batches_total",
526 "Batches that ran on the CPU reference although the operator was placed on a device.",
527 |o| o.fallback_batches,
528 ),
529 (
530 "oxide_operator_elapsed_nanoseconds_total",
531 "Processing time.",
532 |o| o.elapsed_ns,
533 ),
534 (
535 "oxide_operator_bytes_h2d_total",
536 "Bytes copied host to device.",
537 |o| o.bytes_h2d,
538 ),
539 (
540 "oxide_operator_bytes_d2h_total",
541 "Bytes copied device to host.",
542 |o| o.bytes_d2h,
543 ),
544 ];
545 for (name, help, value) in counters {
546 out.push_str(&format!("# HELP {name} {help}\n# TYPE {name} counter\n"));
547 for op in &self.operators {
548 out.push_str(&format!(
549 "{name}{{operator=\"{}\",backend=\"{}\"}} {}\n",
550 escape_label(&op.name),
551 op.backend,
552 value(op)
553 ));
554 }
555 }
556 out.push_str("# HELP oxide_operator_memory_bytes Currently allocated bytes.\n");
557 out.push_str("# TYPE oxide_operator_memory_bytes gauge\n");
558 for op in &self.operators {
559 out.push_str(&format!(
560 "oxide_operator_memory_bytes{{operator=\"{}\",backend=\"{}\"}} {}\n",
561 escape_label(&op.name),
562 op.backend,
563 op.memory_bytes
564 ));
565 }
566 let gauges = [
567 ("oxide_tier_device_bytes", self.tiers.device_bytes),
568 ("oxide_tier_host_bytes", self.tiers.host_bytes),
569 ("oxide_tier_disk_bytes", self.tiers.disk_bytes),
570 ];
571 for (name, value) in gauges {
572 out.push_str(&format!(
573 "# HELP {name} Bytes resident in this memory tier.\n# TYPE {name} gauge\n{name} {value}\n"
574 ));
575 }
576 let spill = [
577 ("oxide_spill_demotions_total", self.spill.demotions),
578 ("oxide_spill_promotions_total", self.spill.promotions),
579 ("oxide_spill_spilled_bytes_total", self.spill.spilled_bytes),
580 (
581 "oxide_spill_reloaded_bytes_total",
582 self.spill.reloaded_bytes,
583 ),
584 ];
585 for (name, value) in spill {
586 out.push_str(&format!(
587 "# HELP {name} Spill manager activity.\n# TYPE {name} counter\n{name} {value}\n"
588 ));
589 }
590 out
591 }
592}
593
594type OperatorCounter = (&'static str, &'static str, fn(&OperatorSnapshot) -> u64);
597
598fn escape_label(value: &str) -> String {
600 let mut out = String::with_capacity(value.len());
601 for c in value.chars() {
602 match c {
603 '\\' => out.push_str("\\\\"),
604 '"' => out.push_str("\\\""),
605 '\n' => out.push_str("\\n"),
606 other => out.push(other),
607 }
608 }
609 out
610}
611
612#[cfg(test)]
613mod tests {
614 use super::*;
615
616 #[test]
617 fn a_snapshot_renders_as_prometheus_text() {
618 let hub = TelemetryHub::new();
619 let op = hub.register_operator("GpuFilterExec", BackendKind::Cuda);
620 op.record_batch(1_000, 400, Duration::from_millis(3));
621 op.record_transfer(2_048, 512);
622 op.record_fallback();
623 hub.tiers().set(MemoryTier::Disk, 4_096);
624 let text = hub.snapshot().to_prometheus();
625
626 assert!(
627 text.contains(
628 "oxide_operator_rows_in_total{operator=\"GpuFilterExec\",backend=\"cuda\"} 1000"
629 ),
630 "{text}"
631 );
632 assert!(
633 text.contains("oxide_operator_fallback_batches_total{operator=\"GpuFilterExec\",backend=\"cuda\"} 1"),
634 "{text}"
635 );
636 assert!(text.contains("oxide_tier_disk_bytes 4096"), "{text}");
637 for line in text.lines().filter(|l| !l.starts_with('#')) {
640 let name = line.split(['{', ' ']).next().unwrap_or_default();
641 assert!(
642 text.contains(&format!("# TYPE {name} ")),
643 "{name} has no TYPE"
644 );
645 }
646 }
647
648 #[test]
651 fn label_values_are_escaped() {
652 let hub = TelemetryHub::new();
653 hub.register_operator("we\"ird\\name", BackendKind::CpuSimd);
654 let text = hub.snapshot().to_prometheus();
655 assert!(text.contains("operator=\"we\\\"ird\\\\name\""), "{text}");
656 }
657
658 #[test]
659 fn operators_register_in_order_and_accumulate() {
660 let hub = TelemetryHub::new();
661 let a = hub.register_operator("GpuFilterExec", BackendKind::Cuda);
662 let b = hub.register_operator("GpuAggregateExec", BackendKind::CpuSimd);
663 assert_eq!((a.id(), b.id()), (0, 1));
664 a.record_batch(100, 40, Duration::from_millis(2));
665 a.record_batch(100, 60, Duration::from_millis(4));
666 a.record_transfer(800, 480);
667 a.set_memory_bytes(4096);
668 let snap = hub.snapshot();
669 assert_eq!(snap.operators.len(), 2);
670 let op = &snap.operators[0];
671 assert_eq!((op.rows_in, op.rows_out, op.batches), (200, 100, 2));
672 assert_eq!(
673 (op.bytes_h2d, op.bytes_d2h, op.memory_bytes),
674 (800, 480, 4096)
675 );
676 assert!((op.mean_batch_latency_ms() - 3.0).abs() < 1e-9);
677 assert_eq!(snap.operators[1].mean_batch_latency_ms(), 0.0);
678 }
679
680 #[test]
681 fn tiers_and_spill_counters() {
682 let hub = TelemetryHub::new();
683 hub.tiers().add(MemoryTier::Host, 1000);
684 hub.tiers().sub(MemoryTier::Host, 300);
685 hub.tiers().sub(MemoryTier::Disk, 5); hub.tiers().set(MemoryTier::Device, 42);
687 hub.spill().record_demotion(700);
688 hub.spill().record_promotion(700);
689 let snap = hub.snapshot();
690 assert_eq!(
691 snap.tiers,
692 TierSnapshot {
693 device_bytes: 42,
694 host_bytes: 700,
695 disk_bytes: 0
696 }
697 );
698 assert_eq!(snap.spill.demotions, 1);
699 assert_eq!(snap.spill.promotions, 1);
700 assert_eq!(snap.spill.spilled_bytes, 700);
701 }
702
703 #[test]
704 fn plan_summary_is_replaced() {
705 let hub = TelemetryHub::new();
706 hub.set_plan(vec![PlanNodeSummary {
707 depth: 0,
708 name: "GpuFilterExec".into(),
709 backend: BackendKind::Metal,
710 detail: "a > 1".into(),
711 }]);
712 assert_eq!(hub.snapshot().plan.len(), 1);
713 hub.set_plan(Vec::new());
714 assert!(hub.snapshot().plan.is_empty());
715 }
716}