Skip to content

Commit

Permalink
Merge pull request #1 from modal-labs/dflemstr/counters
Browse files Browse the repository at this point in the history
Add support for counters
  • Loading branch information
dflemstr authored Oct 29, 2024
2 parents 15eacd6 + 9217962 commit 0a75aed
Show file tree
Hide file tree
Showing 5 changed files with 451 additions and 59 deletions.
223 changes: 169 additions & 54 deletions crates/layer/src/debug_annotations.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,19 @@
//! Internal module used to collect span attrs as Perfetto SDK
//! `DebugAnnotation`s
use std::mem;

use schema::debug_annotation;
use tracing::field;
use tracing_perfetto_sdk_schema as schema;
use tracing_perfetto_sdk_schema::track_event;
use tracing_perfetto_sdk_sys::ffi;

const COUNTER_FIELD_PREFIX: &str = "counter.";

#[derive(Default)]
pub struct FFIDebugAnnotations {
counters: Vec<Counter>,
strings: Vec<ffi::DebugStringAnnotation>,
bools: Vec<ffi::DebugBoolAnnotation>,
ints: Vec<ffi::DebugIntAnnotation>,
Expand All @@ -15,9 +22,23 @@ pub struct FFIDebugAnnotations {

#[derive(Clone, Debug, Default, PartialEq)]
pub struct ProtoDebugAnnotations {
counters: Vec<Counter>,
annotations: Vec<schema::DebugAnnotation>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Counter {
pub name: &'static str,
pub unit: Option<&'static str>,
pub value: CounterValue,
}

#[derive(Clone, Debug, PartialEq)]
pub enum CounterValue {
Float(f64),
Int(i64),
}

impl FFIDebugAnnotations {
pub fn as_ffi(&self) -> ffi::DebugAnnotations {
ffi::DebugAnnotations {
Expand All @@ -31,45 +52,55 @@ impl FFIDebugAnnotations {

impl field::Visit for FFIDebugAnnotations {
fn record_f64(&mut self, field: &field::Field, value: f64) {
self.doubles.push(ffi::DebugDoubleAnnotation {
key: field.name(),
value,
});
if !populate_counter(&mut self.counters, field, value) {
self.doubles.push(ffi::DebugDoubleAnnotation {
key: field.name(),
value,
});
}
}

fn record_i64(&mut self, field: &field::Field, value: i64) {
self.ints.push(ffi::DebugIntAnnotation {
key: field.name(),
value,
});
}

fn record_u64(&mut self, field: &field::Field, value: u64) {
if let Ok(v) = i64::try_from(value) {
if !populate_counter(&mut self.counters, field, value) {
self.ints.push(ffi::DebugIntAnnotation {
key: field.name(),
value: v,
value,
});
} else {
}
}

fn record_u64(&mut self, field: &field::Field, value: u64) {
if !populate_counter(&mut self.counters, field, value) {
if let Ok(v) = i64::try_from(value) {
self.ints.push(ffi::DebugIntAnnotation {
key: field.name(),
value: v,
});
} else {
self.strings.push(ffi::DebugStringAnnotation {
key: field.name(),
value: value.to_string(),
});
}
}
}

fn record_i128(&mut self, field: &field::Field, value: i128) {
if !populate_counter(&mut self.counters, field, value) {
self.strings.push(ffi::DebugStringAnnotation {
key: field.name(),
value: value.to_string(),
});
}
}

fn record_i128(&mut self, field: &field::Field, value: i128) {
self.strings.push(ffi::DebugStringAnnotation {
key: field.name(),
value: value.to_string(),
});
}

fn record_u128(&mut self, field: &field::Field, value: u128) {
self.strings.push(ffi::DebugStringAnnotation {
key: field.name(),
value: value.to_string(),
});
if !populate_counter(&mut self.counters, field, value) {
self.strings.push(ffi::DebugStringAnnotation {
key: field.name(),
value: value.to_string(),
});
}
}

fn record_bool(&mut self, field: &field::Field, value: bool) {
Expand Down Expand Up @@ -106,36 +137,56 @@ impl ProtoDebugAnnotations {
self.annotations
}

pub fn take_counters(&mut self) -> Vec<Counter> {
mem::take(&mut self.counters)
}

fn name_field(field: &field::Field) -> Option<debug_annotation::NameField> {
Some(debug_annotation::NameField::Name(field.name().to_string()))
}
}

impl field::Visit for ProtoDebugAnnotations {
fn record_f64(&mut self, field: &field::Field, value: f64) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::DoubleValue(value)),
..Default::default()
});
if !populate_counter(&mut self.counters, field, value) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::DoubleValue(value)),
..Default::default()
});
}
}

fn record_i64(&mut self, field: &field::Field, value: i64) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::IntValue(value)),
..Default::default()
});
}

