Skip to content

Commit

Permalink
feat(cogs): Track categories in COGS
Browse files Browse the repository at this point in the history
  • Loading branch information
@Dav1dde
Dav1dde committed Feb 12, 2025
1 parent 860374d commit a4d1fa8
Show file tree
Hide file tree
Showing 8 changed files with 418 additions and 91 deletions.
3 changes: 3 additions & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions relay-cogs/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,3 +12,6 @@ autobenches = false

[lints]
workspace = true

[dev-dependencies]
insta = { workspace = true }
219 changes: 189 additions & 30 deletions relay-cogs/src/cogs.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
use core::fmt;
use crate::time::Instant;
use std::collections::BTreeMap;
use std::fmt;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Instant;

use crate::{AppFeature, ResourceId};
use crate::{AppFeature, Measurements, ResourceId};
use crate::{CogsMeasurement, CogsRecorder, Value};

/// COGS measurements collector.
Expand Down Expand Up @@ -64,7 +64,7 @@ impl Cogs {
Token {
resource,
features: weights.into(),
start: Instant::now(),
measurements: Measurements::start(),
recorder: Arc::clone(&self.recorder),
}
}
Expand All @@ -77,7 +77,7 @@ impl Cogs {
pub struct Token {
resource: ResourceId,
features: FeatureWeights,
start: Instant,
measurements: Measurements,
recorder: Arc<dyn CogsRecorder>,
}

Expand All @@ -88,6 +88,20 @@ impl Token {
self.update(FeatureWeights::none());
}

/// Starts a categorized measurement.
///
/// The measurement is finalized when the returned [`CategoryToken`] is dropped.
///
/// Instead of manually starting a categorized measurement, the [`crate::with`]
/// macro can be used.
pub fn start_category(&mut self, category: impl Category) -> CategoryToken<'_> {
CategoryToken {
parent: self,
start: Instant::now(),
category: category.name(),
}
}

/// Updates the app features to which the active measurement is attributed to.
///
/// # Example:
Expand Down Expand Up @@ -115,15 +129,18 @@ impl Token {

impl Drop for Token {
fn drop(&mut self) {
let elapsed = self.start.elapsed();

for (feature, ratio) in self.features.weights() {
let time = elapsed.mul_f32(ratio);
self.recorder.record(CogsMeasurement {
resource: self.resource,
feature,
value: Value::Time(time),
});
let measurements = self.measurements.finish();

for measurement in measurements {
for (feature, ratio) in self.features.weights() {
let time = measurement.duration.mul_f32(ratio);
self.recorder.record(CogsMeasurement {
resource: self.resource,
feature,
category: measurement.category,
value: Value::Time(time),
});
}
}
}
}
Expand All @@ -137,6 +154,46 @@ impl fmt::Debug for Token {
}
}

/// A COGS category.
pub trait Category {
/// String representation of the category.
fn name(&self) -> &'static str;
}

impl Category for &'static str {
fn name(&self) -> &'static str {
self
}
}

/// A categorized COGS measurement.
///
/// Must be started with [`Token::start_category`].
#[must_use]
pub struct CategoryToken<'a> {
parent: &'a mut Token,
start: Instant,
category: &'static str,
}

impl Drop for CategoryToken<'_> {
fn drop(&mut self) {
self.parent
.measurements
.add(self.start.elapsed(), self.category);
}
}

impl fmt::Debug for CategoryToken<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CategoryToken")
.field("resource", &self.parent.resource)
.field("features", &self.parent.features)
.field("category", &self.category)
.finish()
}
}

