diff --git a/Cargo.lock b/Cargo.lock index 2917cbba83..3fd5276f74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1954,9 +1954,9 @@ checksum = "673464e1e314dd67a0fd9544abc99e8eb28d0c7e3b69b033bcff9b2d00b87333" [[package]] name = "futures" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" dependencies = [ "futures-channel", "futures-core", @@ -1969,9 +1969,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" dependencies = [ "futures-core", "futures-sink", @@ -1979,15 +1979,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" [[package]] name = "futures-executor" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" dependencies = [ "futures-core", "futures-task", @@ -1996,9 +1996,9 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" [[package]] name = "futures-lite" @@ -2015,9 +2015,9 @@ dependencies = [ [[package]] name = "futures-macro" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", @@ -2026,15 +2026,15 @@ dependencies = [ [[package]] name = "futures-sink" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" [[package]] name = "futures-task" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" [[package]] name = "futures-timer" @@ -2044,9 +2044,9 @@ checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" [[package]] name = "futures-util" -version = "0.3.31" +version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-channel", "futures-core", @@ -2056,7 +2056,6 @@ dependencies = [ "futures-task", "memchr", "pin-project-lite", - "pin-utils", "slab", ] @@ -3495,6 +3494,7 @@ dependencies = [ "arc-swap", "async-trait", "criterion", + "futures", "hashbrown 0.15.1", "http", "httpmock", diff --git a/datadog-ipc/src/shm_stats.rs b/datadog-ipc/src/shm_stats.rs index 6a7c513fda..87b55981da 100644 --- a/datadog-ipc/src/shm_stats.rs +++ b/datadog-ipc/src/shm_stats.rs @@ -56,7 +56,9 @@ use zwohash::ZwoHasher; use libdd_ddsketch::DDSketch; use libdd_trace_protobuf::pb; -use libdd_trace_stats::span_concentrator::{FixedAggregationKey, FlushableConcentrator}; +use libdd_trace_stats::span_concentrator::{ + FixedAggregationKey, FlushResult, FlushableConcentrator, +}; use crate::platform::{FileBackedHandle, MappedMem, NamedShmHandle}; @@ -820,8 +822,13 @@ impl ShmSpanConcentrator { } impl FlushableConcentrator for ShmSpanConcentrator { - fn flush_buckets(&mut self, force: bool) -> (Vec, u64) { - (self.drain_buckets(force), 0) + fn flush_buckets(&mut self, force: bool) -> FlushResult { + // The SHM concentrator does not perform client-side obfuscation. + FlushResult { + obfuscated_buckets: vec![], + unobfuscated_buckets: self.drain_buckets(force), + collapsed_spans: 0, + } } } diff --git a/datadog-sidecar/src/service/stats_flusher.rs b/datadog-sidecar/src/service/stats_flusher.rs index 13332bbb82..ac144e669d 100644 --- a/datadog-sidecar/src/service/stats_flusher.rs +++ b/datadog-sidecar/src/service/stats_flusher.rs @@ -111,12 +111,6 @@ fn make_exporter( s.meta.clone(), endpoint, NativeCapabilities::new_client(), - // Sidecar does not perform client-side stats obfuscation. Pass a disabled - // default so the `datadog-obfuscation-version` header is never sent. - #[cfg(feature = "stats-obfuscation")] - Arc::new(arc_swap::ArcSwap::from_pointee( - libdd_trace_stats::span_concentrator::StatsComputationObfuscationConfig::default(), - )), #[cfg(feature = "stats-obfuscation")] "0", None, diff --git a/libdd-data-pipeline/results.xml b/libdd-data-pipeline/results.xml deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/libdd-data-pipeline/src/otlp/metrics.rs b/libdd-data-pipeline/src/otlp/metrics.rs index 46104b12f1..21c30c1785 100644 --- a/libdd-data-pipeline/src/otlp/metrics.rs +++ b/libdd-data-pipeline/src/otlp/metrics.rs @@ -255,7 +255,7 @@ pub struct OtlpStatsExporter { impl OtlpStatsExporter { /// Flush the concentrator and export stats; returns `Ok(true)` if anything was sent. async fn send(&self, force_flush: bool, max_retries: u32) -> anyhow::Result { - let (buckets, _collapsed_count) = { + let buckets = { #[allow(clippy::unwrap_used)] let mut c = self.concentrator.lock().unwrap(); c.flush_with_otlp_exact(SystemTime::now(), force_flush) diff --git a/libdd-data-pipeline/src/trace_exporter/stats.rs b/libdd-data-pipeline/src/trace_exporter/stats.rs index 8d7c7a47f1..eaff519d3d 100644 --- a/libdd-data-pipeline/src/trace_exporter/stats.rs +++ b/libdd-data-pipeline/src/trace_exporter/stats.rs @@ -107,7 +107,7 @@ fn is_obfuscation_active(agent_info: &AgentInfo) -> bool { agent_info .info .obfuscation_version - .is_some_and(|v| v >= 1 && v <= SUPPORTED_OBFUSCATION_VERSION) + .is_some_and(|v| v >= 1 && v == SUPPORTED_OBFUSCATION_VERSION) } #[cfg(not(target_arch = "wasm32"))] @@ -170,8 +170,6 @@ fn create_and_start_stats_worker< Endpoint::from_url(add_path(ctx.endpoint_url, STATS_ENDPOINT)), capabilities.clone(), #[cfg(feature = "stats-obfuscation")] - client_side_stats.obfuscation_config.clone(), - #[cfg(feature = "stats-obfuscation")] SUPPORTED_OBFUSCATION_VERSION_STR, #[cfg(feature = "telemetry")] ctx.telemetry.clone(), @@ -258,6 +256,7 @@ fn update_obfuscation_config( ) { let obfuscation_active = client_side_stats.obfuscation_enabled && is_obfuscation_active(agent_info); + // FIXME(APMSP-3720): there is more than this to obfuscation config let sql_obfuscation_mode = (|| { agent_info .info diff --git a/libdd-trace-stats/Cargo.toml b/libdd-trace-stats/Cargo.toml index fd2cf45b78..bb87e1ee0f 100644 --- a/libdd-trace-stats/Cargo.toml +++ b/libdd-trace-stats/Cargo.toml @@ -29,6 +29,7 @@ tokio = { version = "1.23", features = ["macros", "time"], default-features = fa tokio-util = "0.7.11" tracing = { version = "0.1", default-features = false } async-trait = "0.1.85" +futures = { version = "0.3.32", default-features = false, features = ["alloc"] } [target.'cfg(not(target_arch = "wasm32"))'.dependencies] libdd-capabilities-impl = { version = "3.0.0", path = "../libdd-capabilities-impl", default-features = false } diff --git a/libdd-trace-stats/src/span_concentrator/aggregation.rs b/libdd-trace-stats/src/span_concentrator/aggregation.rs index 5d17d54056..d42a1d9294 100644 --- a/libdd-trace-stats/src/span_concentrator/aggregation.rs +++ b/libdd-trace-stats/src/span_concentrator/aggregation.rs @@ -431,6 +431,10 @@ pub(super) struct StatsBucket { max_entries: usize, /// Number of spans collapsed into the overflow bucket due to cardinality limiting. collapsed_count: u64, + /// Indicates if stats obfuscated in this bucket. This is set once at creation and stays + /// constant per bucket + #[cfg(feature = "stats-obfuscation")] + pub(super) obfuscated: bool, } impl StatsBucket { @@ -438,12 +442,18 @@ impl StatsBucket { /// /// `max_entries` is the maximum number of distinct aggregation keys the bucket will hold. /// Once the limit is reached, new distinct keys are collapsed into the overflow sentinel key. - pub(super) fn new(start_timestamp: u64, max_entries: usize) -> Self { + pub(super) fn new( + start_timestamp: u64, + max_entries: usize, + #[cfg(feature = "stats-obfuscation")] obfuscation_enabled: bool, + ) -> Self { Self { data: HashMap::new(), start: start_timestamp, max_entries, collapsed_count: 0, + #[cfg(feature = "stats-obfuscation")] + obfuscated: obfuscation_enabled, } } diff --git a/libdd-trace-stats/src/span_concentrator/mod.rs b/libdd-trace-stats/src/span_concentrator/mod.rs index f1beb5eab9..0f233daca2 100644 --- a/libdd-trace-stats/src/span_concentrator/mod.rs +++ b/libdd-trace-stats/src/span_concentrator/mod.rs @@ -16,18 +16,42 @@ pub use aggregation::{FixedAggregationKey, OtlpExactCell, OtlpExactGroup, OtlpSt pub mod stat_span; pub use stat_span::StatSpan; +/// Result of flushing a concentrator. +/// +/// Obfuscated and un-obfuscated buckets are kept separate because they must be sent in distinct +/// stats payloads: only the obfuscated payload carries the `datadog-obfuscation-version` header. +pub struct FlushResult { + /// Buckets whose resource names were obfuscated client-side. + pub obfuscated_buckets: Vec, + /// Buckets whose resource names were left as-is. + pub unobfuscated_buckets: Vec, + /// Total number of spans that were collapsed into the overflow sentinel bucket due to + /// cardinality limiting across all flushed time buckets. + pub collapsed_spans: u64, +} + +impl FlushResult { + /// All flushed buckets regardless of obfuscation. + #[cfg(test)] + pub fn all_buckets(self) -> Vec { + let mut buckets = self.obfuscated_buckets; + buckets.extend(self.unobfuscated_buckets); + buckets + } +} + /// Concentrators that can provide raw time buckets for export implement this trait. /// /// `StatsExporter` is generic over `C: FlushableConcentrator` so it can work with /// both the in-process [`SpanConcentrator`] and the SHM-backed `ShmSpanConcentrator`. pub trait FlushableConcentrator { - /// Flush time buckets and return them together with the number of spans that were - /// collapsed into the overflow sentinel bucket due to cardinality limiting. - fn flush_buckets(&mut self, force: bool) -> (Vec, u64); + /// Flush time buckets and return them together with flush metadata. If `force` is true, flush + /// all buckets. See [`FlushResult`] for the returned data. + fn flush_buckets(&mut self, force: bool) -> FlushResult; } impl FlushableConcentrator for SpanConcentrator { - fn flush_buckets(&mut self, force: bool) -> (Vec, u64) { + fn flush_buckets(&mut self, force: bool) -> FlushResult { self.flush(SystemTime::now(), force) } } @@ -190,7 +214,23 @@ impl SpanConcentrator { if bucket_timestamp < self.oldest_timestamp { bucket_timestamp = self.oldest_timestamp; } - let obfuscated_resource = self.compute_obfuscated_span(span); + + let target_bucket = self.buckets.entry(bucket_timestamp).or_insert_with(|| { + StatsBucket::new( + bucket_timestamp, + self.max_entries_per_bucket, + #[cfg(feature = "stats-obfuscation")] + self.obfuscation_config.load().enabled, + ) + }); + #[cfg(feature = "stats-obfuscation")] + let obfuscated_resource = if target_bucket.obfuscated { + Self::compute_obfuscated_span(self.obfuscation_config.load().sql_obfuscation_mode, span) + } else { + None + }; + #[cfg(not(feature = "stats-obfuscation"))] + let obfuscated_resource: Option = None; let agg_key = match obfuscated_resource.as_deref() { Some(res) => BorrowedAggregationKey::from_obfuscated_span( res, @@ -199,65 +239,70 @@ impl SpanConcentrator { ), None => BorrowedAggregationKey::from_span(span, self.peer_tag_keys.as_slice()), }; - self.buckets - .entry(bucket_timestamp) - .or_insert_with(|| StatsBucket::new(bucket_timestamp, self.max_entries_per_bucket)) - .insert( - agg_key, - span.duration(), - span.is_error(), - span.has_top_level(), - ); + target_bucket.insert( + agg_key, + span.duration(), + span.is_error(), + span.has_top_level(), + ); } + #[cfg(feature = "stats-obfuscation")] fn compute_obfuscated_span<'a>( - &self, - #[allow(unused)] span: &'a impl StatSpan<'a>, + sql_obfuscation_mode: libdd_trace_obfuscation::sql::SqlObfuscationMode, + span: &'a impl StatSpan<'a>, ) -> Option { - #[cfg(feature = "stats-obfuscation")] - if self.obfuscation_config.load().enabled { - let dbms_hint: Option<&str> = span.get_meta("db.type"); - return libdd_trace_obfuscation::obfuscate::obfuscate_resource_for_stats( - span.r#type(), - span.resource(), - dbms_hint, - self.obfuscation_config.load().sql_obfuscation_mode, - ); - } - None + let dbms_hint: Option<&str> = span.get_meta("db.type"); + libdd_trace_obfuscation::obfuscate::obfuscate_resource_for_stats( + span.r#type(), + span.resource(), + dbms_hint, + sql_obfuscation_mode, + ) } /// Flush all stats bucket except for the `buffer_len` most recent. If `force` is true, flush /// all buckets. /// - /// Returns a tuple of `(buckets, collapsed_spans)` where `collapsed_spans` is the total number - /// of spans that were collapsed into the overflow sentinel bucket due to cardinality limiting - /// across all flushed time buckets. - pub fn flush(&mut self, now: SystemTime, force: bool) -> (Vec, u64) { - self.drain_due_buckets(now, force, StatsBucket::flush) + /// Obfuscated and un-obfuscated buckets are returned separately, see [`FlushResult`]. + pub fn flush(&mut self, now: SystemTime, force: bool) -> FlushResult { + let (buckets, collapsed_spans) = self.drain_due_buckets(now, force, StatsBucket::flush); + let mut obfuscated_buckets = Vec::new(); + let mut unobfuscated_buckets = Vec::new(); + for (obfuscated, bucket) in buckets { + if obfuscated { + obfuscated_buckets.push(bucket); + } else { + unobfuscated_buckets.push(bucket); + } + } + FlushResult { + obfuscated_buckets, + unobfuscated_buckets, + collapsed_spans, + } } /// Like [`Self::flush`], but also emits exact per-cell scalars alongside each bucket for the /// OTLP trace-metrics path. The protobuf bucket inside each [`OtlpStatsBucket`] is identical /// to what [`Self::flush`] would produce, so the /v0.6/stats agent path is unaffected. - /// - /// Returns a tuple of `(buckets, collapsed_spans)` where `collapsed_spans` is the total number - /// of spans that were collapsed into the overflow sentinel bucket due to cardinality limiting - /// across all flushed time buckets. - pub fn flush_with_otlp_exact( - &mut self, - now: SystemTime, - force: bool, - ) -> (Vec, u64) { - self.drain_due_buckets(now, force, StatsBucket::flush_with_otlp_exact) + pub fn flush_with_otlp_exact(&mut self, now: SystemTime, force: bool) -> Vec { + let (buckets, _) = self.drain_due_buckets(now, force, StatsBucket::flush_with_otlp_exact); + buckets.into_iter().map(|(_, bucket)| bucket).collect() } + /// Drain the buckets that are due for flushing, encoding each with `encode`. + /// + /// Returns a tuple `(buckets, collapsed_spans)` where each encoded bucket is paired with a + /// boolean indicating whether it was obfuscated client-side (always `false` when the + /// `stats-obfuscation` feature is disabled), and `collapsed_spans` is the total number of + /// spans collapsed into the overflow sentinel bucket due to cardinality limiting. fn drain_due_buckets( &mut self, now: SystemTime, force: bool, encode: impl Fn(StatsBucket, u64) -> T, - ) -> (Vec, u64) { + ) -> (Vec<(bool, T)>, u64) { // TODO: Wait for HashMap::extract_if to be stabilized to avoid a full drain let now_timestamp = system_time_to_unix_duration(now).as_nanos() as u64; let buckets: Vec<(u64, StatsBucket)> = self.buckets.drain().collect(); @@ -278,13 +323,18 @@ impl SpanConcentrator { // if the tracer stops while the latest buckets aren't old enough to be flushed. // The "force" boolean skips the delay and flushes all buckets, typically on // shutdown. - if !force && timestamp > (now_timestamp - self.buffer_len as u64 * self.bucket_size) - { + let keep = !force + && timestamp > (now_timestamp - self.buffer_len as u64 * self.bucket_size); + if keep { self.buckets.insert(timestamp, bucket); return None; } total_collapsed += bucket.collapsed_count(); - Some(encode(bucket, self.bucket_size)) + #[cfg(feature = "stats-obfuscation")] + let obfuscated = bucket.obfuscated; + #[cfg(not(feature = "stats-obfuscation"))] + let obfuscated = false; + Some((obfuscated, encode(bucket, self.bucket_size))) }) .collect(); if total_collapsed > 0 { diff --git a/libdd-trace-stats/src/span_concentrator/tests.rs b/libdd-trace-stats/src/span_concentrator/tests.rs index 3fd088e09b..81d7e94778 100644 --- a/libdd-trace-stats/src/span_concentrator/tests.rs +++ b/libdd-trace-stats/src/span_concentrator/tests.rs @@ -126,12 +126,12 @@ fn test_concentrator_oldest_timestamp_cold() { // Assert we didn't insert spans in older buckets for _ in 0..concentrator.buffer_len { - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!(stats.len(), 0, "We should get 0 time buckets"); flushtime += Duration::from_nanos(concentrator.bucket_size); } - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!(stats.len(), 1, "We should get exactly one time bucket"); @@ -188,12 +188,12 @@ fn test_concentrator_oldest_timestamp_hot() { } for _ in 0..(concentrator.buffer_len - 1) { - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert!(stats.is_empty(), "We should get 0 time buckets"); flushtime += Duration::from_nanos(concentrator.bucket_size); } - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!(stats.len(), 1, "We should get exactly one time bucket"); // First oldest bucket aggregates, it should have it all except the @@ -213,7 +213,7 @@ fn test_concentrator_oldest_timestamp_hot() { assert_counts_equal(expected, stats.first().unwrap().stats.clone()); flushtime += Duration::from_nanos(concentrator.bucket_size); - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!(stats.len(), 1, "We should get exactly one time bucket"); // Stats of the last four spans. @@ -278,7 +278,7 @@ fn test_concentrator_stats_totals() { let mut flushtime = now; for _ in 0..=concentrator.buffer_len { - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); if stats.is_empty() { continue; } @@ -528,7 +528,7 @@ fn test_concentrator_stats_counts() { let mut flushtime = now; for _ in 0..=concentrator.buffer_len + 2 { - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); let expected_flushed_timestamps = align_timestamp( system_time_to_unix_duration(flushtime).as_nanos() as u64, concentrator.bucket_size, @@ -553,7 +553,7 @@ fn test_concentrator_stats_counts() { stats.first().unwrap().stats.clone(), ); - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!( stats.len(), 0, @@ -658,10 +658,12 @@ fn test_span_should_be_included_in_stats() { }, ]; - let (stats, _) = concentrator.flush( - now + Duration::from_nanos(concentrator.bucket_size * concentrator.buffer_len as u64), - false, - ); + let stats = concentrator + .flush( + now + Duration::from_nanos(concentrator.bucket_size * concentrator.buffer_len as u64), + false, + ) + .all_buckets(); assert_counts_equal( expected, stats @@ -696,10 +698,12 @@ fn test_ignore_partial_spans() { concentrator.add_span(span); } - let (stats, _) = concentrator.flush( - now + Duration::from_nanos(concentrator.bucket_size * concentrator.buffer_len as u64), - false, - ); + let stats = concentrator + .flush( + now + Duration::from_nanos(concentrator.bucket_size * concentrator.buffer_len as u64), + false, + ) + .all_buckets(); assert_eq!(0, stats.len()); } @@ -727,10 +731,10 @@ fn test_force_flush() { let flushtime = now - Duration::from_secs(3600); // Bucket should not be flushed without force flush - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!(0, stats.len()); - let (stats, _) = concentrator.flush(flushtime, true); + let stats = concentrator.flush(flushtime, true).all_buckets(); assert_eq!(1, stats.len()); } @@ -900,7 +904,9 @@ fn test_peer_tags_aggregation() { }, ]; - let (stats_with_peer_tags, _) = concentrator_with_peer_tags.flush(flushtime, false); + let stats_with_peer_tags = concentrator_with_peer_tags + .flush(flushtime, false) + .all_buckets(); assert_counts_equal( expected_with_peer_tags, stats_with_peer_tags @@ -910,7 +916,9 @@ fn test_peer_tags_aggregation() { .clone(), ); - let (stats_without_peer_tags, _) = concentrator_without_peer_tags.flush(flushtime, false); + let stats_without_peer_tags = concentrator_without_peer_tags + .flush(flushtime, false) + .all_buckets(); assert_counts_equal( expected_without_peer_tags, stats_without_peer_tags @@ -1023,7 +1031,9 @@ fn test_peer_tags_quantization_aggregation() { ..Default::default() }]; - let (stats_with_peer_tags, _) = concentrator_with_peer_tags.flush(flushtime, false); + let stats_with_peer_tags = concentrator_with_peer_tags + .flush(flushtime, false) + .all_buckets(); assert_counts_equal( expected_with_peer_tags, stats_with_peer_tags @@ -1193,7 +1203,7 @@ fn test_base_service_peer_tag() { }, ]; - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_counts_equal( expected, stats @@ -1497,7 +1507,7 @@ fn test_pb_span() { // Flush and get stats let flushtime = now + Duration::from_nanos(concentrator.bucket_size * concentrator.buffer_len as u64); - let (stats, _) = concentrator.flush(flushtime, false); + let stats = concentrator.flush(flushtime, false).all_buckets(); assert_eq!(stats.len(), 1, "Should get exactly one time bucket"); let bucket = &stats[0]; @@ -1610,7 +1620,7 @@ fn test_flush_with_otlp_exact_per_cell_scalars() { concentrator.add_span(s); } - let flushed = concentrator.flush_with_otlp_exact(now, true).0; + let flushed = concentrator.flush_with_otlp_exact(now, true); assert_eq!(flushed.len(), 1); let b = &flushed[0]; assert_eq!(b.exact.len(), 1); @@ -1676,7 +1686,7 @@ fn test_cardinality_limit_collapse() { concentrator.add_span(&span); } - let (buckets, _) = concentrator.flush(SystemTime::now(), true); + let buckets = concentrator.flush(SystemTime::now(), true).all_buckets(); assert!(!buckets.is_empty(), "should get at least one time bucket"); let stats = &buckets[0].stats; @@ -1734,7 +1744,7 @@ fn test_overflow_bucket_counts() { concentrator.add_span(&span); } - let (buckets, _) = concentrator.flush(SystemTime::now(), true); + let buckets = concentrator.flush(SystemTime::now(), true).all_buckets(); assert!(!buckets.is_empty()); let stats = &buckets[0].stats; @@ -1783,7 +1793,7 @@ fn test_no_collapse_within_limit() { concentrator.add_span(&span); } - let (buckets, _) = concentrator.flush(SystemTime::now(), true); + let buckets = concentrator.flush(SystemTime::now(), true).all_buckets(); assert!(!buckets.is_empty()); let stats = &buckets[0].stats; @@ -1835,7 +1845,7 @@ fn test_overflow_bucket_key_sentinel_values() { concentrator.add_span(&first); concentrator.add_span(&second); - let (buckets, _) = concentrator.flush(SystemTime::now(), true); + let buckets = concentrator.flush(SystemTime::now(), true).all_buckets(); assert!(!buckets.is_empty()); let stats = &buckets[0].stats; diff --git a/libdd-trace-stats/src/stats_exporter.rs b/libdd-trace-stats/src/stats_exporter.rs index c771fd5b2b..24aa5d4c1b 100644 --- a/libdd-trace-stats/src/stats_exporter.rs +++ b/libdd-trace-stats/src/stats_exporter.rs @@ -9,12 +9,12 @@ use std::{ time, }; -#[cfg(feature = "stats-obfuscation")] -use crate::span_concentrator::SharedStatsComputationObfuscationConfig; use crate::span_concentrator::{FlushableConcentrator, SpanConcentrator}; use async_trait::async_trait; +use futures::stream::FuturesUnordered; +use futures::StreamExt as _; use libdd_capabilities::{HttpClientCapability, MaybeSend, SleepCapability}; -use libdd_common::{tag, Endpoint}; +use libdd_common::Endpoint; use libdd_shared_runtime::Worker; use libdd_trace_protobuf::pb; use libdd_trace_utils::send_with_retry::{send_with_retry, RetryStrategy}; @@ -96,8 +96,6 @@ pub struct StatsExporter< sequence_id: AtomicU64, capabilities: Cap, #[cfg(feature = "stats-obfuscation")] - obfuscation_config: SharedStatsComputationObfuscationConfig, - #[cfg(feature = "stats-obfuscation")] supported_obfuscation_version: &'static str, /// Optional telemetry handle and context key. #[cfg(feature = "telemetry")] @@ -127,8 +125,6 @@ impl meta: StatsMetadata, endpoint: Endpoint, capabilities: Cap, - #[cfg(feature = "stats-obfuscation")] - obfuscation_config: SharedStatsComputationObfuscationConfig, #[cfg(feature = "stats-obfuscation")] supported_obfuscation_version: &'static str, #[cfg(feature = "telemetry")] telemetry: Option< libdd_telemetry::worker::TelemetryWorkerHandle, @@ -139,7 +135,7 @@ impl let telemetry = telemetry.map(|handle| { let key = handle.register_metric_context( COLLAPSED_SPANS_TELEMETRY_METRIC.to_string(), - vec![tag!("collapsed_spans", "whole_key")], + vec![libdd_common::tag!("collapsed_spans", "whole_key")], libdd_telemetry::data::metrics::MetricType::Count, true, libdd_telemetry::data::metrics::MetricNamespace::Tracers, @@ -154,8 +150,6 @@ impl sequence_id: AtomicU64::new(0), capabilities, #[cfg(feature = "stats-obfuscation")] - obfuscation_config, - #[cfg(feature = "stats-obfuscation")] supported_obfuscation_version, #[cfg(feature = "telemetry")] telemetry, @@ -181,37 +175,68 @@ impl /// case stats cannot be flushed since the concentrator might be corrupted. /// Returns `Ok(true)` if stats were sent, `Ok(false)` if the concentrator had nothing to send. pub async fn send(&self, force_flush: bool) -> anyhow::Result { - let (payload, collapsed_spans) = self.flush(force_flush); + let flush = { + #[allow(clippy::unwrap_used)] + let mut concentrator = self.concentrator.lock().unwrap(); + concentrator.flush_buckets(force_flush) + }; - if collapsed_spans > 0 { + if flush.collapsed_spans > 0 { #[cfg(feature = "telemetry")] if let Some((handle, key)) = &self.telemetry { - let _ = handle.add_point(collapsed_spans as f64, key, vec![]); + let _ = handle.add_point(flush.collapsed_spans as f64, key, vec![]); } #[cfg(feature = "dogstatsd")] if let Some(client) = &self.dogstatsd { client.send(vec![libdd_dogstatsd_client::DogStatsDAction::Count( COLLAPSED_SPANS_HEALTH_METRIC, - collapsed_spans as i64, - [tag!("collapsed_spans", "whole_key")].iter(), + flush.collapsed_spans as i64, + [libdd_common::tag!("collapsed_spans", "whole_key")].iter(), )]); } } - if payload.stats.is_empty() { - return Ok(false); + let futures = FuturesUnordered::new(); + + if !flush.obfuscated_buckets.is_empty() { + futures.push(self.send_payload(flush.obfuscated_buckets, true)); + } + + if !flush.unobfuscated_buckets.is_empty() { + futures.push(self.send_payload(flush.unobfuscated_buckets, false)); } + + let sent_stats = !futures.is_empty(); + + futures + .collect::>>() + .await + .into_iter() + .collect::>()?; + + Ok(sent_stats) + } + + /// Encode the given buckets into a stats payload and send it to the agent. + /// + /// `obfuscated` indicates whether the buckets were obfuscated client-side, in which case the + /// `datadog-obfuscation-version` header is added. + async fn send_payload( + &self, + buckets: Vec, + #[cfg_attr(not(feature = "stats-obfuscation"), allow(unused))] obfuscated: bool, + ) -> anyhow::Result<()> { + let sequence = self.sequence_id.fetch_add(1, Ordering::Relaxed); + let payload = encode_stats_payload(&self.meta, sequence, buckets); let body = rmp_serde::encode::to_vec_named(&payload)?; let mut headers: http::HeaderMap = TracerHeaderTags::from(&self.meta).into(); - headers.insert( http::header::CONTENT_TYPE, libdd_common::header::APPLICATION_MSGPACK, ); - #[cfg(feature = "stats-obfuscation")] - if self.obfuscation_config.load().enabled { + if obfuscated { headers.insert( http::HeaderName::from_static("datadog-obfuscation-version"), http::HeaderValue::from_static(self.supported_obfuscation_version), @@ -228,31 +253,13 @@ impl .await; match result { - Ok(_) => Ok(true), + Ok(_) => Ok(()), Err(err) => { - error!(?err, "Error with the StateExporter when sending stats"); + error!(?err, "Error with the StatsExporter when sending stats"); anyhow::bail!("Failed to send stats: {err}"); } } } - - /// Flush stats from the concentrator into a payload - /// - /// # Arguments - /// - `force_flush` if true, triggers a force flush on the concentrator causing all buckets to - /// be flushed regardless of their age. - /// - /// # Panic - /// Will panic if another thread panicked while holding the concentrator lock in which - /// case stats cannot be flushed since the concentrator might be corrupted. - fn flush(&self, force_flush: bool) -> (pb::ClientStatsPayload, u64) { - let sequence = self.sequence_id.fetch_add(1, Ordering::Relaxed); - #[allow(clippy::unwrap_used)] - let (buckets, collapsed_spans) = - self.concentrator.lock().unwrap().flush_buckets(force_flush); - let payload = encode_stats_payload(&self.meta, sequence, buckets); - (payload, collapsed_spans) - } } #[cfg_attr(not(target_arch = "wasm32"), async_trait)] @@ -316,8 +323,6 @@ pub fn stats_url_from_agent_url(agent_url: &str) -> anyhow::Result { #[cfg(test)] mod tests { use super::*; - #[cfg(feature = "stats-obfuscation")] - use crate::span_concentrator::StatsComputationObfuscationConfig; use httpmock::prelude::*; use httpmock::MockServer; use libdd_capabilities_impl::NativeCapabilities; @@ -353,6 +358,17 @@ mod tests { } fn get_test_concentrator() -> SpanConcentrator { + get_test_concentrator_with_obfuscation_config( + #[cfg(feature = "stats-obfuscation")] + None, + ) + } + + fn get_test_concentrator_with_obfuscation_config( + #[cfg(feature = "stats-obfuscation")] obfuscation_config: Option< + crate::span_concentrator::SharedStatsComputationObfuscationConfig, + >, + ) -> SpanConcentrator { let mut concentrator = SpanConcentrator::new( BUCKETS_DURATION, // Make sure the oldest bucket will be flushed on next send @@ -361,7 +377,7 @@ mod tests { vec![], None, #[cfg(feature = "stats-obfuscation")] - None, + obfuscation_config, ); let mut trace = vec![]; @@ -404,8 +420,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), NativeCapabilities::new_client(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -439,8 +453,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), NativeCapabilities::new_client(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -482,8 +494,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), caps.clone(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -531,8 +541,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), caps.clone(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -581,6 +589,7 @@ mod tests { #[cfg_attr(miri, ignore)] #[tokio::test] async fn test_send_stats_with_obfuscation_header() { + use crate::span_concentrator::StatsComputationObfuscationConfig; use arc_swap::ArcSwap; let server = MockServer::start_async().await; @@ -596,18 +605,20 @@ mod tests { }) .await; + let concentrator = get_test_concentrator_with_obfuscation_config(Some(Arc::new( + ArcSwap::from_pointee(StatsComputationObfuscationConfig { + enabled: true, + ..Default::default() + }), + ))); + let stats_exporter = StatsExporter::new( BUCKETS_DURATION, - Arc::new(Mutex::new(get_test_concentrator())), + Arc::new(Mutex::new(concentrator)), get_test_metadata(), Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), NativeCapabilities::new_client(), #[cfg(feature = "stats-obfuscation")] - Arc::new(ArcSwap::from_pointee(StatsComputationObfuscationConfig { - enabled: true, - ..Default::default() - })), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -703,8 +714,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), NativeCapabilities::new_client(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -754,8 +763,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), NativeCapabilities::new_client(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] None, @@ -806,8 +813,6 @@ mod tests { Endpoint::from_url(stats_url_from_agent_url(&server.url("/")).unwrap()), NativeCapabilities::new_client(), #[cfg(feature = "stats-obfuscation")] - StatsComputationObfuscationConfig::disabled(), - #[cfg(feature = "stats-obfuscation")] "1", #[cfg(feature = "telemetry")] Some(handle),