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