/// A collection of weighted [app features](AppFeature).
///
/// Used to attribute a single COGS measurement to multiple features.
Expand Down Expand Up @@ -278,7 +335,7 @@ impl FeatureWeightsBuilder {

#[cfg(test)]
mod tests {
use std::{collections::HashMap, time::Duration};
use std::collections::HashMap;

use super::*;
use crate::test::TestRecorder;
Expand All @@ -291,17 +348,25 @@ mod tests {
drop(cogs.timed(ResourceId::Relay, AppFeature::Spans));

let measurements = recorder.measurements();
assert_eq!(measurements.len(), 1);
assert_eq!(measurements[0].resource, ResourceId::Relay);
assert_eq!(measurements[0].feature, AppFeature::Spans);
insta::assert_debug_snapshot!(measurements, @r###"
[
CogsMeasurement {
resource: Relay,
feature: Spans,
category: None,
value: Time(
100ns,
),
},
]
"###);
}

#[test]
fn test_cogs_multiple_weights() {
let recorder = TestRecorder::default();
let cogs = Cogs::new(recorder.clone());

let start = Instant::now();
let f = FeatureWeights::builder()
.weight(AppFeature::Spans, 1)
.weight(AppFeature::Transactions, 1)
Expand All @@ -311,20 +376,114 @@ mod tests {
.build();
{
let _token = cogs.timed(ResourceId::Relay, f);
std::thread::sleep(Duration::from_millis(50));
crate::time::advance_millis(50);
}

let measurements = recorder.measurements();
insta::assert_debug_snapshot!(measurements, @r###"
[
CogsMeasurement {
resource: Relay,
feature: Spans,
category: None,
value: Time(
25ms,
),
},
CogsMeasurement {
resource: Relay,
feature: MetricsSpans,
category: None,
value: Time(
25ms,
),
},
]
"###);
}

#[test]
fn test_cogs_categorized() {
let recorder = TestRecorder::default();
let cogs = Cogs::new(recorder.clone());

let features = FeatureWeights::builder()
.weight(AppFeature::Spans, 1)
.weight(AppFeature::Errors, 1)
.build();

{
let mut token = cogs.timed(ResourceId::Relay, features);
crate::time::advance_millis(10);
crate::with!(token, "s1", {
crate::time::advance_millis(6);
});
crate::time::advance_millis(20);
let _category = token.start_category("s2");
crate::time::advance_millis(12);
}
let elapsed = start.elapsed();

let measurements = recorder.measurements();
assert_eq!(measurements.len(), 2);
assert_eq!(measurements[0].resource, ResourceId::Relay);
assert_eq!(measurements[0].feature, AppFeature::Spans);
assert_eq!(measurements[1].resource, ResourceId::Relay);
assert_eq!(measurements[1].feature, AppFeature::MetricsSpans);
assert_eq!(measurements[0].value, measurements[1].value);
let Value::Time(time) = measurements[0].value;
assert!(time >= Duration::from_millis(25), "{time:?}");
assert!(time <= elapsed.div_f32(1.99), "{time:?}");
insta::assert_debug_snapshot!(measurements, @r###"
[
CogsMeasurement {
resource: Relay,
feature: Errors,
category: None,
value: Time(
15ms,
),
},
CogsMeasurement {
resource: Relay,
feature: Spans,
category: None,
value: Time(
15ms,
),
},
CogsMeasurement {
resource: Relay,
feature: Errors,
category: Some(
"s1",
),
value: Time(
3ms,
),
},
CogsMeasurement {
resource: Relay,
feature: Spans,
category: Some(
"s1",
),
value: Time(
3ms,
),
},
CogsMeasurement {
resource: Relay,
feature: Errors,
category: Some(
"s2",
),
value: Time(
6ms,
),
},
CogsMeasurement {
resource: Relay,
feature: Spans,
category: Some(
"s2",
),
value: Time(
6ms,
),
},
]
"###);
}

#[test]
Expand Down
52 changes: 46 additions & 6 deletions relay-cogs/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,17 +63,53 @@
html_favicon_url = "https://raw.githubusercontent.com/getsentry/relay/master/artwork/relay-icon.png"
)]

use std::time::Duration;

mod cogs;
mod measurement;
mod recorder;
#[cfg(test)]
mod test;

pub use cogs::*;
pub use recorder::*;
pub(crate) mod time;

pub use self::cogs::*;
pub use self::recorder::*;
#[cfg(test)]
pub use test::*;
pub use self::test::*;

pub(crate) use self::measurement::*;

/// Records a categorized measurement of the passed `body`, in `category` on `token`.
///
/// # Example:
///
/// ```
/// # use relay_cogs::{AppFeature, Cogs, ResourceId};
/// # struct Item;
/// # fn do_something(_: &Item) -> bool { true };
/// # fn do_something_else(_: &Item) -> bool { true };
///
/// fn process(cogs: &Cogs, item: &Item) {
/// let mut token = cogs.timed(ResourceId::Relay, AppFeature::Transactions);
///
/// // The entire body is categorized as `processing`.
/// relay_cogs::with!(token, "processing", {
/// let success = do_something(&item);
/// });
///
/// // Not categorized.
/// if success {
/// do_something_else(&item);
/// }
/// }
/// ```
#[macro_export]
macro_rules! with {
($token:expr, $category:expr, { $($body:tt)* }) => {
let token = $token.start_category($category);
$($body)*
drop(token);
};
}

/// Resource ID as tracked in COGS.
///
Expand Down Expand Up @@ -186,6 +222,10 @@ pub struct CogsMeasurement {
pub resource: ResourceId,
/// The measured app feature.
pub feature: AppFeature,
/// Optional category for this measurement.
///
/// A category further subdivides a measurement for a specific feature.
pub category: Option<&'static str>,
/// The measurement value.
pub value: Value,
}
Expand All @@ -194,5 +234,5 @@ pub struct CogsMeasurement {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Value {
/// A time measurement.
Time(Duration),
Time(std::time::Duration),
}
Loading

0 comments on commit a4d1fa8

Please sign in to comment.