use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock, PoisonError, RwLock};
use std::time::Duration;
use serde::{Deserialize, Serialize};
use crate::BackendKind;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[non_exhaustive]
pub enum MemoryTier {
Device,
Host,
Disk,
}
impl MemoryTier {
const fn index(self) -> usize {
match self {
MemoryTier::Device => 0,
MemoryTier::Host => 1,
MemoryTier::Disk => 2,
}
}
}
#[derive(Debug)]
pub struct OperatorStats {
id: usize,
name: String,
backend: BackendKind,
rows_in: AtomicU64,
rows_out: AtomicU64,
batches: AtomicU64,
elapsed_ns: AtomicU64,
bytes_h2d: AtomicU64,
bytes_d2h: AtomicU64,
memory_bytes: AtomicU64,
fallback_batches: AtomicU64,
}
impl OperatorStats {
fn new(id: usize, name: String, backend: BackendKind) -> Self {
Self {
id,
name,
backend,
rows_in: AtomicU64::new(0),
rows_out: AtomicU64::new(0),
batches: AtomicU64::new(0),
elapsed_ns: AtomicU64::new(0),
bytes_h2d: AtomicU64::new(0),
bytes_d2h: AtomicU64::new(0),
memory_bytes: AtomicU64::new(0),
fallback_batches: AtomicU64::new(0),
}
}
pub const fn id(&self) -> usize {
self.id
}
pub fn name(&self) -> &str {
&self.name
}
pub const fn backend(&self) -> BackendKind {
self.backend
}
pub fn record_batch(&self, rows_in: u64, rows_out: u64, elapsed: Duration) {
self.rows_in.fetch_add(rows_in, Ordering::Relaxed);
self.rows_out.fetch_add(rows_out, Ordering::Relaxed);
self.batches.fetch_add(1, Ordering::Relaxed);
let ns = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
self.elapsed_ns.fetch_add(ns, Ordering::Relaxed);
}
pub fn record_transfer(&self, h2d: u64, d2h: u64) {
self.bytes_h2d.fetch_add(h2d, Ordering::Relaxed);
self.bytes_d2h.fetch_add(d2h, Ordering::Relaxed);
}
pub fn set_memory_bytes(&self, bytes: u64) {
self.memory_bytes.store(bytes, Ordering::Relaxed);
}
pub fn record_fallback(&self) {
self.fallback_batches.fetch_add(1, Ordering::Relaxed);
}
pub fn fallback_batches(&self) -> u64 {
self.fallback_batches.load(Ordering::Relaxed)
}
pub fn snapshot(&self) -> OperatorSnapshot {
OperatorSnapshot {
id: self.id,
name: self.name.clone(),
backend: self.backend,
rows_in: self.rows_in.load(Ordering::Relaxed),
rows_out: self.rows_out.load(Ordering::Relaxed),
batches: self.batches.load(Ordering::Relaxed),
elapsed_ns: self.elapsed_ns.load(Ordering::Relaxed),
bytes_h2d: self.bytes_h2d.load(Ordering::Relaxed),
bytes_d2h: self.bytes_d2h.load(Ordering::Relaxed),
memory_bytes: self.memory_bytes.load(Ordering::Relaxed),
fallback_batches: self.fallback_batches.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct OperatorSnapshot {
pub id: usize,
pub name: String,
pub backend: BackendKind,
pub rows_in: u64,
pub rows_out: u64,
pub batches: u64,
pub elapsed_ns: u64,
pub bytes_h2d: u64,
pub bytes_d2h: u64,
pub memory_bytes: u64,
pub fallback_batches: u64,
}
impl OperatorSnapshot {
pub fn fell_back_entirely(&self) -> bool {
self.batches > 0 && self.fallback_batches >= self.batches
}
pub fn mean_batch_latency_ms(&self) -> f64 {
if self.batches == 0 {
0.0
} else {
self.elapsed_ns as f64 / self.batches as f64 / 1_000_000.0
}
}
}
#[derive(Debug, Default)]
pub struct TierGauges {
bytes: [AtomicU64; 3],
}
impl TierGauges {
pub fn set(&self, tier: MemoryTier, bytes: u64) {
self.bytes[tier.index()].store(bytes, Ordering::Relaxed);
}
pub fn add(&self, tier: MemoryTier, bytes: u64) {
self.bytes[tier.index()].fetch_add(bytes, Ordering::Relaxed);
}
pub fn sub(&self, tier: MemoryTier, bytes: u64) {
let _ = self.bytes[tier.index()].fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
Some(v.saturating_sub(bytes))
});
}
pub fn get(&self, tier: MemoryTier) -> u64 {
self.bytes[tier.index()].load(Ordering::Relaxed)
}
pub fn snapshot(&self) -> TierSnapshot {
TierSnapshot {
device_bytes: self.get(MemoryTier::Device),
host_bytes: self.get(MemoryTier::Host),
disk_bytes: self.get(MemoryTier::Disk),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct TierSnapshot {
pub device_bytes: u64,
pub host_bytes: u64,
pub disk_bytes: u64,
}
#[derive(Debug, Default)]
pub struct SpillCounters {
demotions: AtomicU64,
promotions: AtomicU64,
spilled_bytes: AtomicU64,
reloaded_bytes: AtomicU64,
}
impl SpillCounters {
pub fn record_demotion(&self, bytes: u64) {
self.demotions.fetch_add(1, Ordering::Relaxed);
self.spilled_bytes.fetch_add(bytes, Ordering::Relaxed);
}
pub fn record_promotion(&self, bytes: u64) {
self.promotions.fetch_add(1, Ordering::Relaxed);
self.reloaded_bytes.fetch_add(bytes, Ordering::Relaxed);
}
pub fn snapshot(&self) -> SpillSnapshot {
SpillSnapshot {
demotions: self.demotions.load(Ordering::Relaxed),
promotions: self.promotions.load(Ordering::Relaxed),
spilled_bytes: self.spilled_bytes.load(Ordering::Relaxed),
reloaded_bytes: self.reloaded_bytes.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct TierCapacity {
pub device_bytes: Option<u64>,
pub host_bytes: Option<u64>,
pub spill_on_query_path: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct SpillSnapshot {
pub demotions: u64,
pub promotions: u64,
pub spilled_bytes: u64,
pub reloaded_bytes: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PlanNodeSummary {
pub depth: usize,
pub name: String,
pub backend: BackendKind,
pub detail: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PlacementSkip {
pub node: String,
pub reason: String,
}
pub const MAX_PLACEMENT_SKIPS: usize = 256;
pub const MAX_OPERATORS: usize = 4096;
#[derive(Debug, Default)]
pub struct TelemetryHub {
operators: RwLock<VecDeque<Arc<OperatorStats>>>,
next_operator_id: AtomicUsize,
tiers: TierGauges,
spill: SpillCounters,
plan: RwLock<Vec<PlanNodeSummary>>,
skips: RwLock<Vec<PlacementSkip>>,
capacity: RwLock<TierCapacity>,
}
static GLOBAL: OnceLock<Arc<TelemetryHub>> = OnceLock::new();
impl TelemetryHub {
pub fn new() -> Arc<Self> {
Arc::new(Self::default())
}
pub fn global() -> &'static Arc<TelemetryHub> {
GLOBAL.get_or_init(TelemetryHub::new)
}
pub fn register_operator(
&self,
name: impl Into<String>,
backend: BackendKind,
) -> Arc<OperatorStats> {
let id = self.next_operator_id.fetch_add(1, Ordering::Relaxed);
let stats = Arc::new(OperatorStats::new(id, name.into(), backend));
let mut ops = self
.operators
.write()
.unwrap_or_else(PoisonError::into_inner);
if ops.len() >= MAX_OPERATORS {
ops.pop_front();
}
ops.push_back(Arc::clone(&stats));
stats
}
pub fn clear_operators(&self) {
self.operators
.write()
.unwrap_or_else(PoisonError::into_inner)
.clear();
}
pub fn record_skip(&self, node: impl Into<String>, reason: impl Into<String>) {
let mut skips = self.skips.write().unwrap_or_else(PoisonError::into_inner);
if skips.len() >= MAX_PLACEMENT_SKIPS {
return;
}
skips.push(PlacementSkip {
node: node.into(),
reason: reason.into(),
});
}
pub fn skips(&self) -> Vec<PlacementSkip> {
self.skips
.read()
.unwrap_or_else(PoisonError::into_inner)
.clone()
}
pub fn clear_skips(&self) {
self.skips
.write()
.unwrap_or_else(PoisonError::into_inner)
.clear();
}
pub fn set_plan(&self, nodes: Vec<PlanNodeSummary>) {
*self.plan.write().unwrap_or_else(PoisonError::into_inner) = nodes;
}
pub fn set_capacity(&self, capacity: TierCapacity) {
*self
.capacity
.write()
.unwrap_or_else(PoisonError::into_inner) = capacity;
}
pub fn capacity(&self) -> TierCapacity {
*self.capacity.read().unwrap_or_else(PoisonError::into_inner)
}
pub const fn tiers(&self) -> &TierGauges {
&self.tiers
}
pub const fn spill(&self) -> &SpillCounters {
&self.spill
}
pub fn snapshot(&self) -> TelemetrySnapshot {
let operators = self
.operators
.read()
.unwrap_or_else(PoisonError::into_inner)
.iter()
.map(|o| o.snapshot())
.collect();
let plan = self
.plan
.read()
.unwrap_or_else(PoisonError::into_inner)
.clone();
TelemetrySnapshot {
operators,
tiers: self.tiers.snapshot(),
spill: self.spill.snapshot(),
plan,
capacity: self.capacity(),
}
}
}
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct TelemetrySnapshot {
pub operators: Vec<OperatorSnapshot>,
pub tiers: TierSnapshot,
pub spill: SpillSnapshot,
pub plan: Vec<PlanNodeSummary>,
pub capacity: TierCapacity,
}
impl TelemetrySnapshot {
pub fn to_prometheus(&self) -> String {
let mut out = String::new();
let counters: [OperatorCounter; 7] = [
("oxide_operator_rows_in_total", "Input rows.", |o| o.rows_in),
("oxide_operator_rows_out_total", "Output rows.", |o| {
o.rows_out
}),
("oxide_operator_batches_total", "Batches processed.", |o| {
o.batches
}),
(
"oxide_operator_fallback_batches_total",
"Batches that ran on the CPU reference although the operator was placed on a device.",
|o| o.fallback_batches,
),
(
"oxide_operator_elapsed_nanoseconds_total",
"Processing time.",
|o| o.elapsed_ns,
),
(
"oxide_operator_bytes_h2d_total",
"Bytes copied host to device.",
|o| o.bytes_h2d,
),
(
"oxide_operator_bytes_d2h_total",
"Bytes copied device to host.",
|o| o.bytes_d2h,
),
];
for (name, help, value) in counters {
out.push_str(&format!("# HELP {name} {help}\n# TYPE {name} counter\n"));
for op in &self.operators {
out.push_str(&format!(
"{name}{{operator=\"{}\",backend=\"{}\"}} {}\n",
escape_label(&op.name),
op.backend,
value(op)
));
}
}
out.push_str("# HELP oxide_operator_memory_bytes Currently allocated bytes.\n");
out.push_str("# TYPE oxide_operator_memory_bytes gauge\n");
for op in &self.operators {
out.push_str(&format!(
"oxide_operator_memory_bytes{{operator=\"{}\",backend=\"{}\"}} {}\n",
escape_label(&op.name),
op.backend,
op.memory_bytes
));
}
let gauges = [
("oxide_tier_device_bytes", self.tiers.device_bytes),
("oxide_tier_host_bytes", self.tiers.host_bytes),
("oxide_tier_disk_bytes", self.tiers.disk_bytes),
];
for (name, value) in gauges {
out.push_str(&format!(
"# HELP {name} Bytes resident in this memory tier.\n# TYPE {name} gauge\n{name} {value}\n"
));
}
let spill = [
("oxide_spill_demotions_total", self.spill.demotions),
("oxide_spill_promotions_total", self.spill.promotions),
("oxide_spill_spilled_bytes_total", self.spill.spilled_bytes),
(
"oxide_spill_reloaded_bytes_total",
self.spill.reloaded_bytes,
),
];
for (name, value) in spill {
out.push_str(&format!(
"# HELP {name} Spill manager activity.\n# TYPE {name} counter\n{name} {value}\n"
));
}
out
}
}
type OperatorCounter = (&'static str, &'static str, fn(&OperatorSnapshot) -> u64);
fn escape_label(value: &str) -> String {
let mut out = String::with_capacity(value.len());
for c in value.chars() {
match c {
'\\' => out.push_str("\\\\"),
'"' => out.push_str("\\\""),
'\n' => out.push_str("\\n"),
other => out.push(other),
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_snapshot_renders_as_prometheus_text() {
let hub = TelemetryHub::new();
let op = hub.register_operator("GpuFilterExec", BackendKind::Cuda);
op.record_batch(1_000, 400, Duration::from_millis(3));
op.record_transfer(2_048, 512);
op.record_fallback();
hub.tiers().set(MemoryTier::Disk, 4_096);
let text = hub.snapshot().to_prometheus();
assert!(
text.contains(
"oxide_operator_rows_in_total{operator=\"GpuFilterExec\",backend=\"cuda\"} 1000"
),
"{text}"
);
assert!(
text.contains("oxide_operator_fallback_batches_total{operator=\"GpuFilterExec\",backend=\"cuda\"} 1"),
"{text}"
);
assert!(text.contains("oxide_tier_disk_bytes 4096"), "{text}");
for line in text.lines().filter(|l| !l.starts_with('#')) {
let name = line.split(['{', ' ']).next().unwrap_or_default();
assert!(
text.contains(&format!("# TYPE {name} ")),
"{name} has no TYPE"
);
}
}
#[test]
fn label_values_are_escaped() {
let hub = TelemetryHub::new();
hub.register_operator("we\"ird\\name", BackendKind::CpuSimd);
let text = hub.snapshot().to_prometheus();
assert!(text.contains("operator=\"we\\\"ird\\\\name\""), "{text}");
}
#[test]
fn operators_register_in_order_and_accumulate() {
let hub = TelemetryHub::new();
let a = hub.register_operator("GpuFilterExec", BackendKind::Cuda);
let b = hub.register_operator("GpuAggregateExec", BackendKind::CpuSimd);
assert_eq!((a.id(), b.id()), (0, 1));
a.record_batch(100, 40, Duration::from_millis(2));
a.record_batch(100, 60, Duration::from_millis(4));
a.record_transfer(800, 480);
a.set_memory_bytes(4096);
let snap = hub.snapshot();
assert_eq!(snap.operators.len(), 2);
let op = &snap.operators[0];
assert_eq!((op.rows_in, op.rows_out, op.batches), (200, 100, 2));
assert_eq!(
(op.bytes_h2d, op.bytes_d2h, op.memory_bytes),
(800, 480, 4096)
);
assert!((op.mean_batch_latency_ms() - 3.0).abs() < 1e-9);
assert_eq!(snap.operators[1].mean_batch_latency_ms(), 0.0);
}
#[test]
fn tiers_and_spill_counters() {
let hub = TelemetryHub::new();
hub.tiers().add(MemoryTier::Host, 1000);
hub.tiers().sub(MemoryTier::Host, 300);
hub.tiers().sub(MemoryTier::Disk, 5); hub.tiers().set(MemoryTier::Device, 42);
hub.spill().record_demotion(700);
hub.spill().record_promotion(700);
let snap = hub.snapshot();
assert_eq!(
snap.tiers,
TierSnapshot {
device_bytes: 42,
host_bytes: 700,
disk_bytes: 0
}
);
assert_eq!(snap.spill.demotions, 1);
assert_eq!(snap.spill.promotions, 1);
assert_eq!(snap.spill.spilled_bytes, 700);
}
#[test]
fn plan_summary_is_replaced() {
let hub = TelemetryHub::new();
hub.set_plan(vec![PlanNodeSummary {
depth: 0,
name: "GpuFilterExec".into(),
backend: BackendKind::Metal,
detail: "a > 1".into(),
}]);
assert_eq!(hub.snapshot().plan.len(), 1);
hub.set_plan(Vec::new());
assert!(hub.snapshot().plan.is_empty());
}
}