fn record_u64(&mut self, field: &field::Field, value: u64) {
if let Ok(v) = i64::try_from(value) {
if !populate_counter(&mut self.counters, field, value) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::IntValue(v)),
value: Some(debug_annotation::Value::IntValue(value)),
..Default::default()
});
} else {
}
}

fn record_u64(&mut self, field: &field::Field, value: u64) {
if !populate_counter(&mut self.counters, field, value) {
if let Ok(v) = i64::try_from(value) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::IntValue(v)),
..Default::default()
});
} else {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::StringValue(value.to_string())),
..Default::default()
});
}
}
}

fn record_i128(&mut self, field: &field::Field, value: i128) {
if !populate_counter(&mut self.counters, field, value) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::StringValue(value.to_string())),
Expand All @@ -144,20 +195,14 @@ impl field::Visit for ProtoDebugAnnotations {
}
}

fn record_i128(&mut self, field: &field::Field, value: i128) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::StringValue(value.to_string())),
..Default::default()
});
}

fn record_u128(&mut self, field: &field::Field, value: u128) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::StringValue(value.to_string())),
..Default::default()
});
if !populate_counter(&mut self.counters, field, value) {
self.annotations.push(schema::DebugAnnotation {
name_field: Self::name_field(field),
value: Some(debug_annotation::Value::StringValue(value.to_string())),
..Default::default()
});
}
}

fn record_bool(&mut self, field: &field::Field, value: bool) {
Expand Down Expand Up @@ -192,3 +237,73 @@ impl field::Visit for ProtoDebugAnnotations {
});
}
}

impl CounterValue {
pub fn to_proto(self) -> track_event::CounterValueField {
match self {
CounterValue::Float(v) => track_event::CounterValueField::DoubleCounterValue(v),
CounterValue::Int(v) => track_event::CounterValueField::CounterValue(v),
}
}
}

impl From<f64> for CounterValue {
fn from(value: f64) -> Self {
CounterValue::Float(value)
}
}

impl From<i64> for CounterValue {
fn from(value: i64) -> Self {
CounterValue::Int(value)
}
}

impl From<u64> for CounterValue {
fn from(value: u64) -> Self {
if let Ok(v) = i64::try_from(value) {
CounterValue::Int(v)
} else {
CounterValue::Float(value as f64)
}
}
}

impl From<i128> for CounterValue {
fn from(value: i128) -> Self {
if let Ok(v) = i64::try_from(value) {
CounterValue::Int(v)
} else {
CounterValue::Float(value as f64)
}
}
}

impl From<u128> for CounterValue {
fn from(value: u128) -> Self {
if let Ok(v) = i64::try_from(value) {
CounterValue::Int(v)
} else {
CounterValue::Float(value as f64)
}
}
}

fn populate_counter(
dest_counters: &mut Vec<Counter>,
field: &field::Field,
value: impl Into<CounterValue>,
) -> bool {
if let Some(name) = field.name().strip_prefix(COUNTER_FIELD_PREFIX) {
let (name, unit) = name
.rsplit_once('.')
.map(|(n, u)| (n, Some(u)))
.unwrap_or_else(|| (name, None));

let value = value.into();
dest_counters.push(Counter { name, value, unit });
true
} else {
false
}
}
7 changes: 7 additions & 0 deletions crates/layer/src/ids.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ const THREAD_NS: u32 = 2;
const TOKIO_NS: u32 = 3;
#[cfg(feature = "tokio")]
const TASK_NS: u32 = 4;
const COUNTER_NS: u32 = 5;

#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[repr(transparent)]
Expand Down Expand Up @@ -44,6 +45,12 @@ impl TrackUuid {
TrackUuid(h.finish())
}

pub fn for_counter(counter_name: &str) -> TrackUuid {
let mut h = hash::DefaultHasher::new();
(TRACK_UUID_NS, COUNTER_NS, counter_name).hash(&mut h);
TrackUuid(h.finish())
}

pub fn as_raw(self) -> u64 {
self.0
}
Expand Down
Loading

0 comments on commit 0a75aed

Please sign in to comment.