Skip to main content

libdd_trace_stats/span_concentrator/
aggregation.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4//! This module implement the logic for stats aggregation into time buckets and stats group.
5//! This includes the aggregation key to group spans together and the computation of stats from a
6//! span.
7
8use hashbrown::HashMap;
9use libdd_trace_obfuscation::ip_address::quantize_peer_ip_addresses;
10use libdd_trace_protobuf::pb;
11use libdd_trace_utils::span::SpanText;
12use std::borrow::{Borrow, Cow};
13
14use crate::span_concentrator::StatSpan;
15
16const TAG_STATUS_CODE: &str = "http.status_code";
17const TAG_SYNTHETICS: &str = "synthetics";
18const TAG_SPANKIND: &str = "span.kind";
19const TAG_ORIGIN: &str = "_dd.origin";
20const TAG_SVC_SRC: &str = "_dd.svc_src";
21const GRPC_STATUS_CODE_FIELD: &[&str] = &[
22    "rpc.grpc.status_code",
23    "grpc.code",
24    "rpc.grpc.status.code",
25    "grpc.status.code",
26];
27
28/// Aggregation key fields shared across all concentrator implementations — everything
29/// **except** peer tags.
30///
31/// `T` is the string representation:
32/// * `&'a str`   — borrowed references used in [`BorrowedAggregationKey`]
33/// * `String`    — owned values used in `OwnedAggregationKey`
34/// * `StringRef` — offset+len into a SHM string pool, used in `ShmKeyHeader`
35#[derive(
36    Clone, Default, Hash, Eq, PartialEq, Debug, PartialOrd, serde::Serialize, serde::Deserialize,
37)]
38pub struct FixedAggregationKey<T> {
39    pub resource_name: T,
40    pub service_name: T,
41    pub operation_name: T,
42    pub span_type: T,
43    pub span_kind: T,
44    pub http_method: T,
45    pub http_endpoint: T,
46    pub service_source: T,
47    pub http_status_code: u32,
48    pub grpc_status_code: Option<u8>,
49    pub is_synthetics_request: bool,
50    pub is_trace_root: bool,
51}
52
53impl<T> FixedAggregationKey<T> {
54    /// Map all string fields through `f`, preserving scalar fields.
55    pub fn convert<'a, V: 'a, I: ?Sized + 'a, F: Fn(&'a I) -> V>(
56        &'a self,
57        f: F,
58    ) -> FixedAggregationKey<V>
59    where
60        T: Borrow<I>,
61    {
62        FixedAggregationKey {
63            resource_name: f(self.resource_name.borrow()),
64            service_name: f(self.service_name.borrow()),
65            operation_name: f(self.operation_name.borrow()),
66            span_type: f(self.span_type.borrow()),
67            span_kind: f(self.span_kind.borrow()),
68            http_method: f(self.http_method.borrow()),
69            http_endpoint: f(self.http_endpoint.borrow()),
70            service_source: f(self.service_source.borrow()),
71            http_status_code: self.http_status_code,
72            grpc_status_code: self.grpc_status_code,
73            is_synthetics_request: self.is_synthetics_request,
74            is_trace_root: self.is_trace_root,
75        }
76    }
77}
78
79#[derive(Clone, Hash, PartialEq, Eq)]
80/// Represent a stats aggregation key borrowed from span data
81pub(super) struct BorrowedAggregationKey<'a> {
82    fixed: FixedAggregationKey<&'a str>,
83    peer_tags: Vec<(&'a str, Cow<'a, str>)>,
84}
85
86impl hashbrown::Equivalent<OwnedAggregationKey> for BorrowedAggregationKey<'_> {
87    #[inline]
88    fn equivalent(&self, other: &OwnedAggregationKey) -> bool {
89        self.fixed == other.fixed.convert(|s| s)
90            && self.peer_tags.len() == other.peer_tags.len()
91            && self
92                .peer_tags
93                .iter()
94                .zip(other.peer_tags.iter())
95                .all(|((k1, v1), (k2, v2))| k1 == k2 && v1 == v2)
96    }
97}
98
99#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Default)]
100/// Represents a span aggregation key with owned data
101///
102/// To be able to use BorrowedAggregationKey to index into a stats bucket hashmap two
103/// conditions must stay true:
104/// * Hashing an owned key derived from a borrowed key should produce the same hash as hashing the
105///   borrowed key
106/// * Running the Equivalent trait on an owned key derived from a borrowed key should produce true
107pub(super) struct OwnedAggregationKey {
108    fixed: FixedAggregationKey<String>,
109    peer_tags: Vec<(String, String)>,
110}
111
112impl From<&BorrowedAggregationKey<'_>> for OwnedAggregationKey {
113    fn from(value: &BorrowedAggregationKey<'_>) -> Self {
114        OwnedAggregationKey {
115            fixed: value.fixed.convert(str::to_owned),
116            peer_tags: value
117                .peer_tags
118                .iter()
119                .map(|(k, v)| (k.to_string(), v.to_string()))
120                .collect(),
121        }
122    }
123}
124
125fn float_to_int(f: f64) -> Option<u8> {
126    if f.floor() != f {
127        return None;
128    }
129    if f < 0.0 || (u8::MAX as f64) < f {
130        return None;
131    }
132    Some(f as u8)
133}
134
135fn get_grpc_status_code<'a>(span: &'a impl StatSpan<'a>) -> Option<u8> {
136    for key in GRPC_STATUS_CODE_FIELD {
137        if let Some(val) = span.get_meta(key) {
138            if let Some(code) = grpc_status_str_to_int_value(val) {
139                return Some(code);
140            }
141        }
142    }
143
144    for key in GRPC_STATUS_CODE_FIELD {
145        if let Some(val) = span.get_metrics(key) {
146            if let Some(code) = float_to_int(val) {
147                return Some(code);
148            }
149        }
150    }
151
152    None
153}
154
155fn grpc_status_str_to_int_value(v: &str) -> Option<u8> {
156    if let Ok(status) = v.parse() {
157        return Some(status);
158    }
159    let mut status_uppercase = [0u8; 32];
160    let mut status = v.trim_start_matches("StatusCode.");
161
162    let mut needs_upcasing = false;
163    for b in status.as_bytes() {
164        if !b.is_ascii() {
165            return None;
166        }
167        needs_upcasing |= b.is_ascii_lowercase()
168    }
169    if needs_upcasing {
170        for (c, d) in status.as_bytes().iter().zip(&mut status_uppercase) {
171            *d = c.to_ascii_uppercase();
172        }
173        status = std::str::from_utf8(&status_uppercase[0..status.len().min(status_uppercase.len())])
174            .ok()?
175    }
176
177    match status {
178        "OK" => return Some(0),
179        "CANCELLED" | "CANCELED" => return Some(1),
180        "UNKNOWN" => return Some(2),
181        "INVALID_ARGUMENT" | "INVALIDARGUMENT" => return Some(3),
182        "DEADLINE_EXCEEDED" | "DEADLINEEXCEEDED" => return Some(4),
183        "NOT_FOUND" | "NOTFOUND" => return Some(5),
184        "ALREADY_EXISTS" | "ALREADYEXISTS" => return Some(6),
185        "PERMISSION_DENIED" | "PERMISSIONDENIED" => return Some(7),
186        "UNAUTHENTICATED" => return Some(16),
187        "RESOURCE_EXHAUSTED" | "RESOURCEEXHAUSTED" => return Some(8),
188        "FAILED_PRECONDITION" | "FAILEDPRECONDITION" => return Some(9),
189        "ABORTED" => return Some(10),
190        "OUT_OF_RANGE" | "OUTOFRANGE" => return Some(11),
191        "UNIMPLEMENTED" => return Some(12),
192        "INTERNAL" => return Some(13),
193        "UNAVAILABLE" => return Some(14),
194        "DATA_LOSS" | "DATALOSS" => return Some(15),
195        _ => {}
196    }
197    None
198}
199
200impl<'a> BorrowedAggregationKey<'a> {
201    /// Return an AggregationKey matching the given span.
202    ///
203    /// If `peer_tags_keys` is not empty then the peer tags of the span will be included in the
204    /// key.
205    pub(super) fn from_span<T: StatSpan<'a>>(span: &'a T, peer_tag_keys: &'a [String]) -> Self {
206        Self::from_obfuscated_span(span.resource(), span, peer_tag_keys)
207    }
208
209    pub(crate) fn from_obfuscated_span<'b, T>(
210        resource_name: &'a str,
211        span: &'b T,
212        peer_tag_keys: &'b [String],
213    ) -> BorrowedAggregationKey<'a>
214    where
215        T: StatSpan<'b>,
216        // resource_name is a temporary string on the stack the span will outlive it
217        'b: 'a,
218    {
219        let span_kind = span.get_meta(TAG_SPANKIND).unwrap_or_default();
220        let peer_tags = if should_track_peer_tags(span_kind) {
221            // Parse the meta tags of the span and return a list of the peer tags based on the list
222            // of `peer_tag_keys`. IP address values are quantized to reduce cardinality.
223            peer_tag_keys
224                .iter()
225                .filter_map(|key| {
226                    let value = span.get_meta(key.as_str())?;
227                    Some((key.as_str(), quantize_peer_ip_addresses(value)))
228                })
229                .collect()
230        } else if let Some(base_service) = span.get_meta("_dd.base_service") {
231            // Internal spans with a base service override use only _dd.base_service as peer tag
232            vec![("_dd.base_service", Cow::Borrowed(base_service))]
233        } else {
234            vec![]
235        };
236
237        let http_method = span.get_meta("http.method").unwrap_or_default();
238
239        let http_endpoint = span
240            .get_meta("http.endpoint")
241            .or_else(|| span.get_meta("http.route"))
242            .unwrap_or_default();
243
244        let status_code = if let Some(status_code) = span.get_metrics(TAG_STATUS_CODE) {
245            status_code as u32
246        } else if let Some(status_code) = span.get_meta(TAG_STATUS_CODE) {
247            status_code.parse().unwrap_or_default()
248        } else {
249            0
250        };
251
252        let grpc_status_code = get_grpc_status_code(span);
253
254        let service_source = span.get_meta(TAG_SVC_SRC).unwrap_or_default();
255
256        Self {
257            fixed: FixedAggregationKey {
258                resource_name,
259                service_name: span.service(),
260                operation_name: span.name(),
261                span_type: span.r#type(),
262                span_kind,
263                http_method,
264                http_endpoint,
265                service_source,
266                http_status_code: status_code,
267                grpc_status_code,
268                is_synthetics_request: span
269                    .get_meta(TAG_ORIGIN)
270                    .is_some_and(|origin| origin.starts_with(TAG_SYNTHETICS)),
271                is_trace_root: span.is_trace_root(),
272            },
273            peer_tags,
274        }
275    }
276}
277
278impl From<pb::ClientGroupedStats> for OwnedAggregationKey {
279    fn from(value: pb::ClientGroupedStats) -> Self {
280        Self {
281            fixed: FixedAggregationKey {
282                resource_name: value.resource,
283                service_name: value.service,
284                operation_name: value.name,
285                span_type: value.r#type,
286                span_kind: value.span_kind,
287                http_method: value.http_method,
288                http_endpoint: value.http_endpoint,
289                service_source: value.service_source,
290                http_status_code: value.http_status_code,
291                grpc_status_code: value.grpc_status_code.parse().ok(),
292                is_synthetics_request: value.synthetics,
293                is_trace_root: value.is_trace_root == 1,
294            },
295            peer_tags: value
296                .peer_tags
297                .into_iter()
298                .filter_map(|t| {
299                    let (key, value) = t.split_once(':')?;
300                    Some((key.to_string(), value.to_string()))
301                })
302                .collect(),
303        }
304    }
305}
306
307/// Return true if we care about peer tags on the span
308fn should_track_peer_tags<T>(span_kind: T) -> bool
309where
310    T: SpanText,
311{
312    matches!(
313        span_kind.borrow().to_lowercase().as_str(),
314        "client" | "producer" | "consumer"
315    )
316}
317
318/// The stats computed from a group of span with the same AggregationKey
319#[derive(Debug, Default, Clone)]
320pub(super) struct GroupedStats {
321    hits: u64,
322    errors: u64,
323    duration: u64,
324    top_level_hits: u64,
325    ok_summary: libdd_ddsketch::DDSketch,
326    error_summary: libdd_ddsketch::DDSketch,
327}
328
329impl GroupedStats {
330    /// Update the stats of a GroupedStats by inserting a span.
331    fn insert(&mut self, duration: i64, is_error: bool, is_top_level: bool) {
332        self.hits += 1;
333        self.duration += duration as u64;
334
335        if is_error {
336            self.errors += 1;
337            let _ = self.error_summary.add(duration as f64);
338        } else {
339            let _ = self.ok_summary.add(duration as f64);
340        }
341        if is_top_level {
342            self.top_level_hits += 1;
343        }
344    }
345}
346
347/// A time bucket used for stats aggregation. It stores a map of GroupedStats storing the stats of
348/// spans aggregated on their AggregationKey.
349#[derive(Debug, Clone)]
350pub(super) struct StatsBucket {
351    data: HashMap<OwnedAggregationKey, GroupedStats>,
352    start: u64,
353}
354
355impl StatsBucket {
356    /// Return a new StatsBucket starting at the given timestamp
357    pub(super) fn new(start_timestamp: u64) -> Self {
358        Self {
359            data: HashMap::new(),
360            start: start_timestamp,
361        }
362    }
363
364    /// Insert a value as stats in the group corresponding to the aggregation key, if it does
365    /// not exist it creates it.
366    pub(super) fn insert(
367        &mut self,
368        key: BorrowedAggregationKey<'_>,
369        duration: i64,
370        is_error: bool,
371        is_top_level: bool,
372    ) {
373        self.data
374            .entry_ref(&key)
375            .or_default()
376            .insert(duration, is_error, is_top_level);
377    }
378
379    /// Consume the bucket and return a ClientStatsBucket containing the bucket stats.
380    /// `bucket_duration` is the size of buckets for the concentrator containing the bucket.
381    pub(super) fn flush(self, bucket_duration: u64) -> pb::ClientStatsBucket {
382        pb::ClientStatsBucket {
383            start: self.start,
384            duration: bucket_duration,
385            stats: self
386                .data
387                .into_iter()
388                .map(|(k, b)| encode_grouped_stats(k, b))
389                .collect(),
390            // Agent-only field
391            agent_time_shift: 0,
392        }
393    }
394}
395
396/// Create a ClientGroupedStats struct based on the given AggregationKey and GroupedStats
397fn encode_grouped_stats(key: OwnedAggregationKey, group: GroupedStats) -> pb::ClientGroupedStats {
398    let f = key.fixed;
399    pb::ClientGroupedStats {
400        service: f.service_name,
401        name: f.operation_name,
402        resource: f.resource_name,
403        http_status_code: f.http_status_code,
404        r#type: f.span_type,
405        db_type: String::new(), // db_type is not used yet (see proto definition)
406
407        hits: group.hits,
408        errors: group.errors,
409        duration: group.duration,
410
411        ok_summary: group.ok_summary.encode_to_vec(),
412        error_summary: group.error_summary.encode_to_vec(),
413        synthetics: f.is_synthetics_request,
414        top_level_hits: group.top_level_hits,
415        span_kind: f.span_kind,
416
417        peer_tags: key
418            .peer_tags
419            .into_iter()
420            .map(|(k, v)| format!("{k}:{v}"))
421            .collect(),
422        is_trace_root: if f.is_trace_root {
423            pb::Trilean::True.into()
424        } else {
425            pb::Trilean::False.into()
426        },
427        http_method: f.http_method,
428        http_endpoint: f.http_endpoint,
429        grpc_status_code: f
430            .grpc_status_code
431            .map(|c| c.to_string())
432            .unwrap_or_default(),
433        service_source: f.service_source,
434        span_derived_primary_tags: vec![], // Todo
435    }
436}
437
438#[cfg(test)]
439mod tests {
440    use libdd_trace_utils::span::v04::{SpanBytes, SpanSlice};
441
442    use super::*;
443    use std::{collections::HashMap, hash::Hash};
444
445    fn get_hash(v: &impl Hash) -> u64 {
446        use std::hash::Hasher;
447        let mut hasher = std::hash::DefaultHasher::new();
448        v.hash(&mut hasher);
449        hasher.finish()
450    }
451
452    impl FixedAggregationKey<String> {
453        fn into_key(self) -> OwnedAggregationKey {
454            OwnedAggregationKey {
455                fixed: self,
456                peer_tags: vec![],
457            }
458        }
459        fn into_key_with_peers(self, peer_tags: Vec<(String, String)>) -> OwnedAggregationKey {
460            OwnedAggregationKey {
461                fixed: self,
462                peer_tags,
463            }
464        }
465    }
466
467    #[test]
468    fn test_aggregation_key_from_span() {
469        let test_cases: Vec<(SpanBytes, OwnedAggregationKey)> = vec![
470            // Root span
471            (
472                SpanBytes {
473                    service: "service".into(),
474                    name: "op".into(),
475                    resource: "res".into(),
476                    span_id: 1,
477                    parent_id: 0,
478                    ..Default::default()
479                },
480                FixedAggregationKey {
481                    service_name: "service".into(),
482                    operation_name: "op".into(),
483                    resource_name: "res".into(),
484                    is_trace_root: true,
485                    ..Default::default()
486                }
487                .into_key(),
488            ),
489            // Span with span kind
490            (
491                SpanBytes {
492                    service: "service".into(),
493                    name: "op".into(),
494                    resource: "res".into(),
495                    span_id: 1,
496                    parent_id: 0,
497                    meta: HashMap::from([("span.kind".into(), "client".into())]),
498                    ..Default::default()
499                },
500                FixedAggregationKey {
501                    service_name: "service".into(),
502                    operation_name: "op".into(),
503                    resource_name: "res".into(),
504                    span_kind: "client".into(),
505                    is_trace_root: true,
506                    ..Default::default()
507                }
508                .into_key(),
509            ),
510            // Span with peer tags but peertags aggregation disabled
511            (
512                SpanBytes {
513                    service: "service".into(),
514                    name: "op".into(),
515                    resource: "res".into(),
516                    span_id: 1,
517                    parent_id: 0,
518                    meta: HashMap::from([
519                        ("span.kind".into(), "client".into()),
520                        ("aws.s3.bucket".into(), "bucket-a".into()),
521                    ]),
522                    ..Default::default()
523                },
524                FixedAggregationKey {
525                    service_name: "service".into(),
526                    operation_name: "op".into(),
527                    resource_name: "res".into(),
528                    span_kind: "client".into(),
529                    is_trace_root: true,
530                    ..Default::default()
531                }
532                .into_key(),
533            ),
534            // Span with multiple peer tags but peertags aggregation disabled
535            (
536                SpanBytes {
537                    service: "service".into(),
538                    name: "op".into(),
539                    resource: "res".into(),
540                    span_id: 1,
541                    parent_id: 0,
542                    meta: HashMap::from([
543                        ("span.kind".into(), "producer".into()),
544                        ("aws.s3.bucket".into(), "bucket-a".into()),
545                        ("db.instance".into(), "dynamo.test.us1".into()),
546                        ("db.system".into(), "dynamodb".into()),
547                    ]),
548                    ..Default::default()
549                },
550                FixedAggregationKey {
551                    service_name: "service".into(),
552                    operation_name: "op".into(),
553                    resource_name: "res".into(),
554                    span_kind: "producer".into(),
555                    is_trace_root: true,
556                    ..Default::default()
557                }
558                .into_key(),
559            ),
560            // Span with multiple peer tags but peertags aggregation disabled and span kind is
561            // server
562            (
563                SpanBytes {
564                    service: "service".into(),
565                    name: "op".into(),
566                    resource: "res".into(),
567                    span_id: 1,
568                    parent_id: 0,
569                    meta: HashMap::from([
570                        ("span.kind".into(), "server".into()),
571                        ("aws.s3.bucket".into(), "bucket-a".into()),
572                        ("db.instance".into(), "dynamo.test.us1".into()),
573                        ("db.system".into(), "dynamodb".into()),
574                    ]),
575                    ..Default::default()
576                },
577                FixedAggregationKey {
578                    service_name: "service".into(),
579                    operation_name: "op".into(),
580                    resource_name: "res".into(),
581                    span_kind: "server".into(),
582                    is_trace_root: true,
583                    ..Default::default()
584                }
585                .into_key(),
586            ),
587            // Span from synthetics
588            (
589                SpanBytes {
590                    service: "service".into(),
591                    name: "op".into(),
592                    resource: "res".into(),
593                    span_id: 1,
594                    parent_id: 0,
595                    meta: HashMap::from([("_dd.origin".into(), "synthetics-browser".into())]),
596                    ..Default::default()
597                },
598                FixedAggregationKey {
599                    service_name: "service".into(),
600                    operation_name: "op".into(),
601                    resource_name: "res".into(),
602                    is_synthetics_request: true,
603                    is_trace_root: true,
604                    ..Default::default()
605                }
606                .into_key(),
607            ),
608            // Span with status code in meta
609            (
610                SpanBytes {
611                    service: "service".into(),
612                    name: "op".into(),
613                    resource: "res".into(),
614                    span_id: 1,
615                    parent_id: 0,
616                    meta: HashMap::from([("http.status_code".into(), "418".into())]),
617                    ..Default::default()
618                },
619                FixedAggregationKey {
620                    service_name: "service".into(),
621                    operation_name: "op".into(),
622                    resource_name: "res".into(),
623                    is_synthetics_request: false,
624                    is_trace_root: true,
625                    http_status_code: 418,
626                    ..Default::default()
627                }
628                .into_key(),
629            ),
630            // Span with invalid status code in meta
631            (
632                SpanBytes {
633                    service: "service".into(),
634                    name: "op".into(),
635                    resource: "res".into(),
636                    span_id: 1,
637                    parent_id: 0,
638                    meta: HashMap::from([("http.status_code".into(), "x".into())]),
639                    ..Default::default()
640                },
641                FixedAggregationKey {
642                    service_name: "service".into(),
643                    operation_name: "op".into(),
644                    resource_name: "res".into(),
645                    is_synthetics_request: false,
646                    is_trace_root: true,
647                    ..Default::default()
648                }
649                .into_key(),
650            ),
651            // Span with status code in metrics
652            (
653                SpanBytes {
654                    service: "service".into(),
655                    name: "op".into(),
656                    resource: "res".into(),
657                    span_id: 1,
658                    parent_id: 0,
659                    metrics: HashMap::from([("http.status_code".into(), 418.0)]),
660                    ..Default::default()
661                },
662                FixedAggregationKey {
663                    service_name: "service".into(),
664                    operation_name: "op".into(),
665                    resource_name: "res".into(),
666                    is_synthetics_request: false,
667                    is_trace_root: true,
668                    http_status_code: 418,
669                    ..Default::default()
670                }
671                .into_key(),
672            ),
673            // Span with http.method and http.route
674            (
675                SpanBytes {
676                    service: "service".into(),
677                    name: "op".into(),
678                    resource: "GET /api/v1/users".into(),
679                    span_id: 1,
680                    parent_id: 0,
681                    meta: HashMap::from([
682                        ("http.method".into(), "GET".into()),
683                        ("http.route".into(), "/api/v1/users".into()),
684                    ]),
685                    ..Default::default()
686                },
687                FixedAggregationKey {
688                    service_name: "service".into(),
689                    operation_name: "op".into(),
690                    resource_name: "GET /api/v1/users".into(),
691                    http_method: "GET".into(),
692                    http_endpoint: "/api/v1/users".into(),
693                    is_synthetics_request: false,
694                    is_trace_root: true,
695                    ..Default::default()
696                }
697                .into_key(),
698            ),
699            // Span with http.method and http.endpoint (http.endpoint takes precedence)
700            (
701                SpanBytes {
702                    service: "service".into(),
703                    name: "op".into(),
704                    resource: "POST /users/create".into(),
705                    span_id: 1,
706                    parent_id: 0,
707                    meta: HashMap::from([
708                        ("http.method".into(), "POST".into()),
709                        ("http.route".into(), "/users/create".into()),
710                        ("http.endpoint".into(), "/users/create2".into()),
711                    ]),
712                    ..Default::default()
713                },
714                FixedAggregationKey {
715                    service_name: "service".into(),
716                    operation_name: "op".into(),
717                    resource_name: "POST /users/create".into(),
718                    http_method: "POST".into(),
719                    http_endpoint: "/users/create2".into(),
720                    is_synthetics_request: false,
721                    is_trace_root: true,
722                    ..Default::default()
723                }
724                .into_key(),
725            ),
726            // Span with grpc status from meta as named string
727            (
728                SpanBytes {
729                    meta: HashMap::from([("rpc.grpc.status_code".into(), "OK".into())]),
730                    ..Default::default()
731                },
732                FixedAggregationKey {
733                    grpc_status_code: Some(0),
734                    is_trace_root: true,
735                    ..Default::default()
736                }
737                .into_key(),
738            ),
739            // Span with grpc status from meta as numeric string
740            (
741                SpanBytes {
742                    meta: HashMap::from([("rpc.grpc.status_code".into(), "14".into())]),
743                    ..Default::default()
744                },
745                FixedAggregationKey {
746                    grpc_status_code: Some(14),
747                    is_trace_root: true,
748                    ..Default::default()
749                }
750                .into_key(),
751            ),
752            // Span with grpc status from meta with StatusCode. prefix
753            (
754                SpanBytes {
755                    meta: HashMap::from([("grpc.code".into(), "StatusCode.UNAVAILABLE".into())]),
756                    ..Default::default()
757                },
758                FixedAggregationKey {
759                    grpc_status_code: Some(14),
760                    is_trace_root: true,
761                    ..Default::default()
762                }
763                .into_key(),
764            ),
765            // Span with grpc status from metrics takes precedence over meta
766            (
767                SpanBytes {
768                    meta: HashMap::from([(
769                        "rpc.grpc.status_code".into(),
770                        "PERMISSION_DENIED".into(),
771                    )]),
772                    metrics: HashMap::from([("rpc.grpc.status_code".into(), 2.0)]),
773                    ..Default::default()
774                },
775                FixedAggregationKey {
776                    grpc_status_code: Some(7),
777                    is_trace_root: true,
778                    ..Default::default()
779                }
780                .into_key(),
781            ),
782            // Span with grpc status from metrics via secondary key
783            (
784                SpanBytes {
785                    metrics: HashMap::from([("grpc.code".into(), 3.0)]),
786                    ..Default::default()
787                },
788                FixedAggregationKey {
789                    grpc_status_code: Some(3),
790                    is_trace_root: true,
791                    ..Default::default()
792                }
793                .into_key(),
794            ),
795            // Span with invalid grpc status string
796            (
797                SpanBytes {
798                    meta: HashMap::from([("rpc.grpc.status_code".into(), "NOPE".into())]),
799                    ..Default::default()
800                },
801                FixedAggregationKey {
802                    is_trace_root: true,
803                    ..Default::default()
804                }
805                .into_key(),
806            ),
807            // Span with service source set by integration
808            (
809                SpanBytes {
810                    service: "my-service".into(),
811                    name: "op".into(),
812                    resource: "res".into(),
813                    span_id: 1,
814                    parent_id: 0,
815                    meta: HashMap::from([("_dd.svc_src".into(), "redis".into())]),
816                    ..Default::default()
817                },
818                FixedAggregationKey {
819                    service_name: "my-service".into(),
820                    operation_name: "op".into(),
821                    resource_name: "res".into(),
822                    is_trace_root: true,
823                    service_source: "redis".into(),
824                    ..Default::default()
825                }
826                .into_key(),
827            ),
828            // Span with service source set by configuration option
829            (
830                SpanBytes {
831                    service: "my-service".into(),
832                    name: "op".into(),
833                    resource: "res".into(),
834                    span_id: 1,
835                    parent_id: 0,
836                    meta: HashMap::from([("_dd.svc_src".into(), "opt.split_by_tag".into())]),
837                    ..Default::default()
838                },
839                FixedAggregationKey {
840                    service_name: "my-service".into(),
841                    operation_name: "op".into(),
842                    resource_name: "res".into(),
843                    is_trace_root: true,
844                    service_source: "opt.split_by_tag".into(),
845                    ..Default::default()
846                }
847                .into_key(),
848            ),
849            // Span without service source (default service name)
850            (
851                SpanBytes {
852                    service: "my-service".into(),
853                    name: "op".into(),
854                    resource: "res".into(),
855                    span_id: 1,
856                    parent_id: 0,
857                    ..Default::default()
858                },
859                FixedAggregationKey {
860                    service_name: "my-service".into(),
861                    operation_name: "op".into(),
862                    resource_name: "res".into(),
863                    is_trace_root: true,
864                    service_source: "".into(),
865                    ..Default::default()
866                }
867                .into_key(),
868            ),
869        ];
870
871        let test_peer_tags = vec![
872            "aws.s3.bucket".to_string(),
873            "db.instance".to_string(),
874            "db.system".to_string(),
875        ];
876
877        let test_cases_with_peer_tags: Vec<(SpanSlice, OwnedAggregationKey)> = vec![
878            // Span with peer tags with peertags aggregation enabled
879            (
880                SpanSlice {
881                    service: "service",
882                    name: "op",
883                    resource: "res",
884                    span_id: 1,
885                    parent_id: 0,
886                    meta: HashMap::from([("span.kind", "client"), ("aws.s3.bucket", "bucket-a")]),
887                    ..Default::default()
888                },
889                FixedAggregationKey {
890                    service_name: "service".into(),
891                    operation_name: "op".into(),
892                    resource_name: "res".into(),
893                    span_kind: "client".into(),
894                    is_trace_root: true,
895                    ..Default::default()
896                }
897                .into_key_with_peers(vec![("aws.s3.bucket".into(), "bucket-a".into())]),
898            ),
899            // Span with multiple peer tags with peertags aggregation enabled
900            (
901                SpanSlice {
902                    service: "service",
903                    name: "op",
904                    resource: "res",
905                    span_id: 1,
906                    parent_id: 0,
907                    meta: HashMap::from([
908                        ("span.kind", "producer"),
909                        ("aws.s3.bucket", "bucket-a"),
910                        ("db.instance", "dynamo.test.us1"),
911                        ("db.system", "dynamodb"),
912                    ]),
913                    ..Default::default()
914                },
915                FixedAggregationKey {
916                    service_name: "service".into(),
917                    operation_name: "op".into(),
918                    resource_name: "res".into(),
919                    span_kind: "producer".into(),
920                    is_trace_root: true,
921                    ..Default::default()
922                }
923                .into_key_with_peers(vec![
924                    ("aws.s3.bucket".into(), "bucket-a".into()),
925                    ("db.instance".into(), "dynamo.test.us1".into()),
926                    ("db.system".into(), "dynamodb".into()),
927                ]),
928            ),
929            // Span with multiple peer tags with peertags aggregation enabled and span kind is
930            // server
931            (
932                SpanSlice {
933                    service: "service",
934                    name: "op",
935                    resource: "res",
936                    span_id: 1,
937                    parent_id: 0,
938                    meta: HashMap::from([
939                        ("span.kind", "server"),
940                        ("aws.s3.bucket", "bucket-a"),
941                        ("db.instance", "dynamo.test.us1"),
942                        ("db.system", "dynamodb"),
943                    ]),
944                    ..Default::default()
945                },
946                FixedAggregationKey {
947                    service_name: "service".into(),
948                    operation_name: "op".into(),
949                    resource_name: "res".into(),
950                    span_kind: "server".into(),
951                    is_trace_root: true,
952                    ..Default::default()
953                }
954                .into_key(),
955            ),
956        ];
957
958        for (span, expected_key) in test_cases {
959            let borrowed_key = BorrowedAggregationKey::from_span(&span, &[]);
960            assert_eq!(
961                OwnedAggregationKey::from(&borrowed_key),
962                expected_key,
963                "for span {span:?}"
964            );
965            assert_eq!(
966                get_hash(&borrowed_key),
967                get_hash(&OwnedAggregationKey::from(&borrowed_key))
968            );
969        }
970
971        for (span, expected_key) in test_cases_with_peer_tags {
972            let borrowed_key = BorrowedAggregationKey::from_span(&span, test_peer_tags.as_slice());
973            assert_eq!(OwnedAggregationKey::from(&borrowed_key), expected_key);
974            assert_eq!(
975                get_hash(&borrowed_key),
976                get_hash(&OwnedAggregationKey::from(&borrowed_key))
977            );
978        }
979    }
980
981    #[test]
982    fn test_peer_tag_ip_quantization_in_aggregation_key() {
983        let peer_tag_keys = vec!["peer.hostname".to_string(), "db.instance".to_string()];
984
985        // IPv4 address peer tag gets replaced with blocked-ip-address
986        let span_ipv4 = SpanSlice {
987            service: "service",
988            name: "op",
989            resource: "res",
990            span_id: 1,
991            parent_id: 0,
992            meta: HashMap::from([
993                ("span.kind", "client"),
994                ("peer.hostname", "10.1.2.3"),
995                ("db.instance", "my-db"),
996            ]),
997            ..Default::default()
998        };
999        let key = BorrowedAggregationKey::from_span(&span_ipv4, &peer_tag_keys);
1000        let owned = OwnedAggregationKey::from(&key);
1001        assert_eq!(
1002            owned.peer_tags,
1003            vec![
1004                (
1005                    "peer.hostname".to_string(),
1006                    "blocked-ip-address".to_string()
1007                ),
1008                ("db.instance".to_string(), "my-db".to_string()),
1009            ]
1010        );
1011
1012        // IPv6 address peer tag gets replaced with blocked-ip-address
1013        let span_ipv6 = SpanSlice {
1014            service: "service",
1015            name: "op",
1016            resource: "res",
1017            span_id: 1,
1018            parent_id: 0,
1019            meta: HashMap::from([
1020                ("span.kind", "client"),
1021                ("peer.hostname", "2001:db8:3333:4444:CCCC:DDDD:EEEE:FFFF"),
1022            ]),
1023            ..Default::default()
1024        };
1025        let ipv6_keys = vec!["peer.hostname".to_string()];
1026        let key = BorrowedAggregationKey::from_span(&span_ipv6, &ipv6_keys);
1027        let owned = OwnedAggregationKey::from(&key);
1028        assert_eq!(
1029            owned.peer_tags,
1030            vec![(
1031                "peer.hostname".to_string(),
1032                "blocked-ip-address".to_string()
1033            )]
1034        );
1035
1036        // Non-IP peer tags pass through unchanged
1037        let span_non_ip = SpanSlice {
1038            service: "service",
1039            name: "op",
1040            resource: "res",
1041            span_id: 1,
1042            parent_id: 0,
1043            meta: HashMap::from([("span.kind", "client"), ("db.instance", "dynamo.test.us1")]),
1044            ..Default::default()
1045        };
1046        let non_ip_keys = vec!["db.instance".to_string()];
1047        let key = BorrowedAggregationKey::from_span(&span_non_ip, &non_ip_keys);
1048        let owned = OwnedAggregationKey::from(&key);
1049        assert_eq!(
1050            owned.peer_tags,
1051            vec![("db.instance".to_string(), "dynamo.test.us1".to_string())]
1052        );
1053    }
1054
1055    #[test]
1056    fn test_grpc_status_str_to_int_value() {
1057        // Numeric strings parse directly
1058        assert_eq!(grpc_status_str_to_int_value("0"), Some(0));
1059        assert_eq!(grpc_status_str_to_int_value("14"), Some(14));
1060        assert_eq!(grpc_status_str_to_int_value("255"), Some(255));
1061        assert_eq!(grpc_status_str_to_int_value("256"), None);
1062        assert_eq!(grpc_status_str_to_int_value("-1"), None);
1063
1064        // Named status codes (uppercase)
1065        assert_eq!(grpc_status_str_to_int_value("OK"), Some(0));
1066        assert_eq!(grpc_status_str_to_int_value("CANCELLED"), Some(1));
1067        assert_eq!(grpc_status_str_to_int_value("UNKNOWN"), Some(2));
1068        assert_eq!(grpc_status_str_to_int_value("INVALID_ARGUMENT"), Some(3));
1069        assert_eq!(grpc_status_str_to_int_value("DEADLINE_EXCEEDED"), Some(4));
1070        assert_eq!(grpc_status_str_to_int_value("NOT_FOUND"), Some(5));
1071        assert_eq!(grpc_status_str_to_int_value("ALREADY_EXISTS"), Some(6));
1072        assert_eq!(grpc_status_str_to_int_value("PERMISSION_DENIED"), Some(7));
1073        assert_eq!(grpc_status_str_to_int_value("UNAUTHENTICATED"), Some(16));
1074        assert_eq!(grpc_status_str_to_int_value("RESOURCE_EXHAUSTED"), Some(8));
1075        assert_eq!(grpc_status_str_to_int_value("FAILED_PRECONDITION"), Some(9));
1076        assert_eq!(grpc_status_str_to_int_value("ABORTED"), Some(10));
1077        assert_eq!(grpc_status_str_to_int_value("OUT_OF_RANGE"), Some(11));
1078        assert_eq!(grpc_status_str_to_int_value("UNIMPLEMENTED"), Some(12));
1079        assert_eq!(grpc_status_str_to_int_value("INTERNAL"), Some(13));
1080        assert_eq!(grpc_status_str_to_int_value("UNAVAILABLE"), Some(14));
1081        assert_eq!(grpc_status_str_to_int_value("DATA_LOSS"), Some(15));
1082
1083        // Case-insensitive matching
1084        assert_eq!(grpc_status_str_to_int_value("ok"), Some(0));
1085        assert_eq!(grpc_status_str_to_int_value("Cancelled"), Some(1));
1086        assert_eq!(grpc_status_str_to_int_value("not_found"), Some(5));
1087
1088        // StatusCode. prefix is stripped
1089        assert_eq!(grpc_status_str_to_int_value("StatusCode.OK"), Some(0));
1090        assert_eq!(
1091            grpc_status_str_to_int_value("StatusCode.UNAVAILABLE"),
1092            Some(14)
1093        );
1094        assert_eq!(
1095            grpc_status_str_to_int_value("StatusCode.not_found"),
1096            Some(5)
1097        );
1098
1099        // Alternate spellings
1100        assert_eq!(grpc_status_str_to_int_value("CANCELED"), Some(1));
1101
1102        // Unknown / empty strings
1103        assert_eq!(grpc_status_str_to_int_value("NOPE"), None);
1104        assert_eq!(grpc_status_str_to_int_value(""), None);
1105        assert_eq!(
1106            grpc_status_str_to_int_value("this_is_a_kinda_long_string_that_needs_upcasing"),
1107            None
1108        );
1109
1110        // Non ascii
1111        assert_eq!(grpc_status_str_to_int_value("🤣"), None);
1112    }
1113}