relay_cogs/
cogs.rs

1use crate::time::{Clock, SystemClock};
2use std::collections::BTreeMap;
3use std::fmt;
4use std::num::NonZeroUsize;
5use std::sync::Arc;
6
7use crate::{AppFeature, CatgegoryMeasurement, Measurements, ResourceId};
8use crate::{CogsMeasurement, CogsRecorder, Value};
9
10/// COGS measurements collector.
11///
12/// The collector is cheap to clone.
13#[derive(Clone)]
14pub struct Cogs {
15    recorder: Arc<dyn CogsRecorder>,
16}
17
18impl Cogs {
19    /// Creates a new [`Cogs`] from a [`recorder`](CogsRecorder).
20    pub fn new<T>(recorder: T) -> Self
21    where
22        T: CogsRecorder + 'static,
23    {
24        Self {
25            recorder: Arc::new(recorder),
26        }
27    }
28
29    /// Shortcut for creating a [`Cogs`] from a [`crate::NoopRecorder`].
30    ///
31    /// All collected measurements will be dropped.
32    pub fn noop() -> Self {
33        Self::new(crate::NoopRecorder)
34    }
35}
36
37impl Cogs {
38    /// Starts a recording for a COGS measurement, using the [`SystemClock`].
39    ///
40    /// When the returned token is dropped the measurement will be recorded
41    /// with the configured [recorder](CogsRecorder).
42    ///
43    /// The recorded measurement can be attributed to multiple features by supplying a
44    /// weighted [`FeatureWeights`]. A single [`AppFeature`] attributes the entire measurement
45    /// to the feature.
46    ///
47    /// See also: [`Self::timed_with`].
48    ///
49    /// # Example:
50    ///
51    /// ```
52    /// # use relay_cogs::{AppFeature, Cogs, ResourceId};
53    /// # struct Span;
54    /// # fn scrub_sql(_: &mut Span) {}
55    /// # fn extract_tags(_: &mut Span) {};
56    /// #
57    /// fn process_span(cogs: &Cogs, span: &mut Span) {
58    ///     let _token = cogs.timed(ResourceId::Relay, AppFeature::Spans);
59    ///
60    ///     scrub_sql(span);
61    ///     extract_tags(span);
62    /// }
63    ///
64    /// ```
65    pub fn timed<F: Into<FeatureWeights>>(&self, resource: ResourceId, weights: F) -> Token {
66        self.timed_with(resource, weights, SystemClock)
67    }
68
69    /// Starts a recording for a COGS measurement with a specific [`Clock`].
70    ///
71    /// When the returned token is dropped the measurement will be recorded
72    /// with the configured [recorder](CogsRecorder).
73    ///
74    /// The recorded measurement can be attributed to multiple features by supplying a
75    /// weighted [`FeatureWeights`]. A single [`AppFeature`] attributes the entire measurement
76    /// to the feature.
77    ///
78    /// # Example:
79    ///
80    /// ```
81    /// # use relay_cogs::{AppFeature, Cogs, ResourceId};
82    /// # use relay_cogs::time::{Clock, Instant};
83    /// # struct Span;
84    /// # fn scrub_sql(_: &mut Span) {}
85    /// # fn extract_tags(_: &mut Span) {};
86    /// # struct CustomClock;
87    /// # impl Clock for CustomClock { fn now(&self) -> Instant { todo!() } }
88    /// #
89    /// fn process_span(cogs: &Cogs, span: &mut Span) {
90    ///     let _token = cogs.timed_with(ResourceId::Relay, AppFeature::Spans, CustomClock);
91    ///
92    ///     scrub_sql(span);
93    ///     extract_tags(span);
94    /// }
95    ///
96    /// ```
97    pub fn timed_with<F: Into<FeatureWeights>, C: Clock>(
98        &self,
99        resource: ResourceId,
100        weights: F,
101        clock: C,
102    ) -> Token<C> {
103        Token {
104            resource,
105            features: weights.into(),
106            measurements: Measurements::start(clock),
107            recorder: Some(Arc::clone(&self.recorder)),
108        }
109    }
110}
111
112impl fmt::Debug for Cogs {
113    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114        f.debug_struct("Cogs").finish()
115    }
116}
117
118/// An in progress COGS measurement.
119///
120/// The measurement is recorded when the token is dropped.
121#[must_use]
122pub struct Token<C: Clock = SystemClock> {
123    resource: ResourceId,
124    features: FeatureWeights,
125    measurements: Measurements<C>,
126    recorder: Option<Arc<dyn CogsRecorder>>,
127}
128
129impl<C: Clock> Token<C> {
130    /// Creates a new no-op token, which records nothing.
131    ///
132    /// This is primarily useful for testing.
133    pub fn noop() -> Self
134    where
135        C: Default,
136    {
137        Self {
138            resource: ResourceId::Relay,
139            features: FeatureWeights::none(),
140            measurements: Measurements::start(C::default()),
141            recorder: None,
142        }
143    }
144
145    /// Cancels the COGS measurement.
146    pub fn cancel(&mut self) {
147        // No features -> nothing gets attributed.
148        self.update(FeatureWeights::none());
149    }
150
151    /// Starts a categorized measurement.
152    ///
153    /// The measurement is finalized when the returned [`CategoryToken`] is dropped.
154    ///
155    /// Instead of manually starting a categorized measurement, the [`crate::with`]
156    /// macro can be used.
157    pub fn start_category(&mut self, category: impl Category) -> CategoryToken<'_, C> {
158        CategoryToken {
159            measurement: self.measurements.start_category(),
160            parent: self,
161            category: category.name(),
162        }
163    }
164
165    /// Updates the app features to which the active measurement is attributed to.
166    ///
167    /// # Example:
168    ///
169    /// ```
170    /// # use relay_cogs::{AppFeature, Cogs, ResourceId};
171    /// # struct Item;
172    /// # fn do_something(_: &Item) -> bool { true };
173    ///
174    /// fn process(cogs: &Cogs, item: &Item) {
175    ///     let mut token = cogs.timed(ResourceId::Relay, AppFeature::Unattributed);
176    ///
177    ///     // App feature is only known after some computations.
178    ///     if do_something(item) {
179    ///         token.update(AppFeature::Spans);
180    ///     } else {
181    ///         token.update(AppFeature::Transactions);
182    ///     }
183    /// }
184    /// ```
185    pub fn update<T: Into<FeatureWeights>>(&mut self, features: T) {
186        self.features = features.into();
187    }
188}
189
190impl<C: Clock> Drop for Token<C> {
191    fn drop(&mut self) {
192        let Some(recorder) = self.recorder.as_mut() else {
193            return;
194        };
195
196        let measurements = self.measurements.finish();
197
198        for measurement in measurements {
199            for (feature, ratio) in self.features.weights() {
200                let time = measurement.duration.mul_f32(ratio);
201                recorder.record(CogsMeasurement {
202                    resource: self.resource,
203                    feature,
204                    category: measurement.category,
205                    value: Value::Time(time),
206                });
207            }
208        }
209    }
210}
211
212impl fmt::Debug for Token {
213    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
214        f.debug_struct("CogsToken")
215            .field("resource", &self.resource)
216            .field("features", &self.features)
217            .finish()
218    }
219}
220
221/// A COGS category.
222pub trait Category {
223    /// String representation of the category.
224    fn name(&self) -> &'static str;
225}
226
227impl Category for &'static str {
228    fn name(&self) -> &'static str {
229        self
230    }
231}
232
233/// A categorized COGS measurement.
234///
235/// Must be started with [`Token::start_category`].
236#[must_use]
237pub struct CategoryToken<'a, C: Clock> {
238    parent: &'a mut Token<C>,
239    category: &'static str,
240    measurement: CatgegoryMeasurement,
241}
242
243impl<C: Clock> Drop for CategoryToken<'_, C> {
244    fn drop(&mut self) {
245        self.parent
246            .measurements
247            .finish_category(self.measurement, self.category);
248    }
249}
250
251impl<C: Clock> fmt::Debug for CategoryToken<'_, C> {
252    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253        f.debug_struct("CategoryToken")
254            .field("resource", &self.parent.resource)
255            .field("features", &self.parent.features)
256            .field("category", &self.category)
257            .finish()
258    }
259}
260
261/// A collection of weighted [app features](AppFeature).
262///
263/// Used to attribute a single COGS measurement to multiple features.
264#[derive(Clone)]
265pub struct FeatureWeights(BTreeMap<AppFeature, NonZeroUsize>);
266
267impl FeatureWeights {
268    /// Attributes all measurements to a single [`AppFeature`].
269    pub fn new(feature: AppFeature) -> Self {
270        Self::builder().weight(feature, 1).build()
271    }
272
273    /// Attributes the measurement to nothing.
274    pub fn none() -> Self {
275        Self::builder().build()
276    }
277
278    /// Returns an [`FeatureWeights`] builder.
279    pub fn builder() -> FeatureWeightsBuilder {
280        FeatureWeightsBuilder(Self(Default::default()))
281    }
282
283    /// Merges two instances of [`FeatureWeights`] and sums the contained weights.
284    pub fn merge(mut self, other: Self) -> Self {
285        for (feature, weight) in other.0.into_iter() {
286            if let Some(w) = self.0.get_mut(&feature) {
287                *w = w.saturating_add(weight.get());
288            } else {
289                self.0.insert(feature, weight);
290            }
291        }
292
293        self
294    }
295
296    /// Returns an iterator yielding an app feature and it's associated weight.
297    ///
298    /// Weights are normalized to the total stored weights in the range between `0.0` and `1.0`.
299    /// Used to divide a measurement by the stored weights.
300    ///
301    /// # Examples
302    ///
303    /// ```
304    /// use relay_cogs::{AppFeature, FeatureWeights};
305    /// use std::collections::HashMap;
306    ///
307    /// let app_features = FeatureWeights::builder()
308    ///     .weight(AppFeature::Transactions, 1)
309    ///     .weight(AppFeature::Spans, 1)
310    ///     .build();
311    ///
312    /// let weights: HashMap<AppFeature, f32> = app_features.weights().collect();
313    /// assert_eq!(weights, HashMap::from([(AppFeature::Transactions, 0.5), (AppFeature::Spans, 0.5)]))
314    /// ```
315    pub fn weights(&self) -> impl Iterator<Item = (AppFeature, f32)> + '_ {
316        let total_weight: usize = self.0.values().map(|weight| weight.get()).sum();
317
318        self.0.iter().filter_map(move |(feature, weight)| {
319            if total_weight == 0 {
320                return None;
321            }
322
323            let ratio = (weight.get() as f32 / total_weight as f32).clamp(0.0, 1.0);
324            Some((*feature, ratio))
325        })
326    }
327
328    /// Returns `true` if there are no weights contained.
329    ///
330    /// # Examples
331    ///
332    /// ```
333    /// use relay_cogs::{AppFeature, FeatureWeights};
334    ///
335    /// assert!(FeatureWeights::none().is_empty());
336    /// assert!(!FeatureWeights::new(AppFeature::Spans).is_empty());
337    /// ```
338    pub fn is_empty(&self) -> bool {
339        self.0.is_empty()
340    }
341}
342
343impl fmt::Debug for FeatureWeights {
344    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
345        write!(f, "FeatureWeights(")?;
346
347        let mut first = true;
348        for (feature, weight) in self.weights() {
349            if !first {
350                first = false;
351                write!(f, ", ")?;
352            }
353            write!(f, "{feature:?}={weight:.2}")?;
354        }
355        write!(f, ")")
356    }
357}
358
359impl From<AppFeature> for FeatureWeights {
360    fn from(value: AppFeature) -> Self {
361        Self::new(value)
362    }
363}
364
365/// A builder for [`FeatureWeights`] which can be used to configure different weights per [`AppFeature`].
366pub struct FeatureWeightsBuilder(FeatureWeights);
367
368impl FeatureWeightsBuilder {
369    /// Increases the `weight` of an [`AppFeature`].
370    pub fn add_weight(&mut self, feature: AppFeature, weight: usize) -> &mut Self {
371        let Some(weight) = NonZeroUsize::new(weight) else {
372            return self;
373        };
374
375        if let Some(previous) = self.0.0.get_mut(&feature) {
376            *previous = previous.saturating_add(weight.get());
377        } else {
378            self.0.0.insert(feature, weight);
379        }
380
381        self
382    }
383
384    /// Sets the specified `weight` for an [`AppFeature`].
385    pub fn weight(&mut self, feature: AppFeature, weight: usize) -> &mut Self {
386        if let Some(weight) = NonZeroUsize::new(weight) {
387            self.0.0.insert(feature, weight);
388        } else {
389            self.0.0.remove(&feature);
390        }
391        self
392    }
393
394    /// Builds and returns the [`FeatureWeights`].
395    pub fn build(&mut self) -> FeatureWeights {
396        std::mem::replace(self, FeatureWeights::builder()).0
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use std::collections::HashMap;
403    use std::time::Duration;
404
405    use super::*;
406    use crate::test::TestRecorder;
407    use crate::time::StaticClock;
408
409    #[test]
410    fn test_cogs_simple() {
411        let recorder = TestRecorder::default();
412        let cogs = Cogs::new(recorder.clone());
413
414        let start = std::time::Instant::now();
415        drop(cogs.timed(ResourceId::Relay, AppFeature::Spans));
416        let elapsed = start.elapsed();
417
418        let mut measurements = recorder.measurements();
419
420        // Take out the measurement, so we can assert the snapshot and validate the measurement
421        // separately.
422        let Value::Time(measurement) = std::mem::replace(
423            &mut measurements[0].value,
424            Value::Time(Duration::from_nanos(100)),
425        );
426        assert!(measurement > Duration::ZERO);
427        assert!(measurement <= elapsed);
428
429        insta::assert_debug_snapshot!(measurements, @r###"
430        [
431            CogsMeasurement {
432                resource: Relay,
433                feature: Spans,
434                category: None,
435                value: Time(
436                    100ns,
437                ),
438            },
439        ]
440        "###);
441    }
442
443    #[test]
444    fn test_cogs_multiple_weights() {
445        let recorder = TestRecorder::default();
446        let clock = StaticClock::default();
447        let cogs = Cogs::new(recorder.clone());
448
449        let f = FeatureWeights::builder()
450            .weight(AppFeature::Spans, 1)
451            .weight(AppFeature::Transactions, 1)
452            .weight(AppFeature::MetricsSpans, 0) // Noop
453            .add_weight(AppFeature::MetricsSpans, 1)
454            .weight(AppFeature::Transactions, 0) // Reset
455            .build();
456        {
457            let _token = cogs.timed_with(ResourceId::Relay, f, &clock);
458            clock.advance_millis(50);
459        }
460
461        let measurements = recorder.measurements();
462        insta::assert_debug_snapshot!(measurements, @r###"
463        [
464            CogsMeasurement {
465                resource: Relay,
466                feature: Spans,
467                category: None,
468                value: Time(
469                    25ms,
470                ),
471            },
472            CogsMeasurement {
473                resource: Relay,
474                feature: MetricsSpans,
475                category: None,
476                value: Time(
477                    25ms,
478                ),
479            },
480        ]
481        "###);
482    }
483
484    #[test]
485    fn test_cogs_categorized() {
486        let recorder = TestRecorder::default();
487        let clock = StaticClock::default();
488        let cogs = Cogs::new(recorder.clone());
489
490        let features = FeatureWeights::builder()
491            .weight(AppFeature::Spans, 1)
492            .weight(AppFeature::Errors, 1)
493            .build();
494
495        {
496            let mut token = cogs.timed_with(ResourceId::Relay, features, &clock);
497            clock.advance_millis(10);
498            crate::with!(token, "s1", {
499                clock.advance_millis(6);
500            });
501            clock.advance_millis(20);
502            let _category = token.start_category("s2");
503            clock.advance_millis(12);
504        }
505
506        let measurements = recorder.measurements();
507        insta::assert_debug_snapshot!(measurements, @r###"
508        [
509            CogsMeasurement {
510                resource: Relay,
511                feature: Errors,
512                category: None,
513                value: Time(
514                    15ms,
515                ),
516            },
517            CogsMeasurement {
518                resource: Relay,
519                feature: Spans,
520                category: None,
521                value: Time(
522                    15ms,
523                ),
524            },
525            CogsMeasurement {
526                resource: Relay,
527                feature: Errors,
528                category: Some(
529                    "s1",
530                ),
531                value: Time(
532                    3ms,
533                ),
534            },
535            CogsMeasurement {
536                resource: Relay,
537                feature: Spans,
538                category: Some(
539                    "s1",
540                ),
541                value: Time(
542                    3ms,
543                ),
544            },
545            CogsMeasurement {
546                resource: Relay,
547                feature: Errors,
548                category: Some(
549                    "s2",
550                ),
551                value: Time(
552                    6ms,
553                ),
554            },
555            CogsMeasurement {
556                resource: Relay,
557                feature: Spans,
558                category: Some(
559                    "s2",
560                ),
561                value: Time(
562                    6ms,
563                ),
564            },
565        ]
566        "###);
567    }
568
569    #[test]
570    fn test_app_features_none() {
571        let a = FeatureWeights::none();
572        assert_eq!(a.weights().count(), 0);
573    }
574
575    #[test]
576    fn test_app_features_new() {
577        let a = FeatureWeights::new(AppFeature::Spans);
578        assert_eq!(
579            a.weights().collect::<Vec<_>>(),
580            vec![(AppFeature::Spans, 1.0)]
581        );
582    }
583
584    #[test]
585    fn test_app_features_merge() {
586        let a = FeatureWeights::builder()
587            .weight(AppFeature::Spans, 1)
588            .weight(AppFeature::Transactions, 2)
589            .build();
590
591        let b = FeatureWeights::builder()
592            .weight(AppFeature::Spans, 2)
593            .weight(AppFeature::Unattributed, 5)
594            .build();
595
596        let c = FeatureWeights::merge(FeatureWeights::none(), FeatureWeights::merge(a, b));
597
598        let weights: HashMap<_, _> = c.weights().collect();
599        assert_eq!(
600            weights,
601            HashMap::from([
602                (AppFeature::Spans, 0.3),
603                (AppFeature::Transactions, 0.2),
604                (AppFeature::Unattributed, 0.5),
605            ])
606        )
607    }
608}
relay_cogs/cogs.rs

relay_cogs/
cogs.rs
