diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000..a86bc77912 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# NDJSON fixtures are byte-compared with compressed streams and must not be translated to CRLF. +*.ndjson text eol=lf diff --git a/Cargo.lock b/Cargo.lock index d06491af90..dd807b985f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -64,12 +64,15 @@ dependencies = [ "serde_json", "serde_with", "stringtheory", + "tempfile", "tikv-jemallocator", "tokio", "tokio-util", "tonic", + "tower", "tracing", "uuid", + "zstd", ] [[package]] diff --git a/bin/agent-data-plane/Cargo.toml b/bin/agent-data-plane/Cargo.toml index a9cc5f27a6..a28c6a94f0 100644 --- a/bin/agent-data-plane/Cargo.toml +++ b/bin/agent-data-plane/Cargo.toml @@ -55,6 +55,7 @@ serde = { workspace = true } serde_json = { workspace = true } serde_with = { workspace = true } stringtheory = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = [ "macros", "net", @@ -67,6 +68,7 @@ tokio-util = { workspace = true } tonic = { workspace = true } tracing = { workspace = true } uuid = { workspace = true, features = ["v4"] } +zstd = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] tokio = { workspace = true, features = ["taskdump"] } @@ -82,7 +84,8 @@ chrono = { workspace = true } [dev-dependencies] datadog-agent-config-testing = { workspace = true } derive-where = { workspace = true } -saluki-components = { workspace = true } +saluki-components = { workspace = true, features = ["test-util"] } saluki-config = { workspace = true, features = ["test-util"] } saluki-core = { workspace = true, features = ["test-util"] } saluki-metrics = { workspace = true, features = ["test"] } +tower = { workspace = true, features = ["util"] } diff --git a/bin/agent-data-plane/src/cli/dogstatsd.rs b/bin/agent-data-plane/src/cli/dogstatsd.rs index 2896b4d917..7f89f5e108 100644 --- a/bin/agent-data-plane/src/cli/dogstatsd.rs +++ b/bin/agent-data-plane/src/cli/dogstatsd.rs @@ -17,9 +17,7 @@ use saluki_components::sources::DEFAULT_REPLAY_LOOPS; #[cfg(target_os = "linux")] use saluki_components::sources::REPLAY_CREDENTIALS_GID; use saluki_config::{DurationString, GenericConfiguration}; -#[cfg(any(target_os = "linux", test))] -use saluki_error::generic_error; -use saluki_error::{ErrorContext as _, GenericError}; +use saluki_error::{generic_error, ErrorContext as _, GenericError}; #[cfg(target_os = "linux")] use saluki_io::net::{unix::uds_sendmsg_with_creds, ProcessCredentials}; use serde::Deserialize; @@ -34,6 +32,9 @@ use tracing::{error, info}; use crate::cli::utils::DataPlaneAPIClient; +mod top; +use self::top::{handle_dogstatsd_dump_contexts, handle_dogstatsd_top, DumpContextsCommand, TopCommand}; + /// DogStatsD-specific debugging commands. #[derive(FromArgs, Debug)] #[argh(subcommand, name = "dogstatsd")] @@ -48,6 +49,8 @@ enum DogstatsdSubcommand { Stats(StatsCommand), Capture(CaptureCommand), Replay(ReplayCommand), + Top(TopCommand), + DumpContexts(DumpContextsCommand), } /// Prints basic statistics about the metrics received by the data plane. @@ -168,34 +171,35 @@ struct StatsResponse<'a> { /// Entrypoint for the `dogstatsd` commands. pub async fn handle_dogstatsd_command(bootstrap_config: &GenericConfiguration, cmd: DogstatsdCommand) { - let mut api_client = match DataPlaneAPIClient::from_config(bootstrap_config) { - Ok(client) => client, - Err(e) => { - error!("Failed to create data plane API client: {:#}", e); - std::process::exit(1); - } - }; + if let Err(error) = run_dogstatsd_command(bootstrap_config, cmd).await { + error!("{:#}", error); + std::process::exit(1); + } +} +async fn run_dogstatsd_command( + bootstrap_config: &GenericConfiguration, cmd: DogstatsdCommand, +) -> Result<(), GenericError> { match cmd.subcommand { DogstatsdSubcommand::Stats(config) => { - if let Err(e) = handle_dogstatsd_stats(&mut api_client, config).await { - error!("Failed to run stats subcommand: {:#}", e); - std::process::exit(1); - } + let mut api_client = data_plane_api_client(bootstrap_config)?; + handle_dogstatsd_stats(&mut api_client, config) + .await + .error_context("Failed to run stats subcommand") } DogstatsdSubcommand::Capture(config) => { - if let Err(e) = handle_dogstatsd_capture(&mut api_client, config).await { - error!("Failed to start DogStatsD capture: {:#}", e); - std::process::exit(1); - } + let mut api_client = data_plane_api_client(bootstrap_config)?; + handle_dogstatsd_capture(&mut api_client, config) + .await + .error_context("Failed to start DogStatsD capture") } DogstatsdSubcommand::Replay(config) => { #[cfg(target_os = "linux")] { - if let Err(e) = handle_dogstatsd_replay(&mut api_client, bootstrap_config, config).await { - error!("Failed to replay DogStatsD traffic: {:#}", e); - std::process::exit(1); - } + let mut api_client = data_plane_api_client(bootstrap_config)?; + handle_dogstatsd_replay(&mut api_client, bootstrap_config, config) + .await + .error_context("Failed to replay DogStatsD traffic") } #[cfg(not(target_os = "linux"))] @@ -205,13 +209,31 @@ pub async fn handle_dogstatsd_command(bootstrap_config: &GenericConfiguration, c loops, } = config; let _ = (replay_file_path, loops); - error!("DogStatsD replay is only supported on Linux."); - std::process::exit(1); + Err(generic_error!("DogStatsD replay is only supported on Linux.")) + } + } + DogstatsdSubcommand::Top(config) => { + let config = config.validate(); + let mut output = std::io::stdout(); + if config.is_offline() { + handle_dogstatsd_top(None, config, &mut output).await + } else { + let mut api_client = data_plane_api_client(bootstrap_config)?; + handle_dogstatsd_top(Some(&mut api_client), config, &mut output).await } } + DogstatsdSubcommand::DumpContexts(_) => { + let mut api_client = data_plane_api_client(bootstrap_config)?; + let mut output = std::io::stdout(); + handle_dogstatsd_dump_contexts(&mut api_client, &mut output).await + } } } +fn data_plane_api_client(config: &GenericConfiguration) -> Result { + DataPlaneAPIClient::from_config(config).error_context("Failed to create data plane API client") +} + async fn handle_dogstatsd_stats(api_client: &mut DataPlaneAPIClient, cmd: StatsCommand) -> Result<(), GenericError> { // Trigger a statistics collection and wait for it to complete. info!( diff --git a/bin/agent-data-plane/src/cli/dogstatsd/top.rs b/bin/agent-data-plane/src/cli/dogstatsd/top.rs new file mode 100644 index 0000000000..708ffaeace --- /dev/null +++ b/bin/agent-data-plane/src/cli/dogstatsd/top.rs @@ -0,0 +1,445 @@ +use std::io::Write; +use std::path::{Path, PathBuf}; + +use argh::FromArgs; +use async_trait::async_trait; +use saluki_error::{generic_error, ErrorContext as _, GenericError}; + +use crate::cli::utils::DataPlaneAPIClient; +use crate::dogstatsd_contexts::read_report; + +/// Displays the DogStatsD contexts with the highest cardinality. +#[derive(FromArgs, Debug)] +#[argh(subcommand, name = "top")] +pub(super) struct TopCommand { + /// read a context dump artifact instead of requesting a new dump + #[argh(option, short = 'p', long = "path")] + path: Option, + + /// set the maximum number of metrics to display + #[argh(option, short = 'm', long = "num-metrics", default = "10")] + num_metrics: usize, + + /// set the maximum number of tags to display for each metric + #[argh(option, short = 't', long = "num-tags")] + num_tags: Option, +} + +impl TopCommand { + pub(super) fn validate(self) -> ValidatedTopCommand { + ValidatedTopCommand { + path: self.path, + num_metrics: self.num_metrics, + num_tags: self.num_tags.unwrap_or(5), + } + } +} + +#[derive(Debug)] +pub(super) struct ValidatedTopCommand { + path: Option, + num_metrics: usize, + num_tags: usize, +} + +impl ValidatedTopCommand { + pub(super) fn is_offline(&self) -> bool { + self.path.is_some() + } +} + +/// Writes a DogStatsD context dump artifact. +#[derive(FromArgs, Debug)] +#[argh(subcommand, name = "dump-contexts")] +pub(super) struct DumpContextsCommand {} + +#[async_trait(?Send)] +pub(super) trait DogStatsDContextDumpRequester { + async fn request_context_dump(&mut self) -> Result; +} + +#[async_trait(?Send)] +impl DogStatsDContextDumpRequester for DataPlaneAPIClient { + async fn request_context_dump(&mut self) -> Result { + self.dogstatsd_contexts_dump().await + } +} + +pub(super) async fn handle_dogstatsd_top( + requester: Option<&mut dyn DogStatsDContextDumpRequester>, cmd: ValidatedTopCommand, output: &mut dyn Write, +) -> Result<(), GenericError> { + let path = match cmd.path { + Some(path) => path, + None => { + let requester = + requester.ok_or_else(|| generic_error!("Online DogStatsD top requires a context dump requester."))?; + let path = requester + .request_context_dump() + .await + .error_context("Failed to request a DogStatsD context dump.")?; + write_dump_path(output, &path)?; + path + } + }; + + let report = read_report(&path) + .with_error_context(|| format!("Failed to read DogStatsD context report from '{}'.", path.display()))?; + let rendered = report.render(cmd.num_metrics, cmd.num_tags); + output + .write_all(rendered.as_bytes()) + .with_error_context(|| format!("Failed to write DogStatsD context report for '{}'.", path.display()))?; + output + .flush() + .with_error_context(|| format!("Failed to flush DogStatsD context report for '{}'.", path.display()))?; + Ok(()) +} + +pub(super) async fn handle_dogstatsd_dump_contexts( + requester: &mut dyn DogStatsDContextDumpRequester, output: &mut dyn Write, +) -> Result<(), GenericError> { + let path = requester + .request_context_dump() + .await + .error_context("Failed to request a DogStatsD context dump.")?; + write_dump_path(output, &path) +} + +fn write_dump_path(output: &mut dyn Write, path: &Path) -> Result<(), GenericError> { + writeln!(output, "Wrote {}", path.display()) + .with_error_context(|| format!("Failed to write the DogStatsD context dump path '{}'.", path.display()))?; + output + .flush() + .with_error_context(|| format!("Failed to flush the DogStatsD context dump path '{}'.", path.display()))?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::io; + use std::path::{Path, PathBuf}; + + use argh::FromArgs as _; + use async_trait::async_trait; + use saluki_error::{generic_error, GenericError}; + + use super::{ + handle_dogstatsd_dump_contexts, handle_dogstatsd_top, DogStatsDContextDumpRequester, TopCommand, + ValidatedTopCommand, + }; + use crate::cli::dogstatsd::{DogstatsdCommand, DogstatsdSubcommand}; + + const PLAIN_FIXTURE: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/dogstatsd_contexts_agent.ndjson" + ); + const GOLDEN_REPORT: &str = concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 3\ta.metric\t(2 env, 1 service)\n", + " 1\tz.metric\t(1 bare, 1 env, 1 image)\n", + ); + + #[test] + fn dogstatsd_top_parses_defaults() { + let command = parse_dogstatsd(&["top"]).expect("top defaults should parse"); + let DogstatsdSubcommand::Top(top) = command.subcommand else { + panic!("expected top subcommand"); + }; + + assert_eq!(top.path, None); + assert_eq!(top.num_metrics, 10); + assert_eq!(top.num_tags, None); + let validated = top.validate(); + assert_eq!(validated.path, None); + assert_eq!(validated.num_metrics, 10); + assert_eq!(validated.num_tags, 5); + } + + #[test] + fn dogstatsd_top_parses_all_short_flags() { + let command = + parse_dogstatsd(&["top", "-p", PLAIN_FIXTURE, "-m", "7", "-t", "3"]).expect("top short flags should parse"); + let DogstatsdSubcommand::Top(top) = command.subcommand else { + panic!("expected top subcommand"); + }; + + assert_eq!(top.path, Some(PathBuf::from(PLAIN_FIXTURE))); + assert_eq!(top.num_metrics, 7); + assert_eq!(top.num_tags, Some(3)); + let validated = top.validate(); + assert_eq!(validated.path, Some(PathBuf::from(PLAIN_FIXTURE))); + assert_eq!(validated.num_metrics, 7); + assert_eq!(validated.num_tags, 3); + } + + #[test] + fn dogstatsd_top_parses_corrected_long_num_tags() { + let command = parse_dogstatsd(&["top", "--num-tags", "4"]).expect("corrected option should parse"); + let DogstatsdSubcommand::Top(top) = command.subcommand else { + panic!("expected top subcommand"); + }; + + assert_eq!(top.num_tags, Some(4)); + assert_eq!(top.validate().num_tags, 4); + } + + #[test] + fn dogstatsd_top_rejects_legacy_mum_tags() { + assert!(parse_dogstatsd(&["top", "--mum-tags", "6"]).is_err()); + } + + #[test] + fn dogstatsd_top_rejects_negative_limits_and_extra_arguments() { + for args in [ + &["top", "--num-metrics", "-1"][..], + &["top", "--num-tags", "-1"][..], + &["top", "artifact.ndjson"][..], + ] { + assert!(parse_dogstatsd(args).is_err(), "arguments should be rejected: {args:?}"); + } + } + + #[test] + fn dogstatsd_dump_contexts_parses_without_arguments_and_rejects_extras() { + let command = parse_dogstatsd(&["dump-contexts"]).expect("dump-contexts should parse"); + assert!(matches!(command.subcommand, DogstatsdSubcommand::DumpContexts(_))); + + let error = parse_dogstatsd(&["dump-contexts", "extra"]).expect_err("extra argument should fail"); + assert!( + error.output.contains("Unrecognized argument: extra"), + "{}", + error.output + ); + } + + #[tokio::test] + async fn dogstatsd_top_reads_an_offline_fixture_without_triggering_a_dump() { + let command = top_command(Some(PathBuf::from(PLAIN_FIXTURE)), 10, None); + let mut requester = FakeRequester::returning_path("unused"); + let mut output = RecordingWriter::default(); + + handle_dogstatsd_top(Some(&mut requester), command, &mut output) + .await + .expect("offline top should succeed"); + + assert_eq!(requester.calls, 0); + assert_eq!(output.text(), GOLDEN_REPORT); + } + + #[tokio::test] + async fn dogstatsd_top_triggers_one_online_dump_before_rendering_the_report() { + let mut requester = FakeRequester::returning_path(PLAIN_FIXTURE); + let mut output = RecordingWriter::default(); + + handle_dogstatsd_top(Some(&mut requester), top_command(None, 10, None), &mut output) + .await + .expect("online top should succeed"); + + assert_eq!(requester.calls, 1); + assert_eq!(output.text(), format!("Wrote {PLAIN_FIXTURE}\n{GOLDEN_REPORT}")); + assert_eq!( + output.flushes.first().unwrap(), + format!("Wrote {PLAIN_FIXTURE}\n").as_bytes() + ); + } + + #[tokio::test] + async fn dogstatsd_dump_contexts_triggers_once_and_prints_only_the_path() { + let mut requester = FakeRequester::returning_path(PLAIN_FIXTURE); + let mut output = RecordingWriter::default(); + + handle_dogstatsd_dump_contexts(&mut requester, &mut output) + .await + .expect("dump-contexts should succeed"); + + assert_eq!(requester.calls, 1); + assert_eq!(output.text(), format!("Wrote {PLAIN_FIXTURE}\n")); + assert_eq!(output.flushes, vec![format!("Wrote {PLAIN_FIXTURE}\n").into_bytes()]); + } + + #[tokio::test] + async fn dogstatsd_dump_contexts_prints_nothing_when_transport_fails() { + let mut requester = FakeRequester::returning_error("injected dump transport failure"); + let mut output = RecordingWriter::default(); + + let error = handle_dogstatsd_dump_contexts(&mut requester, &mut output) + .await + .expect_err("transport failure should propagate"); + + assert_eq!(requester.calls, 1); + assert!(format!("{error:#}").contains("injected dump transport failure")); + assert_eq!(output.text(), ""); + assert!(output.flushes.is_empty()); + } + + #[tokio::test] + async fn dogstatsd_top_keeps_the_flushed_path_when_online_artifact_parsing_fails() { + let artifact = tempfile::NamedTempFile::new().expect("temporary artifact should be created"); + std::fs::write(artifact.path(), b"not-json").expect("corrupt artifact should be written"); + let mut requester = FakeRequester::returning_path(artifact.path()); + let mut output = RecordingWriter::default(); + + let error = handle_dogstatsd_top(Some(&mut requester), top_command(None, 10, None), &mut output) + .await + .expect_err("corrupt artifact should fail"); + + let wrote_line = format!("Wrote {}\n", artifact.path().display()); + assert_eq!(requester.calls, 1); + assert_eq!(output.text(), wrote_line); + assert_eq!(output.flushes, vec![wrote_line.into_bytes()]); + let error_chain = format!("{error:#}"); + assert!(error_chain.contains("DogStatsD context report"), "{error_chain}"); + assert!(error_chain.contains("decode record 1"), "{error_chain}"); + } + + #[tokio::test] + async fn dogstatsd_top_redacts_malformed_artifact_values_from_errors() { + const SENTINEL: &str = "SECRET_TENANT_TAG"; + + let artifact = tempfile::NamedTempFile::new().expect("temporary artifact should be created"); + let record = format!( + "{{\"Name\":\"metric\",\"Host\":\"host\",\"Type\":\"Gauge\",\"TaggerTags\":[],\"MetricTags\":[],\"NoIndex\":\"{SENTINEL}\",\"Source\":1}}\n" + ); + std::fs::write(artifact.path(), record).expect("malformed artifact should be written"); + let mut output = RecordingWriter::default(); + + let error = handle_dogstatsd_top( + None, + top_command(Some(artifact.path().to_owned()), 10, None), + &mut output, + ) + .await + .expect_err("wrong-typed artifact field should fail"); + + let error_chain = format!("{error:#}"); + assert!(!error_chain.contains(SENTINEL), "{error_chain}"); + assert!( + error_chain.contains(&artifact.path().display().to_string()), + "{error_chain}" + ); + assert!(error_chain.contains("decode record 1"), "{error_chain}"); + assert!(error_chain.contains("line 1"), "{error_chain}"); + assert!(error_chain.contains("column"), "{error_chain}"); + } + + #[tokio::test] + async fn dogstatsd_top_renders_only_the_heading_for_an_empty_artifact() { + let artifact = tempfile::NamedTempFile::new().expect("temporary artifact should be created"); + let mut output = RecordingWriter::default(); + + handle_dogstatsd_top( + None, + top_command(Some(artifact.path().to_owned()), 10, None), + &mut output, + ) + .await + .expect("empty artifact should render"); + + assert_eq!( + output.text(), + " Contexts\tMetric name\t(number of unique values for each tag)\n" + ); + } + + #[tokio::test] + async fn dogstatsd_top_applies_custom_and_zero_report_limits() { + let cases = [ + ( + top_command(Some(PathBuf::from(PLAIN_FIXTURE)), 0, Some(5)), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 4\t(other 2 metrics)\n", + ), + ), + ( + top_command(Some(PathBuf::from(PLAIN_FIXTURE)), 10, Some(0)), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 3\ta.metric\t(3 values in 2 other tags)\n", + " 1\tz.metric\t(3 values in 3 other tags)\n", + ), + ), + ( + top_command(Some(PathBuf::from(PLAIN_FIXTURE)), 1, Some(1)), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 3\ta.metric\t(2 env, 1 service)\n", + " 1\tz.metric\t(1 bare, 2 values in 2 other tags)\n", + ), + ), + ]; + + for (command, expected) in cases { + let mut output = RecordingWriter::default(); + handle_dogstatsd_top(None, command, &mut output) + .await + .expect("offline report should render"); + assert_eq!(output.text(), expected); + } + } + + fn parse_dogstatsd(args: &[&str]) -> Result { + DogstatsdCommand::from_args(&["agent-data-plane", "dogstatsd"], args) + } + + fn top_command(path: Option, num_metrics: usize, num_tags: Option) -> ValidatedTopCommand { + TopCommand { + path, + num_metrics, + num_tags, + } + .validate() + } + + struct FakeRequester { + calls: usize, + response: Option>, + } + + impl FakeRequester { + fn returning_path(path: impl AsRef) -> Self { + Self { + calls: 0, + response: Some(Ok(path.as_ref().to_owned())), + } + } + + fn returning_error(message: &'static str) -> Self { + Self { + calls: 0, + response: Some(Err(generic_error!(message))), + } + } + } + + #[async_trait(?Send)] + impl DogStatsDContextDumpRequester for FakeRequester { + async fn request_context_dump(&mut self) -> Result { + self.calls += 1; + self.response.take().expect("fake requester called more than once") + } + } + + #[derive(Default)] + struct RecordingWriter { + bytes: Vec, + flushes: Vec>, + } + + impl RecordingWriter { + fn text(&self) -> String { + String::from_utf8(self.bytes.clone()).expect("test output should be UTF-8") + } + } + + impl io::Write for RecordingWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + self.bytes.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.flushes.push(self.bytes.clone()); + Ok(()) + } + } +} diff --git a/bin/agent-data-plane/src/cli/run.rs b/bin/agent-data-plane/src/cli/run.rs index 5afb6bb1a9..3e4102ecbf 100644 --- a/bin/agent-data-plane/src/cli/run.rs +++ b/bin/agent-data-plane/src/cli/run.rs @@ -27,9 +27,10 @@ use saluki_components::{ relays::otlp::OtlpRelayConfiguration, sources::{ChecksIPCConfiguration, DogStatsDConfiguration, OtlpConfiguration}, transforms::{ - AggregateConfiguration, ApmStatsTransformConfiguration, AutoscalingFailoverGatewayConfiguration, - ChainedConfiguration, DogStatsDMapperConfiguration, HostEnrichmentConfiguration, - MrfMetricsGatewayConfiguration, TraceObfuscationConfiguration, TraceSamplerConfiguration, + AggregateConfiguration, AggregateContextSnapshotHandle, ApmStatsTransformConfiguration, + AutoscalingFailoverGatewayConfiguration, ChainedConfiguration, DogStatsDMapperConfiguration, + HostEnrichmentConfiguration, MrfMetricsGatewayConfiguration, TraceObfuscationConfiguration, + TraceSamplerConfiguration, }, }; use saluki_config::GenericConfiguration; @@ -49,6 +50,7 @@ use crate::{ ottl_filter_processor::OttlFilterConfiguration, ottl_transform_processor::OttlTransformConfiguration, tag_filterlist::TagFilterlistConfiguration, }, + dogstatsd_contexts::DogStatsDContextDumpAPIHandler, internal::{ create_internal_supervisor, logging::LoggingConfigurationTranslator, remote_agent::RemoteAgentBootstrap, DogStatsDControlSurface, TopologyControlSurfaces, @@ -677,6 +679,17 @@ async fn add_baseline_traces_pipeline_to_blueprint( Ok(()) } +fn build_dogstatsd_context_dump_api_handler( + config: &GenericConfiguration, snapshot_handle: AggregateContextSnapshotHandle, +) -> Result { + let run_path = config + .try_get_typed::("run_path") + .error_context("Failed to read configured `run_path` for DogStatsD context dumps.")? + .unwrap_or_default(); + + Ok(DogStatsDContextDumpAPIHandler::new(vec![snapshot_handle], run_path)) +} + async fn add_dsd_pipeline_to_blueprint( blueprint: &mut TopologyBlueprint, config: &GenericConfiguration, @@ -738,6 +751,7 @@ async fn add_dsd_pipeline_to_blueprint( .error_context("Failed to configure metric tag filterlist transform.")?; let dsd_agg_config = AggregateConfiguration::from_configuration(config).error_context("Failed to configure aggregate transform.")?; + let dsd_context_snapshot_handle = dsd_agg_config.context_snapshot_handle(); let dsd_post_agg_filter_config = DogStatsDPostAggregateFilterConfiguration::from_configuration(config) .error_context("Failed to configure DogStatsD post-aggregate filter transform.")?; let events_enrich_config = ChainedConfiguration::default().with_transform_builder( @@ -758,6 +772,7 @@ async fn add_dsd_pipeline_to_blueprint( let stats_api_handler = dsd_stats_config.api_handler(); let capture_api_handler = dsd_config.capture_api_handler(); let replay_api_handler = dsd_config.replay_api_handler(); + let context_dump_api_handler = build_dogstatsd_context_dump_api_handler(config, dsd_context_snapshot_handle)?; blueprint // Components. @@ -801,6 +816,7 @@ async fn add_dsd_pipeline_to_blueprint( stats_api_handler, capture_api_handler, replay_api_handler, + context_dump_api_handler, }) } @@ -896,3 +912,348 @@ fn write_sizing_guide(bounds: ComponentBounds) -> Result<(), GenericError> { Ok(()) } + +#[cfg(test)] +mod tests { + use std::{path::Path, sync::Mutex, time::Duration}; + + use async_trait::async_trait; + use http::{Request, StatusCode}; + use http_body_util::{BodyExt as _, Empty}; + use hyper::body::Bytes; + use saluki_api::{response::Response, APIHandler as _}; + use saluki_components::transforms::{ + aggregate_context_snapshot_channel_for_test, AggregateConfiguration, AggregateContextSnapshotEntry, + AggregateMetricType, ChainedConfiguration, DogStatsDMapperConfiguration, + }; + use saluki_config::{config_from, GenericConfiguration}; + use saluki_context::Context; + use saluki_core::{ + accounting::{ComponentRegistry, MemoryBounds, MemoryBoundsBuilder, MemoryLimiter}, + components::{ + destinations::{Destination, DestinationBuilder, DestinationContext}, + sources::{Source, SourceBuilder, SourceContext}, + ComponentContext, + }, + data_model::event::{metric::Metric, Event, EventType}, + health::HealthRegistry, + runtime::Supervisor, + topology::{OutputDefinition, TopologyBlueprint}, + }; + use saluki_error::{generic_error, GenericError}; + use serde_json::json; + use stringtheory::MetaString; + use tokio::sync::{mpsc, oneshot}; + use tower::ServiceExt as _; + + use super::build_dogstatsd_context_dump_api_handler; + use crate::{ + components::{ + dogstatsd_prefix_filter::DogStatsDPrefixFilterConfiguration, tag_filterlist::TagFilterlistConfiguration, + }, + dogstatsd_contexts::DogStatsDContextDumpAPIHandler, + }; + + const CONTEXT_DUMP_FILENAME: &str = "dogstatsd_contexts.json.zstd"; + const CONTEXT_DUMP_ROUTE: &str = crate::dogstatsd_contexts::CONTEXT_DUMP_ROUTE; + + #[tokio::test] + async fn retained_context_identity_follows_dogstatsd_post_processing() { + tokio::time::timeout(Duration::from_secs(5), async { + let config = config_from(json!({ + "dogstatsd_mapper_profiles": [{ + "name": "retained-context-test", + "prefix": "raw.requests.", + "mappings": [{ + "match": "raw.requests.*", + "name": "mapped.requests", + "tags": { "route": "$1" } + }] + }], + "statsd_metric_namespace": "tenant", + "statsd_metric_namespace_blocklist": [], + "metric_filterlist": ["tenant.raw.blocked"], + "metric_filterlist_match_prefix": false, + "metric_tag_filterlist": [{ + "metric_name": "tenant.mapped.requests", + "action": "exclude", + "tags": ["remove"] + }], + "aggregate_flush_interval": { "secs": 60, "nanos": 0 } + })) + .await; + + let mapper = + DogStatsDMapperConfiguration::from_configuration(&config).expect("mapper configuration should parse"); + let mapper_chain = ChainedConfiguration::default().with_transform_builder("dogstatsd_mapper", mapper); + let prefix_filter = DogStatsDPrefixFilterConfiguration::from_configuration(&config) + .expect("prefix filter configuration should parse"); + let tag_filter = + TagFilterlistConfiguration::from_configuration(&config).expect("tag filter configuration should parse"); + let aggregate = + AggregateConfiguration::from_configuration(&config).expect("aggregate configuration should parse"); + let snapshot_handle = aggregate.context_snapshot_handle(); + + let (events_tx, events_rx) = mpsc::channel(2); + let source = ControlledMetricSourceBuilder { + events: Mutex::new(Some(events_rx)), + outputs: vec![OutputDefinition::default_output(EventType::Metric)], + }; + let component_registry = ComponentRegistry::default(); + let mut blueprint = TopologyBlueprint::new("dogstatsd_retained_identity", &component_registry); + blueprint + .add_source("source", source) + .expect("controlled source should be accepted") + .add_transform("mapper", mapper_chain) + .expect("mapper should be accepted") + .add_transform("prefix_filter", prefix_filter) + .expect("prefix filter should be accepted") + .add_transform("tag_filter", tag_filter) + .expect("tag filter should be accepted") + .add_transform("aggregate", aggregate) + .expect("aggregate should be accepted") + .add_destination("destination", DrainingMetricDestinationBuilder) + .expect("draining destination should be accepted"); + blueprint + .connect_components_in_order([ + "source", + "mapper", + "prefix_filter", + "tag_filter", + "aggregate", + "destination", + ]) + .expect("DogStatsD post-processing topology should connect"); + blueprint + .with_health_registry(HealthRegistry::new()) + .with_memory_limiter(MemoryLimiter::noop()) + .with_ambient_worker_pool(); + + let mut supervisor = + Supervisor::new("dogstatsd-retained-identity").expect("test supervisor should be created"); + supervisor.add_worker(blueprint); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let topology_task = tokio::spawn(async move { supervisor.run_with_shutdown(shutdown_rx).await }); + + events_tx + .send(Event::Metric(Metric::counter("raw.blocked", 1.0))) + .await + .expect("controlled source should accept the blocked metric"); + let input_context = Context::from_static_parts("raw.requests.checkout", &["keep:client", "remove:secret"]); + events_tx + .send(Event::Metric(Metric::counter(input_context.clone(), 1.0))) + .await + .expect("controlled source should accept the retained metric"); + + let expected_context = + Context::from_static_parts("tenant.mapped.requests", &["keep:client", "route:checkout"]); + let snapshot = loop { + let snapshot = snapshot_handle + .snapshot() + .await + .expect("running aggregate should fulfill snapshots"); + if snapshot.iter().any(|entry| entry.context() == &expected_context) { + break snapshot; + } + tokio::task::yield_now().await; + }; + + assert_eq!(snapshot.len(), 1); + assert_eq!(snapshot[0].context(), &expected_context); + assert_eq!(snapshot[0].metric_type(), AggregateMetricType::Counter); + assert_eq!(snapshot[0].context().tags().len(), 2); + assert_eq!( + snapshot[0] + .context() + .tags() + .get_single_tag("keep") + .and_then(|tag| tag.value()), + Some("client") + ); + assert_eq!( + snapshot[0] + .context() + .tags() + .get_single_tag("route") + .and_then(|tag| tag.value()), + Some("checkout") + ); + assert!(snapshot[0].context().tags().get_single_tag("remove").is_none()); + assert!(snapshot.iter().all(|entry| entry.context() != &input_context)); + assert!(snapshot + .iter() + .all(|entry| { !matches!(entry.context().name().as_ref(), "raw.blocked" | "tenant.raw.blocked") })); + + drop(events_tx); + drop(snapshot_handle); + shutdown_tx.send(()).expect("test topology should still be running"); + let topology_result = topology_task.await.expect("topology task should not panic"); + assert!( + topology_result.is_ok(), + "topology should stop cleanly: {topology_result:?}" + ); + }) + .await + .expect("DogStatsD retained identity test should complete without hanging"); + } + + #[tokio::test] + async fn context_dump_handler_reads_configured_run_path_and_uses_supplied_owner() { + let run_directory = tempfile::tempdir().expect("run directory should be created"); + let config = context_dump_config(Some(run_directory.path())).await; + let (snapshot_handle, mut snapshot_responder) = aggregate_context_snapshot_channel_for_test(); + + let handler = build_dogstatsd_context_dump_api_handler(&config, snapshot_handle) + .expect("configured handler should build"); + let owner = tokio::spawn(async move { + snapshot_responder + .respond(vec![snapshot_entry("from.supplied.aggregate")]) + .await + }); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::OK); + let artifact_path = run_directory.path().join(CONTEXT_DUMP_FILENAME); + assert_eq!( + response_body(response).await, + serde_json::to_string(&artifact_path).unwrap() + ); + assert!(artifact_path.is_file()); + owner.await.unwrap().unwrap(); + } + + #[tokio::test] + async fn context_dump_handler_keeps_missing_and_empty_run_path_empty_until_publication() { + let cwd_artifact = std::env::current_dir().unwrap().join(CONTEXT_DUMP_FILENAME); + assert!(!cwd_artifact.exists(), "test requires no pre-existing cwd artifact"); + + for run_path in [None, Some(Path::new(""))] { + let config = context_dump_config(run_path).await; + let (snapshot_handle, mut snapshot_responder) = aggregate_context_snapshot_channel_for_test(); + let handler = build_dogstatsd_context_dump_api_handler(&config, snapshot_handle) + .expect("empty run path should reach publication"); + let owner = tokio::spawn(async move { snapshot_responder.respond(Vec::new()).await }); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert!(!cwd_artifact.exists()); + owner.await.unwrap().unwrap(); + } + } + + struct ControlledMetricSource { + events: mpsc::Receiver, + } + + #[async_trait] + impl Source for ControlledMetricSource { + async fn run(mut self: Box, mut context: SourceContext) -> Result<(), GenericError> { + let shutdown = context.take_shutdown_handle(); + tokio::pin!(shutdown); + let mut events_open = true; + + loop { + tokio::select! { + _ = &mut shutdown => break, + maybe_event = self.events.recv(), if events_open => match maybe_event { + Some(event) => context.dispatcher().dispatch_one(event).await?, + None => events_open = false, + } + } + } + Ok(()) + } + } + + struct ControlledMetricSourceBuilder { + events: Mutex>>, + outputs: Vec>, + } + + #[async_trait] + impl SourceBuilder for ControlledMetricSourceBuilder { + fn outputs(&self) -> &[OutputDefinition] { + &self.outputs + } + + async fn build(&self, _context: ComponentContext) -> Result, GenericError> { + let events = self + .events + .lock() + .map_err(|_| generic_error!("controlled metric source receiver lock is poisoned"))? + .take() + .ok_or_else(|| generic_error!("controlled metric source receiver has already been taken"))?; + Ok(Box::new(ControlledMetricSource { events })) + } + } + + impl MemoryBounds for ControlledMetricSourceBuilder { + fn specify_bounds(&self, _builder: &mut MemoryBoundsBuilder) {} + } + + struct DrainingMetricDestination; + + #[async_trait] + impl Destination for DrainingMetricDestination { + async fn run(self: Box, mut context: DestinationContext) -> Result<(), GenericError> { + while context.events().next().await.is_some() {} + Ok(()) + } + } + + struct DrainingMetricDestinationBuilder; + + #[async_trait] + impl DestinationBuilder for DrainingMetricDestinationBuilder { + fn input_event_type(&self) -> EventType { + EventType::Metric + } + + async fn build(&self, _context: ComponentContext) -> Result, GenericError> { + Ok(Box::new(DrainingMetricDestination)) + } + } + + impl MemoryBounds for DrainingMetricDestinationBuilder { + fn specify_bounds(&self, _builder: &mut MemoryBoundsBuilder) {} + } + + async fn context_dump_config(run_path: Option<&Path>) -> GenericConfiguration { + let mut values = json!({}); + if let Some(run_path) = run_path { + values["run_path"] = json!(run_path); + } + config_from(values).await + } + + fn context_dump_post() -> Request> { + Request::builder() + .method("POST") + .uri(CONTEXT_DUMP_ROUTE) + .body(Empty::new()) + .unwrap() + } + + async fn send(handler: &DogStatsDContextDumpAPIHandler, request: Request>) -> Response { + handler + .generate_routes() + .with_state(handler.generate_initial_state()) + .oneshot(request) + .await + .unwrap() + } + + async fn response_body(response: Response) -> String { + String::from_utf8(response.into_body().collect().await.unwrap().to_bytes().to_vec()).unwrap() + } + + fn snapshot_entry(name: &'static str) -> AggregateContextSnapshotEntry { + AggregateContextSnapshotEntry::for_test( + Context::from_static_name(name), + AggregateMetricType::Gauge, + MetaString::empty(), + ) + } +} diff --git a/bin/agent-data-plane/src/cli/utils.rs b/bin/agent-data-plane/src/cli/utils.rs index e8255951e2..fe5aef8f26 100644 --- a/bin/agent-data-plane/src/cli/utils.rs +++ b/bin/agent-data-plane/src/cli/utils.rs @@ -1,3 +1,5 @@ +use std::path::PathBuf; + use futures::TryFutureExt as _; use http::{header::CONTENT_TYPE, uri::PathAndQuery, Request, Response, StatusCode, Uri}; use http_body_util::BodyExt as _; @@ -10,10 +12,13 @@ use hyper::body::Incoming; use prost::Message as _; use saluki_config::GenericConfiguration; use saluki_error::{generic_error, ErrorContext as _, GenericError}; -use saluki_io::net::{client::http::HttpClient, ListenAddress}; +use saluki_io::net::{ + client::http::{HttpClient, HttpClientBuilder}, + ListenAddress, +}; use serde::{Deserialize, Serialize}; -use crate::config::DataPlaneConfiguration; +use crate::{config::DataPlaneConfiguration, dogstatsd_contexts::CONTEXT_DUMP_ROUTE}; /// Typed API client for interacting with the APIs exposed by ADP. pub struct DataPlaneAPIClient { @@ -50,10 +55,11 @@ impl DataPlaneAPIClient { pub fn from_config(config: &GenericConfiguration) -> Result { let dp_config = DataPlaneConfiguration::from_configuration(config)?; - let listen_address = dp_config.secure_api_listen_address(); - let builder = HttpClient::builder().with_tls_config(|b| b.danger_accept_invalid_certs()); + Self::from_builder(builder, dp_config.secure_api_listen_address()) + } + fn from_builder(builder: HttpClientBuilder, listen_address: &ListenAddress) -> Result { let (builder, authority) = match listen_address { ListenAddress::Tcp(_) => { let local_address = listen_address @@ -193,6 +199,26 @@ impl DataPlaneAPIClient { .and_then(body_when_success) } + /// Requests a DogStatsD context dump and returns its path on the server. + /// + /// The response contains only the server-local artifact path; this method does not download the dump contents. + /// + /// # Errors + /// + /// If the request fails, the server rejects it, or the successful response does not contain a JSON string path, an + /// error is returned. + pub async fn dogstatsd_contexts_dump(&mut self) -> Result { + let uri = self.build_uri(CONTEXT_DUMP_ROUTE, None); + let request = build_dogstatsd_contexts_dump_request(uri); + let response = self + .client + .send(request) + .await + .error_context("Failed to request a DogStatsD context dump from the privileged API endpoint.")?; + let response = process_response_body(response).await?; + path_when_context_dump_success(response) + } + /// Starts a DogStatsD traffic capture. /// /// # Errors @@ -331,6 +357,24 @@ impl DataPlaneAPIClient { } } +fn build_dogstatsd_contexts_dump_request(uri: Uri) -> Request { + Request::post(uri) + .body(String::new()) + .expect("valid DogStatsD context dump request") +} + +fn path_when_context_dump_success(resp: Response) -> Result { + match resp.status() { + status if status.is_success() => serde_json::from_str::(resp.body()) + .error_context("Failed to decode DogStatsD context dump path from response JSON."), + status => Err(generic_error!( + "Failed to create DogStatsD context dump ({}): {}.", + status, + resp.into_body() + )), + } +} + async fn collect_body(body: Incoming) -> Option { // `Collected::to_bytes()` merges all frames. Do not use `Buf::chunk().to_vec()` on an aggregated body: `chunk()` // is only the first contiguous slice (often ~16 KiB), which truncates large JSON such as `/config` responses. @@ -416,11 +460,14 @@ fn empty_when_replay_session_success(resp: Response) -> Result<(), Gener #[cfg(test)] mod tests { - use http::{Response, StatusCode}; + use std::path::PathBuf; - use super::body_when_capture_success; + use http::{Method, Response, StatusCode, Uri}; + + use super::{body_when_capture_success, build_dogstatsd_contexts_dump_request, path_when_context_dump_success}; #[cfg(target_os = "linux")] use super::{body_when_replay_session_success, empty_when_replay_session_success}; + use crate::dogstatsd_contexts::CONTEXT_DUMP_ROUTE; #[test] fn dogstatsd_capture_failed_precondition_surfaces_server_message() { @@ -480,4 +527,63 @@ mod tests { assert_eq!(error.to_string(), "session does not own active replay"); } + + #[test] + fn dogstatsd_contexts_request_is_empty_post() { + let uri = Uri::from_static("https://127.0.0.1:5101/agent/dogstatsd-contexts-dump"); + + let request = build_dogstatsd_contexts_dump_request(uri); + + assert_eq!(request.method(), Method::POST); + assert_eq!(request.uri().path(), CONTEXT_DUMP_ROUTE); + assert!(request.body().is_empty()); + assert!(request.headers().is_empty()); + } + + #[test] + fn dogstatsd_contexts_success_decodes_json_path() { + let response = Response::builder() + .status(StatusCode::OK) + .body(r#""/var/run/datadog/dogstatsd_contexts.json.zstd""#.to_string()) + .expect("valid response"); + + let path = path_when_context_dump_success(response).expect("successful response should decode"); + + assert_eq!(path, PathBuf::from("/var/run/datadog/dogstatsd_contexts.json.zstd")); + } + + #[test] + fn dogstatsd_contexts_malformed_success_is_actionable() { + let response = Response::builder() + .status(StatusCode::OK) + .body("not-json".to_string()) + .expect("valid response"); + + let error = path_when_context_dump_success(response).expect_err("malformed JSON should fail"); + + assert!(error.to_string().contains("DogStatsD context dump path")); + } + + #[test] + fn dogstatsd_contexts_server_errors_include_status_and_body() { + for status in [ + StatusCode::INTERNAL_SERVER_ERROR, + StatusCode::SERVICE_UNAVAILABLE, + StatusCode::GATEWAY_TIMEOUT, + ] { + let response = Response::builder() + .status(status) + .body("dump unavailable".to_string()) + .expect("valid response"); + + let error = path_when_context_dump_success(response).expect_err("server error should fail"); + let message = error.to_string(); + + assert!(message.contains(status.as_str()), "missing status in: {message}"); + assert!( + message.contains("dump unavailable"), + "missing response body in: {message}" + ); + } + } } diff --git a/bin/agent-data-plane/src/dogstatsd_contexts/api.rs b/bin/agent-data-plane/src/dogstatsd_contexts/api.rs new file mode 100644 index 0000000000..96191417aa --- /dev/null +++ b/bin/agent-data-plane/src/dogstatsd_contexts/api.rs @@ -0,0 +1,179 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use futures::future::try_join_all; +use http::StatusCode; +use saluki_api::{ + extract::State, + response::{IntoResponse as _, Response}, + routing::{post, Router}, + APIHandler, Json, +}; +use saluki_components::transforms::{AggregateContextSnapshotEntry, AggregateContextSnapshotHandle}; +use saluki_error::GenericError; +use tokio::sync::Mutex; + +use super::{publish_context_dump, CONTEXT_DUMP_ROUTE}; + +// The Agent has no owner-request timeout. ADP uses a finite bound so a stopped topology cannot hold an HTTP request +// indefinitely; 30 seconds is several orders of magnitude above the measured one-million-context snapshot time. +const PRODUCTION_SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(30); +const SNAPSHOT_UNAVAILABLE: &str = + "DogStatsD context snapshot is unavailable; retry after the aggregate owner is running."; +const SNAPSHOT_TIMED_OUT: &str = "Timed out waiting for the DogStatsD context snapshot; retry the request."; +const PUBLICATION_FAILED: &str = + "Failed to publish DogStatsD context dump; check the configured run path and permissions."; +const PUBLICATION_TASK_FAILED: &str = "DogStatsD context dump publication did not complete; retry the request."; +const NON_UTF8_RESULT_PATH: &str = + "The completed DogStatsD context dump path is not valid UTF-8; configure a UTF-8 run path."; + +#[derive(Clone)] +struct ContextSnapshotCoordinator { + handles: Arc<[AggregateContextSnapshotHandle]>, + timeout: Duration, +} + +impl ContextSnapshotCoordinator { + fn new(handles: Vec, timeout: Duration) -> Self { + Self { + handles: handles.into(), + timeout, + } + } + + async fn snapshot(&self) -> Result, SnapshotError> { + if self.handles.is_empty() { + return Err(SnapshotError::SnapshotUnavailable); + } + + let requests = self.handles.iter().map(AggregateContextSnapshotHandle::snapshot); + let owner_snapshots = tokio::time::timeout(self.timeout, try_join_all(requests)) + .await + .map_err(|_| SnapshotError::SnapshotTimedOut)? + .map_err(|_| SnapshotError::SnapshotUnavailable)?; + let total_len = owner_snapshots.iter().map(Vec::len).sum::(); + let mut owner_snapshots = owner_snapshots.into_iter(); + let mut combined = owner_snapshots + .next() + .expect("a non-empty snapshot handle set always returns at least one owner snapshot"); + combined.reserve(total_len - combined.len()); + for snapshot in owner_snapshots { + combined.extend(snapshot); + } + Ok(combined) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SnapshotError { + SnapshotUnavailable, + SnapshotTimedOut, +} + +trait ContextDumpPublisher: Send + Sync { + fn publish(&self, run_path: PathBuf, snapshot: Vec) + -> Result; +} + +struct FileSystemContextDumpPublisher; + +impl ContextDumpPublisher for FileSystemContextDumpPublisher { + fn publish( + &self, run_path: PathBuf, snapshot: Vec, + ) -> Result { + publish_context_dump(&run_path, &snapshot) + } +} + +#[derive(Clone)] +pub(crate) struct DogStatsDContextDumpAPIState { + coordinator: ContextSnapshotCoordinator, + run_path: PathBuf, + publication_lock: Arc>, + publisher: Arc, +} + +#[derive(Clone)] +pub(crate) struct DogStatsDContextDumpAPIHandler { + state: DogStatsDContextDumpAPIState, +} + +impl DogStatsDContextDumpAPIHandler { + pub(crate) fn new(handles: Vec, run_path: impl Into) -> Self { + Self::new_with_services( + handles, + run_path.into(), + PRODUCTION_SNAPSHOT_TIMEOUT, + Arc::new(FileSystemContextDumpPublisher), + ) + } + + #[cfg(test)] + fn new_for_test( + handles: Vec, run_path: PathBuf, snapshot_timeout: Duration, + publisher: Arc, + ) -> Self { + Self::new_with_services(handles, run_path, snapshot_timeout, publisher) + } + + fn new_with_services( + handles: Vec, run_path: PathBuf, snapshot_timeout: Duration, + publisher: Arc, + ) -> Self { + Self { + state: DogStatsDContextDumpAPIState { + coordinator: ContextSnapshotCoordinator::new(handles, snapshot_timeout), + run_path, + publication_lock: Arc::new(Mutex::new(())), + publisher, + }, + } + } + + async fn dump_handler(State(state): State) -> Response { + let publication_guard = state.publication_lock.clone().lock_owned().await; + let snapshot = match state.coordinator.snapshot().await { + Ok(snapshot) => snapshot, + Err(SnapshotError::SnapshotUnavailable) => { + return (StatusCode::SERVICE_UNAVAILABLE, SNAPSHOT_UNAVAILABLE).into_response(); + } + Err(SnapshotError::SnapshotTimedOut) => { + return (StatusCode::GATEWAY_TIMEOUT, SNAPSHOT_TIMED_OUT).into_response(); + } + }; + + let run_path = state.run_path.clone(); + let publisher = state.publisher.clone(); + let publication = tokio::task::spawn_blocking(move || { + let _publication_guard = publication_guard; + publisher.publish(run_path, snapshot) + }) + .await; + + match publication { + Ok(Ok(path)) => match path.into_os_string().into_string() { + Ok(path) => Json(path).into_response(), + Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, NON_UTF8_RESULT_PATH).into_response(), + }, + Ok(Err(_)) => (StatusCode::INTERNAL_SERVER_ERROR, PUBLICATION_FAILED).into_response(), + Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, PUBLICATION_TASK_FAILED).into_response(), + } + } +} + +impl APIHandler for DogStatsDContextDumpAPIHandler { + type State = DogStatsDContextDumpAPIState; + + fn generate_initial_state(&self) -> Self::State { + self.state.clone() + } + + fn generate_routes(&self) -> Router { + Router::new().route(CONTEXT_DUMP_ROUTE, post(Self::dump_handler)) + } +} + +#[cfg(test)] +#[path = "api_tests.rs"] +mod tests; diff --git a/bin/agent-data-plane/src/dogstatsd_contexts/api_tests.rs b/bin/agent-data-plane/src/dogstatsd_contexts/api_tests.rs new file mode 100644 index 0000000000..557b223fec --- /dev/null +++ b/bin/agent-data-plane/src/dogstatsd_contexts/api_tests.rs @@ -0,0 +1,539 @@ +use std::fs; +use std::path::{Path, PathBuf}; +use std::sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + mpsc, Arc, Condvar, Mutex as StdMutex, +}; +use std::time::Duration; + +use http::{header::CONTENT_TYPE, Request, StatusCode}; +use http_body_util::{BodyExt as _, Empty}; +use hyper::body::Bytes; +use saluki_api::{response::Response, APIHandler}; +use saluki_components::transforms::{ + aggregate_context_snapshot_channel_for_test, AggregateContextSnapshotEntry, AggregateContextSnapshotHandle, + AggregateContextSnapshotPendingResponse, AggregateMetricType, +}; +use saluki_context::Context; +use saluki_error::{generic_error, GenericError}; +use stringtheory::MetaString; +use tower::ServiceExt as _; + +use super::{ + ContextDumpPublisher, ContextSnapshotCoordinator, DogStatsDContextDumpAPIHandler, FileSystemContextDumpPublisher, + SnapshotError, +}; +use crate::dogstatsd_contexts::{ + artifact::{for_each_record, CONTEXT_DUMP_FILENAME}, + publish_context_dump, +}; + +const SENSITIVE_PUBLISHER_DETAIL: &str = "sensitive-publisher-detail"; +const ROUTE: &str = super::CONTEXT_DUMP_ROUTE; + +#[tokio::test] +async fn coordinator_reuses_the_one_owner_snapshot_allocation() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let owner_snapshot = vec![snapshot_entry("owner.one.first"), snapshot_entry("owner.one.second")]; + let expected = owner_snapshot.clone(); + let original_pointer = owner_snapshot.as_ptr() as usize; + let original_capacity = owner_snapshot.capacity(); + let owner = tokio::spawn(async move { responder.respond(owner_snapshot).await }); + let coordinator = ContextSnapshotCoordinator::new(vec![handle], Duration::from_secs(1)); + + let actual = coordinator.snapshot().await.expect("snapshot should succeed"); + + assert_eq!(actual, expected); + assert_eq!(actual.as_ptr() as usize, original_pointer); + assert_eq!(actual.capacity(), original_capacity); + owner.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn coordinator_requests_all_owners_concurrently_and_flattens_each_snapshot_once_in_owner_order() { + let (first_handle, mut first_responder) = aggregate_context_snapshot_channel_for_test(); + let (second_handle, mut second_responder) = aggregate_context_snapshot_channel_for_test(); + let first_snapshot = vec![snapshot_entry("owner.one.first"), snapshot_entry("owner.one.second")]; + let second_snapshot = vec![snapshot_entry("owner.two.only")]; + let expected = [first_snapshot.clone(), second_snapshot.clone()].concat(); + let (second_responded_tx, second_responded_rx) = tokio::sync::oneshot::channel(); + let first_owner = tokio::spawn(async move { + second_responded_rx + .await + .expect("second owner should receive its request"); + first_responder.respond(first_snapshot).await + }); + let second_owner = tokio::spawn(async move { + let result = second_responder.respond(second_snapshot).await; + let _ = second_responded_tx.send(()); + result + }); + let coordinator = ContextSnapshotCoordinator::new(vec![first_handle, second_handle], Duration::from_millis(250)); + + let actual = coordinator.snapshot().await.expect("both snapshots should succeed"); + + assert_eq!(actual, expected); + first_owner.await.unwrap().unwrap(); + second_owner.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn coordinator_reports_no_handles_as_typed_unavailable() { + let coordinator = ContextSnapshotCoordinator::new(Vec::new(), Duration::from_secs(1)); + + let error = coordinator.snapshot().await.expect_err("empty owner set should fail"); + + assert_eq!(error, SnapshotError::SnapshotUnavailable); +} + +#[tokio::test] +async fn coordinator_reports_a_closed_owner_as_typed_unavailable() { + let (handle, responder) = aggregate_context_snapshot_channel_for_test(); + drop(responder); + let coordinator = ContextSnapshotCoordinator::new(vec![handle], Duration::from_secs(1)); + + let error = coordinator.snapshot().await.expect_err("closed owner should fail"); + + assert_eq!(error, SnapshotError::SnapshotUnavailable); +} + +#[tokio::test] +async fn coordinator_reports_elapsed_deadline_as_typed_timeout() { + let (handle, _responder) = aggregate_context_snapshot_channel_for_test(); + let coordinator = ContextSnapshotCoordinator::new(vec![handle], Duration::from_millis(10)); + + let error = coordinator + .snapshot() + .await + .expect_err("unresponsive owner should time out"); + + assert_eq!(error, SnapshotError::SnapshotTimedOut); +} + +#[tokio::test] +async fn canceled_coordinator_request_does_not_make_a_late_owner_response_panic_or_fail() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let coordinator = ContextSnapshotCoordinator::new(vec![handle], Duration::from_secs(1)); + let request = tokio::spawn(async move { coordinator.snapshot().await }); + let pending_response: AggregateContextSnapshotPendingResponse = responder + .receive() + .await + .expect("the owner should accept the snapshot request"); + + request.abort(); + let _ = request.await; + + pending_response.respond(vec![snapshot_entry("late.owner.response")]); +} + +#[tokio::test] +async fn generated_router_rejects_get_with_method_not_allowed() { + let handler = test_handler(Vec::new(), PathBuf::from("unused"), Arc::new(NoOpPublisher)); + let request = Request::builder() + .method("GET") + .uri(ROUTE) + .body(Empty::::new()) + .unwrap(); + + let response = send(&handler, request).await; + + assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED); +} + +#[tokio::test] +async fn router_post_returns_only_the_json_path_to_a_complete_decodable_artifact() { + let run_directory = tempfile::tempdir().unwrap(); + let target = run_directory.path().join(CONTEXT_DUMP_FILENAME); + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let snapshot = vec![snapshot_entry("published.first"), snapshot_entry("published.second")]; + let handler = DogStatsDContextDumpAPIHandler::new(vec![handle], run_directory.path()); + let owner = tokio::spawn(async move { responder.respond(snapshot).await }); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers().get(CONTENT_TYPE).unwrap(), "application/json"); + assert_eq!( + response_body(response).await, + serde_json::to_string(target.to_str().unwrap()).unwrap() + ); + owner.await.unwrap().unwrap(); + let mut names = Vec::new(); + for_each_record(&target, |record| names.push(record.name)).expect("published artifact should decode"); + assert_eq!(names, ["published.first", "published.second"]); + assert_eq!( + directory_entries(run_directory.path()), + vec![CONTEXT_DUMP_FILENAME.to_owned()] + ); +} + +#[tokio::test] +async fn no_handles_and_closed_owner_return_service_unavailable() { + let no_handles = test_handler(Vec::new(), PathBuf::from("unused"), Arc::new(NoOpPublisher)); + let response = send(&no_handles, context_dump_post()).await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response_body(response).await, + "DogStatsD context snapshot is unavailable; retry after the aggregate owner is running." + ); + + let (handle, responder) = aggregate_context_snapshot_channel_for_test(); + drop(responder); + let closed_owner = test_handler(vec![handle], PathBuf::from("unused"), Arc::new(NoOpPublisher)); + let response = send(&closed_owner, context_dump_post()).await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response_body(response).await, + "DogStatsD context snapshot is unavailable; retry after the aggregate owner is running." + ); +} + +#[tokio::test] +async fn owner_stopped_after_accepting_request_returns_service_unavailable() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let handler = test_handler(vec![handle], PathBuf::from("unused"), Arc::new(NoOpPublisher)); + let owner = tokio::spawn(async move { responder.stop_after_receiving().await }); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response_body(response).await, + "DogStatsD context snapshot is unavailable; retry after the aggregate owner is running." + ); + owner.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn snapshot_deadline_returns_gateway_timeout() { + let (handle, _responder) = aggregate_context_snapshot_channel_for_test(); + let handler = DogStatsDContextDumpAPIHandler::new_for_test( + vec![handle], + PathBuf::from("unused"), + Duration::from_millis(10), + Arc::new(NoOpPublisher), + ); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::GATEWAY_TIMEOUT); + assert_eq!( + response_body(response).await, + "Timed out waiting for the DogStatsD context snapshot; retry the request." + ); +} + +#[tokio::test] +async fn empty_and_non_directory_run_paths_reach_publication_and_return_safe_internal_errors() { + let temporary_directory = tempfile::tempdir().unwrap(); + let non_directory = temporary_directory.path().join("not-a-directory"); + fs::write(&non_directory, b"fixture").unwrap(); + + for run_path in [PathBuf::new(), non_directory] { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let handler = DogStatsDContextDumpAPIHandler::new(vec![handle], run_path); + let owner = tokio::spawn(async move { responder.respond(Vec::new()).await }); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + let body = response_body(response).await; + assert_eq!( + body, + "Failed to publish DogStatsD context dump; check the configured run path and permissions." + ); + assert!(!body.contains(SENSITIVE_PUBLISHER_DETAIL)); + owner.await.unwrap().unwrap(); + } + assert_eq!( + directory_entries(temporary_directory.path()), + vec!["not-a-directory".to_owned()] + ); +} + +#[tokio::test] +async fn publication_failure_and_blocking_task_panic_return_safe_internal_errors() { + for publisher in [ + Arc::new(FailingPublisher) as Arc, + Arc::new(PanickingPublisher) as Arc, + ] { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let handler = test_handler(vec![handle], PathBuf::from("unused"), publisher); + let owner = tokio::spawn(async move { responder.respond(Vec::new()).await }); + + let response = send(&handler, context_dump_post()).await; + + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + let body = response_body(response).await; + assert!( + body == "Failed to publish DogStatsD context dump; check the configured run path and permissions." + || body == "DogStatsD context dump publication did not complete; retry the request." + ); + assert!(!body.contains("injected")); + assert!(!body.contains(SENSITIVE_PUBLISHER_DETAIL)); + owner.await.unwrap().unwrap(); + } +} + +#[test] +fn blocking_publisher_release_guard_releases_and_notifies_on_drop() { + let release = Arc::new((StdMutex::new(false), Condvar::new())); + let (waiting_tx, waiting_rx) = mpsc::sync_channel(1); + + std::thread::scope(|scope| { + let waiter_release = release.clone(); + let waiter = scope.spawn(move || { + let (released, wake) = &*waiter_release; + let released = released.lock().unwrap(); + waiting_tx.send(()).unwrap(); + let (released, timeout) = wake + .wait_timeout_while(released, Duration::from_secs(1), |released| !*released) + .unwrap(); + assert!(!timeout.timed_out(), "release guard should notify the waiter"); + assert!(*released); + }); + waiting_rx.recv().unwrap(); + + let release_guard = BlockingPublisherReleaseGuard::new(release); + drop(release_guard); + waiter.join().unwrap(); + }); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn requests_serialize_fixed_path_publisher_executions() { + let (started_tx, started_rx) = mpsc::sync_channel(1); + let release = Arc::new((StdMutex::new(false), Condvar::new())); + let publisher = Arc::new(BlockingFirstPublisher { + active: AtomicUsize::new(0), + max_active: AtomicUsize::new(0), + calls: AtomicUsize::new(0), + first_started: started_tx, + release: release.clone(), + }); + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let handler = test_handler(vec![handle], PathBuf::from("unused"), publisher.clone()); + let owner = tokio::spawn(async move { + responder.respond(Vec::new()).await.unwrap(); + responder.respond(Vec::new()).await.unwrap(); + }); + + let first_handler = handler.clone(); + let first = tokio::spawn(async move { send(&first_handler, context_dump_post()).await }); + tokio::task::spawn_blocking(move || started_rx.recv_timeout(Duration::from_secs(1))) + .await + .unwrap() + .expect("first publisher should start"); + let release_guard = BlockingPublisherReleaseGuard::new(release); + let second_handler = handler.clone(); + let second = tokio::spawn(async move { send(&second_handler, context_dump_post()).await }); + tokio::time::sleep(Duration::from_millis(25)).await; + assert_eq!(publisher.calls.load(Ordering::SeqCst), 1); + drop(release_guard); + + assert_eq!(first.await.unwrap().status(), StatusCode::OK); + assert_eq!(second.await.unwrap().status(), StatusCode::OK); + owner.await.unwrap(); + assert_eq!(publisher.calls.load(Ordering::SeqCst), 2); + assert_eq!(publisher.max_active.load(Ordering::SeqCst), 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn canceling_after_spawn_blocking_starts_still_completes_atomic_publication() { + let run_directory = tempfile::tempdir().unwrap(); + let target = run_directory.path().join(CONTEXT_DUMP_FILENAME); + fs::write(&target, b"original canonical artifact").unwrap(); + let (started_tx, started_rx) = mpsc::sync_channel(1); + let release = Arc::new((StdMutex::new(false), Condvar::new())); + let completed = Arc::new(AtomicBool::new(false)); + let publisher = Arc::new(BlockingRealPublisher { + started: started_tx, + release: release.clone(), + completed: completed.clone(), + }); + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let handler = test_handler(vec![handle], run_directory.path().to_owned(), publisher); + let owner = tokio::spawn(async move { responder.respond(vec![snapshot_entry("after.cancel")]).await }); + let request_handler = handler.clone(); + let request = tokio::spawn(async move { send(&request_handler, context_dump_post()).await }); + tokio::task::spawn_blocking(move || started_rx.recv_timeout(Duration::from_secs(1))) + .await + .unwrap() + .expect("publisher should start"); + let release_guard = BlockingPublisherReleaseGuard::new(release); + + request.abort(); + let _ = request.await; + assert_eq!(fs::read(&target).unwrap(), b"original canonical artifact"); + drop(release_guard); + tokio::time::timeout(Duration::from_secs(2), async { + while !completed.load(Ordering::SeqCst) { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("detached publication should complete"); + + owner.await.unwrap().unwrap(); + let mut names = Vec::new(); + for_each_record(&target, |record| names.push(record.name)).expect("canonical artifact should be complete"); + assert_eq!(names, ["after.cancel"]); + assert_eq!( + directory_entries(run_directory.path()), + vec![CONTEXT_DUMP_FILENAME.to_owned()] + ); +} + +fn test_handler( + handles: Vec, run_path: PathBuf, publisher: Arc, +) -> DogStatsDContextDumpAPIHandler { + DogStatsDContextDumpAPIHandler::new_for_test(handles, run_path, Duration::from_millis(100), publisher) +} + +fn context_dump_post() -> Request> { + Request::builder().method("POST").uri(ROUTE).body(Empty::new()).unwrap() +} + +async fn send(handler: &DogStatsDContextDumpAPIHandler, request: Request>) -> Response { + handler + .generate_routes() + .with_state(handler.generate_initial_state()) + .oneshot(request) + .await + .unwrap() +} + +async fn response_body(response: Response) -> String { + String::from_utf8(response.into_body().collect().await.unwrap().to_bytes().to_vec()).unwrap() +} + +fn snapshot_entry(name: &'static str) -> AggregateContextSnapshotEntry { + AggregateContextSnapshotEntry::for_test( + Context::from_static_name(name), + AggregateMetricType::Gauge, + MetaString::empty(), + ) +} + +fn directory_entries(path: &Path) -> Vec { + let mut entries: Vec<_> = fs::read_dir(path) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + entries.sort_unstable(); + entries +} + +struct NoOpPublisher; + +impl ContextDumpPublisher for NoOpPublisher { + fn publish( + &self, run_path: PathBuf, _snapshot: Vec, + ) -> Result { + Ok(run_path.join(CONTEXT_DUMP_FILENAME)) + } +} + +struct FailingPublisher; + +impl ContextDumpPublisher for FailingPublisher { + fn publish( + &self, _run_path: PathBuf, _snapshot: Vec, + ) -> Result { + Err(generic_error!( + "injected publisher failure with {SENSITIVE_PUBLISHER_DETAIL}" + )) + } +} + +struct PanickingPublisher; + +impl ContextDumpPublisher for PanickingPublisher { + fn publish( + &self, _run_path: PathBuf, _snapshot: Vec, + ) -> Result { + panic!("injected publisher panic") + } +} + +struct BlockingPublisherReleaseGuard { + release: Arc<(StdMutex, Condvar)>, +} + +impl BlockingPublisherReleaseGuard { + fn new(release: Arc<(StdMutex, Condvar)>) -> Self { + Self { release } + } +} + +impl Drop for BlockingPublisherReleaseGuard { + fn drop(&mut self) { + let (released, wake) = &*self.release; + let mut released = released.lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + *released = true; + wake.notify_all(); + } +} + +struct BlockingFirstPublisher { + active: AtomicUsize, + max_active: AtomicUsize, + calls: AtomicUsize, + first_started: mpsc::SyncSender<()>, + release: Arc<(StdMutex, Condvar)>, +} + +impl ContextDumpPublisher for BlockingFirstPublisher { + fn publish( + &self, run_path: PathBuf, _snapshot: Vec, + ) -> Result { + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.max_active.fetch_max(active, Ordering::SeqCst); + let call = self.calls.fetch_add(1, Ordering::SeqCst); + if call == 0 { + self.first_started.send(()).unwrap(); + let (released, wake) = &*self.release; + let mut released = released.lock().unwrap(); + while !*released { + released = wake.wait(released).unwrap(); + } + } + self.active.fetch_sub(1, Ordering::SeqCst); + Ok(run_path.join(CONTEXT_DUMP_FILENAME)) + } +} + +struct BlockingRealPublisher { + started: mpsc::SyncSender<()>, + release: Arc<(StdMutex, Condvar)>, + completed: Arc, +} + +impl ContextDumpPublisher for BlockingRealPublisher { + fn publish( + &self, run_path: PathBuf, snapshot: Vec, + ) -> Result { + self.started.send(()).unwrap(); + let (released, wake) = &*self.release; + let mut released = released.lock().unwrap(); + while !*released { + released = wake.wait(released).unwrap(); + } + drop(released); + let result = publish_context_dump(&run_path, &snapshot); + self.completed.store(true, Ordering::SeqCst); + result + } +} + +#[test] +fn production_publisher_type_uses_the_real_artifact_publisher() { + let run_directory = tempfile::tempdir().unwrap(); + let path = FileSystemContextDumpPublisher + .publish(run_directory.path().to_owned(), vec![snapshot_entry("real.publisher")]) + .unwrap(); + + let mut count = 0; + for_each_record(&path, |_| count += 1).unwrap(); + assert_eq!(count, 1); +} diff --git a/bin/agent-data-plane/src/dogstatsd_contexts/artifact.rs b/bin/agent-data-plane/src/dogstatsd_contexts/artifact.rs new file mode 100644 index 0000000000..9fa33571ae --- /dev/null +++ b/bin/agent-data-plane/src/dogstatsd_contexts/artifact.rs @@ -0,0 +1,335 @@ +use std::error::Error; +use std::fmt; +#[cfg(unix)] +use std::fs; +use std::fs::File; +#[cfg(unix)] +use std::io; +use std::io::{BufRead, BufReader, BufWriter, Read, Write}; +use std::path::{Path, PathBuf}; + +use saluki_components::transforms::{AggregateContextSnapshotEntry, AggregateMetricType}; +use saluki_context::tags::TagSet; +use saluki_error::{ErrorContext as _, GenericError}; +use serde::{ser::SerializeStruct as _, Deserialize, Serialize, Serializer}; +use tempfile::Builder as TempFileBuilder; + +// Zstandard frame magic from RFC 8878, section 3.1.1. +const ZSTD_MAGIC: &[u8; 4] = b"\x28\xb5\x2f\xfd"; +// Matches the fixed Agent output path in demultiplexerendpoint/impl/endpoint.go at the pinned compatibility revision. +pub(crate) const CONTEXT_DUMP_FILENAME: &str = "dogstatsd_contexts.json.zstd"; +// `MetricSourceDogstatsd` is the second value after `MetricSourceUnknown` in the Agent's metricsource.go enum. +const AGENT_METRIC_SOURCE_DOGSTATSD: u32 = 1; + +struct AgentContextSnapshotRecord<'a>(&'a AggregateContextSnapshotEntry); + +impl Serialize for AgentContextSnapshotRecord<'_> { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let entry = self.0; + let context = entry.context(); + let name: &str = context.name().as_ref(); + let mut record = serializer.serialize_struct("AgentContextRecord", 7)?; + record.serialize_field("Name", name)?; + record.serialize_field("Host", context.host().unwrap_or_default())?; + record.serialize_field("Type", agent_metric_type(entry))?; + record.serialize_field("TaggerTags", &BorrowedTags(context.origin_tags()))?; + record.serialize_field("MetricTags", &BorrowedTags(context.tags()))?; + record.serialize_field("NoIndex", &false)?; + record.serialize_field("Source", &AGENT_METRIC_SOURCE_DOGSTATSD)?; + record.end() + } +} + +struct BorrowedTags<'a>(&'a TagSet); + +impl Serialize for BorrowedTags<'_> { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.collect_seq(self.0.into_iter().map(|tag| tag.as_str())) + } +} + +// These spellings match `MetricType.String()` in the Agent's pkg/metrics/metric_sample.go. +fn agent_metric_type(entry: &AggregateContextSnapshotEntry) -> &'static str { + match entry.metric_type() { + AggregateMetricType::Counter => "Counter", + AggregateMetricType::Rate => "Rate", + AggregateMetricType::Gauge => "Gauge", + AggregateMetricType::Set => "Set", + AggregateMetricType::Distribution => "Distribution", + AggregateMetricType::Histogram if entry.unit() == Some("millisecond") => "Historate", + AggregateMetricType::Histogram => "Histogram", + } +} + +fn write_snapshot_records( + writer: &mut W, snapshot: &[AggregateContextSnapshotEntry], path: &Path, +) -> Result<(), GenericError> { + for (index, entry) in snapshot.iter().enumerate() { + serde_json::to_writer(&mut *writer, &AgentContextSnapshotRecord(entry)).with_error_context(|| { + format!( + "Failed to serialize DogStatsD context dump record {} for '{}'.", + index + 1, + path.display() + ) + })?; + writer.write_all(b"\n").with_error_context(|| { + format!( + "Failed to write DogStatsD context dump record {} terminator for '{}'.", + index + 1, + path.display() + ) + })?; + } + Ok(()) +} + +pub(crate) fn publish_context_dump( + run_path: &Path, snapshot: &[AggregateContextSnapshotEntry], +) -> Result { + let target = run_path.join(CONTEXT_DUMP_FILENAME); + if run_path.as_os_str().is_empty() { + return Err(saluki_error::generic_error!( + "Failed to validate run path for DogStatsD context dump target '{}': the run path is empty.", + target.display() + )); + } + + let prefix = format!(".{CONTEXT_DUMP_FILENAME}."); + let mut temporary = TempFileBuilder::new() + .prefix(&prefix) + .suffix(".tmp") + .tempfile_in(run_path) + .with_error_context(|| { + format!( + "Failed to create temporary DogStatsD context dump for target '{}'.", + target.display() + ) + })?; + + #[cfg(unix)] + set_owner_only_permissions(temporary.as_file()).with_error_context(|| { + format!( + "Failed to set owner-only permissions on temporary DogStatsD context dump '{}' for target '{}'.", + temporary.path().display(), + target.display() + ) + })?; + + let file = write_compressed_snapshot(BufWriter::new(temporary.as_file_mut()), snapshot, &target)?; + file.sync_all().with_error_context(|| { + format!( + "Failed to sync temporary DogStatsD context dump for target '{}'.", + target.display() + ) + })?; + + temporary + .persist(&target) + .map_err(|error| error.error) + .with_error_context(|| { + format!( + "Failed to atomically replace DogStatsD context dump target '{}'.", + target.display() + ) + })?; + Ok(target) +} + +#[cfg(unix)] +fn set_owner_only_permissions(file: &File) -> io::Result<()> { + use std::os::unix::fs::PermissionsExt as _; + + file.set_permissions(fs::Permissions::from_mode(0o600)) +} + +fn write_compressed_snapshot( + buffer: BufWriter, snapshot: &[AggregateContextSnapshotEntry], target: &Path, +) -> Result { + let mut encoder = zstd::stream::write::Encoder::new(buffer, 0).with_error_context(|| { + format!( + "Failed to initialize zstd stream for DogStatsD context dump target '{}'.", + target.display() + ) + })?; + write_snapshot_records(&mut encoder, snapshot, target)?; + let buffer = finish_zstd_encoder(encoder, target)?; + finish_buffered_writer(buffer, target) +} + +fn finish_zstd_encoder<'a, W: Write>( + encoder: zstd::stream::write::Encoder<'a, W>, target: &Path, +) -> Result { + encoder.finish().with_error_context(|| { + format!( + "Failed to finalize zstd stream for DogStatsD context dump target '{}'.", + target.display() + ) + }) +} + +fn finish_buffered_writer(buffer: BufWriter, target: &Path) -> Result { + buffer + .into_inner() + .map_err(|error| error.into_error()) + .with_error_context(|| { + format!( + "Failed to finalize buffered DogStatsD context dump output for target '{}'.", + target.display() + ) + }) +} + +#[derive(Debug, Deserialize, Eq, PartialEq)] +#[serde(rename_all = "PascalCase")] +pub(crate) struct AgentContextRecord { + pub(super) name: String, + pub(super) host: String, + #[serde(rename = "Type")] + pub(super) metric_type: String, + pub(super) tagger_tags: Vec, + pub(super) metric_tags: Vec, + pub(super) no_index: bool, + pub(super) source: u32, +} + +#[derive(Clone, Copy, Debug)] +enum JsonErrorCategory { + Io, + Syntax, + Data, + Eof, +} + +impl JsonErrorCategory { + fn from_serde(category: serde_json::error::Category) -> Self { + match category { + serde_json::error::Category::Io => Self::Io, + serde_json::error::Category::Syntax => Self::Syntax, + serde_json::error::Category::Data => Self::Data, + serde_json::error::Category::Eof => Self::Eof, + } + } + + fn label(self) -> &'static str { + match self { + Self::Io => "I/O", + Self::Syntax => "syntax", + Self::Data => "data", + Self::Eof => "end-of-file", + } + } +} + +#[derive(Debug)] +struct ArtifactDecodeError { + category: JsonErrorCategory, + line: usize, + column: usize, +} + +impl ArtifactDecodeError { + fn from_serde(error: serde_json::Error) -> Self { + Self { + category: JsonErrorCategory::from_serde(error.classify()), + line: error.line(), + column: error.column(), + } + } +} + +impl fmt::Display for ArtifactDecodeError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + formatter, + "JSON {} error at line {}, column {}", + self.category.label(), + self.line, + self.column + ) + } +} + +impl Error for ArtifactDecodeError {} + +#[derive(Debug)] +pub(crate) struct ArtifactError { + path: PathBuf, + operation: &'static str, + record_index: usize, + source: Box, +} + +impl ArtifactError { + fn new( + path: &Path, operation: &'static str, record_index: usize, source: impl Error + Send + Sync + 'static, + ) -> Self { + Self { + path: path.to_owned(), + operation, + record_index, + source: Box::new(source), + } + } + + fn decode(path: &Path, record_index: usize, source: serde_json::Error) -> Self { + Self::new(path, "decode", record_index, ArtifactDecodeError::from_serde(source)) + } +} + +impl fmt::Display for ArtifactError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + formatter, + "{}: failed to {} record {}: {}", + self.path.display(), + self.operation, + self.record_index, + self.source + ) + } +} + +impl Error for ArtifactError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + Some(self.source.as_ref()) + } +} + +pub(crate) fn for_each_record(path: &Path, mut consume: impl FnMut(AgentContextRecord)) -> Result<(), ArtifactError> { + let file = File::open(path).map_err(|error| ArtifactError::new(path, "open", 0, error))?; + let mut compressed_input = BufReader::new(file); + let is_compressed = compressed_input + .fill_buf() + .map_err(|error| ArtifactError::new(path, "inspect compression magic for", 0, error))? + .starts_with(ZSTD_MAGIC); + + if is_compressed { + let decoder = zstd::stream::read::Decoder::with_buffer(compressed_input) + .map_err(|error| ArtifactError::new(path, "initialize zstd decoder for", 0, error))?; + decode_records(path, BufReader::new(decoder), &mut consume) + } else { + decode_records(path, compressed_input, &mut consume) + } +} + +fn decode_records( + path: &Path, reader: R, consume: &mut impl FnMut(AgentContextRecord), +) -> Result<(), ArtifactError> { + let records = serde_json::Deserializer::from_reader(reader).into_iter::(); + for (index, record) in records.enumerate() { + let record_index = index + 1; + let record = record.map_err(|error| ArtifactError::decode(path, record_index, error))?; + consume(record); + } + Ok(()) +} + +#[cfg(test)] +#[path = "artifact_tests.rs"] +mod tests; diff --git a/bin/agent-data-plane/src/dogstatsd_contexts/artifact_tests.rs b/bin/agent-data-plane/src/dogstatsd_contexts/artifact_tests.rs new file mode 100644 index 0000000000..c6f09d5cf0 --- /dev/null +++ b/bin/agent-data-plane/src/dogstatsd_contexts/artifact_tests.rs @@ -0,0 +1,497 @@ +use std::fs; +use std::io::{self, Write}; +use std::path::PathBuf; + +use saluki_components::transforms::{AggregateContextSnapshotEntry, AggregateMetricType}; +use saluki_context::{ + tags::{Tag, TagSet}, + Context, +}; +use serde_json::json; +use stringtheory::MetaString; + +use super::{ + for_each_record, publish_context_dump, write_snapshot_records, AgentContextRecord, ArtifactError, + CONTEXT_DUMP_FILENAME, +}; +use crate::dogstatsd_contexts::read_report; + +impl ArtifactError { + fn path(&self) -> &std::path::Path { + &self.path + } + + fn operation(&self) -> &'static str { + self.operation + } + + fn record_index(&self) -> usize { + self.record_index + } +} + +fn collect_records(path: &std::path::Path) -> Result, ArtifactError> { + let mut records = Vec::new(); + for_each_record(path, |record| records.push(record))?; + Ok(records) +} + +const PLAIN_FIXTURE_BYTES: &[u8] = include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/dogstatsd_contexts_agent.ndjson" +)); +const COMPRESSED_FIXTURE_BYTES: &[u8] = include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/dogstatsd_contexts_agent.ndjson.zstd" +)); +const PLAIN_FIXTURE_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/dogstatsd_contexts_agent.ndjson" +); +const COMPRESSED_FIXTURE_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/dogstatsd_contexts_agent.ndjson.zstd" +); + +// These fixtures and expectations mirror ContextDebugRepr and the `dogstatsd top` command at Datadog Agent +// commit 9de2a8371cf8d95794f238939709491907893288. +fn expected_fixture_records() -> Vec { + vec![ + AgentContextRecord { + name: "z.metric".to_owned(), + host: "node-a".to_owned(), + metric_type: "Gauge".to_owned(), + tagger_tags: vec!["pod_name:web-a".to_owned()], + metric_tags: vec![ + "env:prod".to_owned(), + "image:registry/repo:v1".to_owned(), + "bare".to_owned(), + ], + no_index: false, + source: 1, + }, + AgentContextRecord { + name: "a.metric".to_owned(), + host: "node-a".to_owned(), + metric_type: "Counter".to_owned(), + tagger_tags: vec!["pod_name:web-a".to_owned()], + metric_tags: vec!["env:prod".to_owned(), "service:web".to_owned()], + no_index: false, + source: 1, + }, + AgentContextRecord { + name: "a.metric".to_owned(), + host: "node-b".to_owned(), + metric_type: "Counter".to_owned(), + tagger_tags: vec!["pod_name:web-b".to_owned()], + metric_tags: vec!["env:prod".to_owned(), "service:web".to_owned()], + no_index: false, + source: 1, + }, + AgentContextRecord { + name: "a.metric".to_owned(), + host: "node-c".to_owned(), + metric_type: "Counter".to_owned(), + tagger_tags: vec!["pod_name:web-c".to_owned()], + metric_tags: vec!["env:staging".to_owned(), "service:web".to_owned()], + no_index: false, + source: 1, + }, + ] +} + +#[test] +fn streams_exact_agent_schema_for_every_metric_type() { + let cases = [ + (AggregateMetricType::Counter, None, "Counter"), + (AggregateMetricType::Rate, None, "Rate"), + (AggregateMetricType::Gauge, None, "Gauge"), + (AggregateMetricType::Set, None, "Set"), + (AggregateMetricType::Distribution, None, "Distribution"), + (AggregateMetricType::Histogram, None, "Histogram"), + (AggregateMetricType::Histogram, Some("second"), "Histogram"), + (AggregateMetricType::Histogram, Some("millisecond"), "Historate"), + ]; + let snapshot: Vec<_> = cases + .iter() + .enumerate() + .map(|(index, (metric_type, unit, _))| { + snapshot_entry( + Box::leak(format!("metric.{index}").into_boxed_str()), + (index == 0).then_some("node-a"), + *metric_type, + unit.unwrap_or_default(), + &["client:value", "bare"], + &["origin:value"], + ) + }) + .collect(); + let mut output = Vec::new(); + + write_snapshot_records(&mut output, &snapshot, PathBuf::from("plain-stream").as_path()) + .expect("snapshot should serialize"); + + let output = String::from_utf8(output).expect("serialized output should be UTF-8"); + let lines: Vec<_> = output.lines().collect(); + assert_eq!(lines.len(), cases.len()); + assert_eq!(output.bytes().filter(|byte| *byte == b'\n').count(), cases.len()); + assert!(output.ends_with('\n')); + for (index, ((_, _, expected_type), line)) in cases.iter().zip(&lines).enumerate() { + let expected_host = if index == 0 { "node-a" } else { "" }; + let expected = json!({ + "Name": format!("metric.{index}"), + "Host": expected_host, + "Type": expected_type, + "TaggerTags": ["origin:value"], + "MetricTags": ["client:value", "bare"], + "NoIndex": false, + "Source": 1, + }); + assert_eq!(serde_json::from_str::(line).unwrap(), expected); + assert_eq!( + *line, + format!( + "{{\"Name\":\"metric.{index}\",\"Host\":\"{expected_host}\",\"Type\":\"{expected_type}\",\"TaggerTags\":[\"origin:value\"],\"MetricTags\":[\"client:value\",\"bare\"],\"NoIndex\":false,\"Source\":1}}" + ) + ); + } +} + +#[test] +fn plain_stream_round_trips_through_reader_and_report() { + let snapshot = vec![ + snapshot_entry( + "shared.metric", + Some("node-a"), + AggregateMetricType::Gauge, + "", + &["env:prod", "service:web"], + &["pod_name:web-a"], + ), + snapshot_entry( + "shared.metric", + None, + AggregateMetricType::Counter, + "", + &["env:staging", "service:web"], + &["pod_name:web-b"], + ), + ]; + let artifact = tempfile::NamedTempFile::new().expect("temporary artifact should be created"); + let mut file = artifact.reopen().expect("temporary artifact should reopen"); + + write_snapshot_records(&mut file, &snapshot, artifact.path()).expect("snapshot should serialize"); + file.flush().expect("plain stream should flush"); + + assert_eq!( + collect_records(artifact.path()).expect("plain stream should decode"), + vec![ + AgentContextRecord { + name: "shared.metric".to_owned(), + host: "node-a".to_owned(), + metric_type: "Gauge".to_owned(), + tagger_tags: vec!["pod_name:web-a".to_owned()], + metric_tags: vec!["env:prod".to_owned(), "service:web".to_owned()], + no_index: false, + source: 1, + }, + AgentContextRecord { + name: "shared.metric".to_owned(), + host: String::new(), + metric_type: "Counter".to_owned(), + tagger_tags: vec!["pod_name:web-b".to_owned()], + metric_tags: vec!["env:staging".to_owned(), "service:web".to_owned()], + no_index: false, + source: 1, + }, + ] + ); + assert_eq!( + read_report(artifact.path()) + .expect("plain stream report should build") + .render(10, 10), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 2\tshared.metric\t(2 env, 1 service)\n", + ) + ); +} + +#[test] +fn surfaces_a_partial_writer_failure_with_record_and_path_context() { + let snapshot = [snapshot_entry( + "write.failure", + None, + AggregateMetricType::Gauge, + "", + &["env:test"], + &[], + )]; + let path = PathBuf::from("/run/test/dogstatsd_contexts.json.zstd"); + let mut writer = FailAfterBytes::new(12); + + let error = write_snapshot_records(&mut writer, &snapshot, &path).expect_err("write failure should surface"); + + assert_eq!(writer.written.len(), 12); + assert!(error.to_string().contains("serialize DogStatsD context dump record 1")); + assert!(error.to_string().contains(&path.display().to_string())); + assert!(format!("{error:#}").contains("injected write failure")); +} + +#[test] +fn publishes_a_decodable_mode_0600_dump_and_replaces_the_canonical_artifact() { + let run_directory = tempfile::tempdir().expect("temporary run directory should be created"); + let target = run_directory.path().join(CONTEXT_DUMP_FILENAME); + fs::write(&target, b"old canonical artifact").expect("old canonical artifact should be written"); + let snapshot = vec![ + snapshot_entry( + "published.metric", + Some("node-a"), + AggregateMetricType::Histogram, + "millisecond", + &["env:prod"], + &["pod_name:web-a"], + ), + snapshot_entry( + "published.metric", + None, + AggregateMetricType::Distribution, + "", + &["env:staging"], + &["pod_name:web-b"], + ), + ]; + + let published = publish_context_dump(run_directory.path(), &snapshot).expect("snapshot should publish"); + + assert_eq!(published, target); + assert_eq!( + collect_records(&target) + .expect("published artifact should decode") + .len(), + 2 + ); + assert_eq!( + read_report(&target) + .expect("published report should build") + .render(10, 10), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 2\tpublished.metric\t(2 env)\n", + ) + ); + assert_eq!( + directory_entries(run_directory.path()), + vec![CONTEXT_DUMP_FILENAME.to_owned()] + ); + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + + assert_eq!(fs::metadata(&target).unwrap().permissions().mode() & 0o777, 0o600); + } +} + +#[test] +fn rejects_empty_missing_and_non_directory_run_paths_with_target_and_operation_context() { + let temporary_directory = tempfile::tempdir().expect("temporary directory should be created"); + let missing = temporary_directory.path().join("missing"); + let not_directory = temporary_directory.path().join("not-a-directory"); + fs::write(¬_directory, b"file").expect("non-directory fixture should be written"); + let cases = [ + (PathBuf::new(), "validate run path"), + (missing, "create temporary"), + (not_directory, "create temporary"), + ]; + + for (run_path, expected_operation) in cases { + let error = publish_context_dump(&run_path, &[]).expect_err("invalid run path should fail"); + let message = format!("{error:#}"); + assert!(message.contains(CONTEXT_DUMP_FILENAME), "{message}"); + assert!(message.contains(expected_operation), "{message}"); + } +} + +#[test] +fn persist_failure_is_contextual_and_removes_the_temporary_file() { + let run_directory = tempfile::tempdir().expect("temporary run directory should be created"); + let target = run_directory.path().join(CONTEXT_DUMP_FILENAME); + fs::create_dir(&target).expect("directory should block file persistence"); + + let error = publish_context_dump(run_directory.path(), &[]).expect_err("persistence should fail"); + + let message = format!("{error:#}"); + assert!(message.contains("atomically replace"), "{message}"); + assert!(message.contains(&target.display().to_string()), "{message}"); + assert_eq!( + directory_entries(run_directory.path()), + vec![CONTEXT_DUMP_FILENAME.to_owned()] + ); +} + +fn directory_entries(path: &std::path::Path) -> Vec { + let mut entries: Vec<_> = fs::read_dir(path) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .collect(); + entries.sort_unstable(); + entries +} + +struct FailAfterBytes { + remaining: usize, + written: Vec, +} + +impl FailAfterBytes { + fn new(remaining: usize) -> Self { + Self { + remaining, + written: Vec::new(), + } + } +} + +impl io::Write for FailAfterBytes { + fn write(&mut self, bytes: &[u8]) -> io::Result { + if self.remaining == 0 { + return Err(io::Error::other("injected write failure")); + } + let written = self.remaining.min(bytes.len()); + self.written.extend_from_slice(&bytes[..written]); + self.remaining -= written; + Ok(written) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +fn snapshot_entry( + name: &'static str, host: Option<&'static str>, metric_type: AggregateMetricType, unit: &'static str, + metric_tags: &[&'static str], origin_tags: &[&'static str], +) -> AggregateContextSnapshotEntry { + let context = Context::from_static_parts(name, metric_tags) + .with_host(host.map(MetaString::from_static)) + .with_origin_tags(origin_tags.iter().map(|tag| Tag::from_static(tag)).collect::()); + AggregateContextSnapshotEntry::for_test(context, metric_type, MetaString::from_static(unit)) +} + +#[test] +fn detects_plain_and_compressed_artifacts_by_magic_bytes() { + let temporary_directory = tempfile::tempdir().expect("temporary directory should be created"); + let compressed_without_suffix = temporary_directory.path().join("compressed-contexts"); + let plain_with_compressed_suffix = temporary_directory.path().join("plain-contexts.zstd"); + fs::write(&compressed_without_suffix, COMPRESSED_FIXTURE_BYTES).expect("compressed fixture should be copied"); + fs::write(&plain_with_compressed_suffix, PLAIN_FIXTURE_BYTES).expect("plain fixture should be copied"); + + let cases = [ + PathBuf::from(PLAIN_FIXTURE_PATH), + PathBuf::from(COMPRESSED_FIXTURE_PATH), + compressed_without_suffix, + plain_with_compressed_suffix, + ]; + let expected = expected_fixture_records(); + + for path in cases { + assert_eq!( + collect_records(&path).expect("artifact should decode"), + expected, + "{path:?}" + ); + } +} + +#[test] +fn checked_in_compressed_fixture_expands_to_the_plain_fixture_exactly() { + let decompressed = + zstd::stream::decode_all(COMPRESSED_FIXTURE_BYTES).expect("checked-in compressed fixture should decompress"); + + assert_eq!(decompressed, PLAIN_FIXTURE_BYTES); + assert_eq!(PLAIN_FIXTURE_BYTES.last(), Some(&b'\n')); +} + +#[test] +fn rejects_invalid_artifacts_with_path_operation_and_record_index() { + let temporary_directory = tempfile::tempdir().expect("temporary directory should be created"); + let valid_record = concat!( + "{\"Name\":\"metric\",\"Host\":\"host\",\"Type\":\"Gauge\",", + "\"TaggerTags\":[],\"MetricTags\":[],\"NoIndex\":false,\"Source\":1}" + ); + let truncated_length = COMPRESSED_FIXTURE_BYTES.len() - 4; + let cases = [ + InvalidArtifactCase { + filename: "array.ndjson", + bytes: format!("[{valid_record}]").into_bytes(), + expected_record_index: 1, + }, + InvalidArtifactCase { + filename: "missing-field.ndjson", + bytes: br#"{"Name":"metric"}"#.to_vec(), + expected_record_index: 1, + }, + InvalidArtifactCase { + filename: "malformed-trailing.ndjson", + bytes: format!("{valid_record}\nnot-json").into_bytes(), + expected_record_index: 2, + }, + InvalidArtifactCase { + filename: "truncated.zstd", + bytes: COMPRESSED_FIXTURE_BYTES[..truncated_length].to_vec(), + expected_record_index: 5, + }, + ]; + + for case in cases { + let path = temporary_directory.path().join(case.filename); + fs::write(&path, case.bytes).expect("invalid artifact should be written"); + + let error = collect_records(&path).expect_err("invalid artifact should be rejected"); + + assert_eq!(error.path(), path.as_path(), "{}", case.filename); + assert_eq!(error.operation(), "decode", "{}", case.filename); + assert_eq!(error.record_index(), case.expected_record_index, "{}", case.filename); + assert!(error.to_string().contains(&path.display().to_string()), "{error}"); + assert!( + error + .to_string() + .contains(&format!("decode record {}", case.expected_record_index)), + "{error}" + ); + assert!(error.to_string().contains("JSON"), "{error}"); + assert!(error.to_string().contains("line"), "{error}"); + assert!(error.to_string().contains("column"), "{error}"); + } +} + +struct InvalidArtifactCase { + filename: &'static str, + bytes: Vec, + expected_record_index: usize, +} + +#[test] +fn ignores_unknown_fields_but_requires_every_agent_field() { + let artifact = tempfile::NamedTempFile::new().expect("temporary artifact should be created"); + let record = concat!( + "{\"Name\":\"metric\",\"Host\":\"host\",\"Type\":\"Gauge\",", + "\"TaggerTags\":[],\"MetricTags\":[\"env:prod\"],\"NoIndex\":false,\"Source\":7,", + "\"FutureField\":\"ignored\"}\n" + ); + fs::write(artifact.path(), record).expect("artifact should be written"); + + assert_eq!( + collect_records(artifact.path()).expect("record with unknown field should decode"), + vec![AgentContextRecord { + name: "metric".to_owned(), + host: "host".to_owned(), + metric_type: "Gauge".to_owned(), + tagger_tags: Vec::new(), + metric_tags: vec!["env:prod".to_owned()], + no_index: false, + source: 7, + }] + ); +} diff --git a/bin/agent-data-plane/src/dogstatsd_contexts/mod.rs b/bin/agent-data-plane/src/dogstatsd_contexts/mod.rs new file mode 100644 index 0000000000..7ed7581e5e --- /dev/null +++ b/bin/agent-data-plane/src/dogstatsd_contexts/mod.rs @@ -0,0 +1,55 @@ +use std::path::Path; + +use saluki_error::GenericError; + +pub(crate) use self::api::DogStatsDContextDumpAPIHandler; +pub(crate) use self::artifact::{for_each_record, publish_context_dump}; +pub(crate) use self::report::ContextReport; + +mod api; +mod artifact; +mod report; + +/// Agent-compatible command API route for requesting a DogStatsD context dump. +pub(crate) const CONTEXT_DUMP_ROUTE: &str = "/agent/dogstatsd-contexts-dump"; + +pub(crate) fn read_report(path: &Path) -> Result { + let mut report = ContextReport::new(); + for_each_record(path, |record| report.ingest(record))?; + Ok(report) +} + +#[cfg(test)] +mod tests { + use std::path::Path; + + use super::read_report; + + const HEADING: &str = " Contexts\tMetric name\t(number of unique values for each tag)\n"; + const PLAIN_FIXTURE: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/dogstatsd_contexts_agent.ndjson" + ); + + #[test] + fn renders_agent_fixture_as_a_golden_report() { + let report = read_report(Path::new(PLAIN_FIXTURE)).expect("fixture should decode"); + + assert_eq!( + report.render(usize::MAX, usize::MAX), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 3\ta.metric\t(2 env, 1 service)\n", + " 1\tz.metric\t(1 bare, 1 env, 1 image)\n", + ) + ); + } + + #[test] + fn empty_artifact_renders_only_the_heading() { + let artifact = tempfile::NamedTempFile::new().expect("temporary artifact should be created"); + let report = read_report(artifact.path()).expect("empty artifact should decode"); + + assert_eq!(report.render(10, 10), HEADING); + } +} diff --git a/bin/agent-data-plane/src/dogstatsd_contexts/report.rs b/bin/agent-data-plane/src/dogstatsd_contexts/report.rs new file mode 100644 index 0000000000..720ac1e4e1 --- /dev/null +++ b/bin/agent-data-plane/src/dogstatsd_contexts/report.rs @@ -0,0 +1,332 @@ +use std::collections::{HashMap, HashSet}; +use std::fmt::Write as _; + +use super::artifact::AgentContextRecord; + +const HEADING: &str = " Contexts\tMetric name\t(number of unique values for each tag)\n"; + +#[derive(Default)] +pub(crate) struct ContextReport { + metrics: HashMap, +} + +impl ContextReport { + pub(super) fn new() -> Self { + Self::default() + } + + pub(super) fn ingest(&mut self, record: AgentContextRecord) { + let AgentContextRecord { name, metric_tags, .. } = record; + let summary = self.metrics.entry(name).or_default(); + summary.context_count += 1; + summary.metric_tags.extend(metric_tags); + } + + pub(crate) fn render(&self, metric_limit: usize, tag_limit: usize) -> String { + let mut output = String::from(HEADING); + let mut metrics: Vec<_> = self.metrics.iter().collect(); + metrics.sort_unstable_by(|(left_name, left), (right_name, right)| { + right + .context_count + .cmp(&left.context_count) + .then_with(|| left_name.cmp(right_name)) + }); + + let visible_count = visible_item_count(metrics.len(), metric_limit); + for (name, summary) in &metrics[..visible_count] { + write!(output, " {:>10}\t{}\t(", summary.context_count, name).expect("writing to a String should not fail"); + summary.write_tags(&mut output, tag_limit); + output.push_str(")\n"); + } + + let hidden_metrics = &metrics[visible_count..]; + if !hidden_metrics.is_empty() { + let hidden_context_count: usize = hidden_metrics.iter().map(|(_, summary)| summary.context_count).sum(); + writeln!( + output, + " {hidden_context_count:>10}\t(other {} metrics)", + hidden_metrics.len() + ) + .expect("writing to a String should not fail"); + } + + output + } +} + +#[derive(Default)] +struct MetricSummary { + context_count: usize, + metric_tags: HashSet, +} + +impl MetricSummary { + fn write_tags(&self, output: &mut String, tag_limit: usize) { + let mut cardinalities: HashMap<&str, usize> = HashMap::new(); + for tag in &self.metric_tags { + let key = tag.split_once(':').map_or(tag.as_str(), |(key, _)| key); + *cardinalities.entry(key).or_default() += 1; + } + + let mut tags: Vec<_> = cardinalities.into_iter().collect(); + tags.sort_unstable_by(|(left_key, left_count), (right_key, right_count)| { + right_count.cmp(left_count).then_with(|| left_key.cmp(right_key)) + }); + + let visible_count = visible_item_count(tags.len(), tag_limit); + for (index, (key, cardinality)) in tags[..visible_count].iter().enumerate() { + if index > 0 { + output.push_str(", "); + } + write!(output, "{cardinality} {key}").expect("writing to a String should not fail"); + } + + let hidden_tags = &tags[visible_count..]; + if !hidden_tags.is_empty() { + let hidden_value_count: usize = hidden_tags.iter().map(|(_, cardinality)| cardinality).sum(); + if visible_count > 0 { + output.push_str(", "); + } + write!( + output, + "{hidden_value_count} values in {} other tags", + hidden_tags.len() + ) + .expect("writing to a String should not fail"); + } + } +} + +fn visible_item_count(item_count: usize, limit: usize) -> usize { + if item_count.saturating_sub(limit) > 1 { + limit + } else { + item_count + } +} + +#[cfg(test)] +mod tests { + use super::ContextReport; + use crate::dogstatsd_contexts::artifact::AgentContextRecord; + + const HEADING: &str = " Contexts\tMetric name\t(number of unique values for each tag)\n"; + + #[test] + fn sorts_metric_and_tag_ties_lexically() { + let mut report = ContextReport::new(); + report.ingest(record("beta", "host", &[], &["z:value", "a:value"])); + report.ingest(record("alpha", "host", &[], &["z:value", "a:value"])); + + assert_eq!( + report.render(10, 10), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 1\talpha\t(1 a, 1 z)\n", + " 1\tbeta\t(1 a, 1 z)\n", + ) + ); + } + + #[test] + fn deduplicates_full_metric_tags_before_calculating_cardinality() { + let mut report = ContextReport::new(); + report.ingest(record( + "metric", + "host-a", + &[], + &["env:prod", "env:prod", "env:staging"], + )); + report.ingest(record("metric", "host-b", &[], &["env:prod"])); + + assert_eq!( + report.render(10, 10), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 2\tmetric\t(2 env)\n", + ) + ); + } + + #[test] + fn counts_host_and_origin_variants_but_excludes_tagger_tags() { + let mut report = ContextReport::new(); + report.ingest(record("metric", "host-a", &["pod_name:a"], &["env:prod"])); + + let mut second = record("metric", "host-b", &["pod_name:b"], &["env:prod"]); + second.metric_type = "Counter".to_owned(); + second.no_index = true; + second.source = 99; + report.ingest(second); + + report.ingest(record("metric", "host-c", &["pod_name:c"], &["env:prod"])); + + assert_eq!( + report.render(10, 10), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 3\tmetric\t(1 env)\n", + ) + ); + } + + #[test] + fn uses_the_whole_colonless_tag_or_the_prefix_before_only_the_first_colon() { + let mut report = ContextReport::new(); + report.ingest(record( + "metric", + "host", + &[], + &["bare", "image:registry/repo:v1", "image:registry/repo:v2"], + )); + + assert_eq!( + report.render(10, 10), + concat!( + " Contexts\tMetric name\t(number of unique values for each tag)\n", + " 1\tmetric\t(2 image, 1 bare)\n", + ) + ); + } + + #[test] + fn applies_agent_metric_limit_plus_one_rule_without_overflow() { + let cases = [ + MetricLimitCase { + name: "zero limit shows its one allowed extra item", + counts: &[("a", 1)], + limit: 0, + expected_body: " 1\ta\t()\n", + }, + MetricLimitCase { + name: "zero limit summarizes two items", + counts: &[("a", 2), ("b", 1)], + limit: 0, + expected_body: " 3\t(other 2 metrics)\n", + }, + MetricLimitCase { + name: "item count equal to limit is shown", + counts: &[("a", 2), ("b", 1)], + limit: 2, + expected_body: " 2\ta\t()\n 1\tb\t()\n", + }, + MetricLimitCase { + name: "limit plus one items are all shown", + counts: &[("a", 2), ("b", 1)], + limit: 1, + expected_body: " 2\ta\t()\n 1\tb\t()\n", + }, + MetricLimitCase { + name: "limit plus two items hide the remainder and sum their contexts", + counts: &[("a", 3), ("b", 2), ("c", 1)], + limit: 1, + expected_body: " 3\ta\t()\n 3\t(other 2 metrics)\n", + }, + MetricLimitCase { + name: "maximum limit does not overflow", + counts: &[("a", 3), ("b", 2), ("c", 1)], + limit: usize::MAX, + expected_body: " 3\ta\t()\n 2\tb\t()\n 1\tc\t()\n", + }, + ]; + + for case in cases { + let report = report_with_counts(case.counts); + assert_eq!( + report.render(case.limit, usize::MAX), + format!("{HEADING}{}", case.expected_body), + "{}", + case.name + ); + } + } + + #[test] + fn applies_agent_tag_limit_plus_one_rule_without_overflow() { + let cases = [ + TagLimitCase { + name: "zero limit shows its one allowed extra tag", + tags: &["a:1"], + limit: 0, + expected_tags: "1 a", + }, + TagLimitCase { + name: "zero limit summarizes two tags and their values", + tags: &["a:1", "a:2", "b:1"], + limit: 0, + expected_tags: "3 values in 2 other tags", + }, + TagLimitCase { + name: "tag count equal to limit is shown", + tags: &["a:1", "b:1"], + limit: 2, + expected_tags: "1 a, 1 b", + }, + TagLimitCase { + name: "limit plus one tags are all shown", + tags: &["a:1", "b:1"], + limit: 1, + expected_tags: "1 a, 1 b", + }, + TagLimitCase { + name: "limit plus two tags hide the remainder and sum their values", + tags: &["a:1", "a:2", "a:3", "b:1", "b:2", "c:1"], + limit: 1, + expected_tags: "3 a, 3 values in 2 other tags", + }, + TagLimitCase { + name: "maximum limit does not overflow", + tags: &["a:1", "a:2", "a:3", "b:1", "b:2", "c:1"], + limit: usize::MAX, + expected_tags: "3 a, 2 b, 1 c", + }, + ]; + + for case in cases { + let mut report = ContextReport::new(); + report.ingest(record("metric", "host", &[], case.tags)); + assert_eq!( + report.render(usize::MAX, case.limit), + format!("{HEADING} 1\tmetric\t({})\n", case.expected_tags), + "{}", + case.name + ); + } + } + + fn report_with_counts(counts: &[(&str, usize)]) -> ContextReport { + let mut report = ContextReport::new(); + for (name, count) in counts { + for index in 0..*count { + report.ingest(record(name, &format!("host-{index}"), &[], &[])); + } + } + report + } + + fn record(name: &str, host: &str, tagger_tags: &[&str], metric_tags: &[&str]) -> AgentContextRecord { + AgentContextRecord { + name: name.to_owned(), + host: host.to_owned(), + metric_type: "Gauge".to_owned(), + tagger_tags: tagger_tags.iter().map(|tag| (*tag).to_owned()).collect(), + metric_tags: metric_tags.iter().map(|tag| (*tag).to_owned()).collect(), + no_index: false, + source: 1, + } + } + + struct MetricLimitCase { + name: &'static str, + counts: &'static [(&'static str, usize)], + limit: usize, + expected_body: &'static str, + } + + struct TagLimitCase { + name: &'static str, + tags: &'static [&'static str], + limit: usize, + expected_tags: &'static str, + } +} diff --git a/bin/agent-data-plane/src/internal/control_surfaces.rs b/bin/agent-data-plane/src/internal/control_surfaces.rs index 6bc45d4319..60cc559223 100644 --- a/bin/agent-data-plane/src/internal/control_surfaces.rs +++ b/bin/agent-data-plane/src/internal/control_surfaces.rs @@ -4,6 +4,8 @@ use saluki_components::{ sources::{DogStatsDCaptureAPIHandler, DogStatsDReplayAPIHandler}, }; +use crate::dogstatsd_contexts::DogStatsDContextDumpAPIHandler; + /// Combined set of control surfaces to expose from the privileged API endpoint. #[derive(Default)] pub struct TopologyControlSurfaces { @@ -34,6 +36,8 @@ pub struct DogStatsDControlSurface { pub(crate) capture_api_handler: DogStatsDCaptureAPIHandler, /// API handler for the `/dogstatsd/replay/session` endpoints. pub(crate) replay_api_handler: DogStatsDReplayAPIHandler, + /// API handler for the Agent-compatible `/agent/dogstatsd-contexts-dump` endpoint. + pub(crate) context_dump_api_handler: DogStatsDContextDumpAPIHandler, } impl DogStatsDControlSurface { @@ -42,5 +46,6 @@ impl DogStatsDControlSurface { .with_handler(self.stats_api_handler) .with_handler(self.capture_api_handler) .with_handler(self.replay_api_handler) + .with_handler(self.context_dump_api_handler) } } diff --git a/bin/agent-data-plane/src/main.rs b/bin/agent-data-plane/src/main.rs index 8cec1ad27e..be4e525a0b 100644 --- a/bin/agent-data-plane/src/main.rs +++ b/bin/agent-data-plane/src/main.rs @@ -29,6 +29,7 @@ use crate::internal::logging::LoggingConfigurationTranslator; mod components; mod config; +mod dogstatsd_contexts; mod internal; pub(crate) mod state; diff --git a/bin/agent-data-plane/tests/fixtures/dogstatsd_contexts_agent.ndjson b/bin/agent-data-plane/tests/fixtures/dogstatsd_contexts_agent.ndjson new file mode 100644 index 0000000000..0d91ad8024 --- /dev/null +++ b/bin/agent-data-plane/tests/fixtures/dogstatsd_contexts_agent.ndjson @@ -0,0 +1,4 @@ +{"Name":"z.metric","Host":"node-a","Type":"Gauge","TaggerTags":["pod_name:web-a"],"MetricTags":["env:prod","image:registry/repo:v1","bare"],"NoIndex":false,"Source":1} +{"Name":"a.metric","Host":"node-a","Type":"Counter","TaggerTags":["pod_name:web-a"],"MetricTags":["env:prod","service:web"],"NoIndex":false,"Source":1} +{"Name":"a.metric","Host":"node-b","Type":"Counter","TaggerTags":["pod_name:web-b"],"MetricTags":["env:prod","service:web"],"NoIndex":false,"Source":1} +{"Name":"a.metric","Host":"node-c","Type":"Counter","TaggerTags":["pod_name:web-c"],"MetricTags":["env:staging","service:web"],"NoIndex":false,"Source":1} diff --git a/bin/agent-data-plane/tests/fixtures/dogstatsd_contexts_agent.ndjson.zstd b/bin/agent-data-plane/tests/fixtures/dogstatsd_contexts_agent.ndjson.zstd new file mode 100644 index 0000000000..0dc488722b Binary files /dev/null and b/bin/agent-data-plane/tests/fixtures/dogstatsd_contexts_agent.ndjson.zstd differ diff --git a/docs/agent-data-plane/dogstatsd-top.md b/docs/agent-data-plane/dogstatsd-top.md new file mode 100644 index 0000000000..ef0b386cb7 --- /dev/null +++ b/docs/agent-data-plane/dogstatsd-top.md @@ -0,0 +1,111 @@ +# Inspect retained DogStatsD contexts + +Use `agent-data-plane dogstatsd top` to inspect which DogStatsD metric contexts Agent Data Plane (ADP) currently retains and to identify metric names with high cardinality. + +## Understand what `top` reports + +`top` snapshots the contexts currently retained by the `dsd_agg` aggregate transform. It does not measure packet or sample frequency, and it does not collect data over a time window. Use the separate `agent-data-plane dogstatsd stats` command when you need time-windowed sample counts and last-seen times. + +A context's full identity consists of its metric name, host, client tags, and origin tags. The report groups these contexts by metric name, so contexts that differ only by host or origin tags still increase the context count for that name. The displayed tag cardinalities include only client tags; they do not include hosts or origin tags. + +The report reflects Saluki's normal aggregate lifecycle. Contexts disappear after their aggregation state is flushed and no longer retained. Sparse counters can remain as idle counters while the aggregate transform emits zero values, then disappear according to the configured Saluki counter expiry. Do not interpret this behavior as an Agent-identical context time to live. + +## Request and analyze a dump + +Run `top` without `--path` to request a new dump and analyze it: + +```console +$ agent-data-plane dogstatsd top +Wrote /var/run/datadog/dogstatsd_contexts.json.zstd + Contexts Metric name (number of unique values for each tag) + 3 request.count (2 env, 1 service) +``` + +Online requests use the same privileged API client path and access-control boundary as existing ADP commands. The API creates and returns only the server-local path `/dogstatsd_contexts.json.zstd`. + +The API does not return the artifact contents. The CLI reads the returned path directly, so the CLI process and ADP must see the same filesystem and path. If ADP runs in a container, run the CLI with the same mount or use the copy workflow in the next section. The API does not provide a remote download endpoint. + +Use `--num-metrics` (`-m`) and `--num-tags` (`-t`) to change the default limits of 10 metric names and 5 tag keys: + +```console +$ agent-data-plane dogstatsd top --num-metrics 20 --num-tags 8 +$ agent-data-plane dogstatsd top -m 20 -t 8 +``` + +Limits must be non-negative integers, and `top` rejects extra positional arguments. + +The report orders metric names by descending context count and tag keys by descending unique-value count. It resolves ties lexically. When only one item exceeds a limit, the report displays that item; when two or more exceed it, the report replaces them with one remainder row or tag summary. + +To create the artifact without rendering a report, run: + +```console +$ agent-data-plane dogstatsd dump-contexts +Wrote /var/run/datadog/dogstatsd_contexts.json.zstd +``` + +`dump-contexts` takes no additional arguments. A successful request atomically replaces any prior file at the fixed path. The completed artifact remains there until an operator removes it. + +## Analyze a copied dump offline + +Wait for `top` or `dump-contexts` to report success, then securely copy the completed artifact to the analysis host. Analyze the copy by passing its path explicitly: + +```console +$ agent-data-plane dogstatsd top --path /secure/cases/incident-123/dogstatsd_contexts.json.zstd +``` + +You can combine `--path` (`-p`) with the metric and tag limits: + +```console +$ agent-data-plane dogstatsd top -p /secure/cases/incident-123/dogstatsd_contexts.json.zstd -m 25 -t 10 +``` + +Offline mode reads only the supplied file and does not contact ADP. `top` requires the `--path` value and does not fall back to a file in the current working directory. + +## Interpret context and tag counts + +Each artifact record represents one full retained context. The **Contexts** column counts records with the same metric name, including contexts distinguished only by host or origin tags. + +For the parenthesized tag summary, `top` deduplicates client tags across all records for the metric name. It treats the text before the first colon as the tag key and counts unique complete tag values. It treats a tag without a colon as its own key. For example, `(3 pod_name, 2 env)` means the retained contexts contain three distinct client tags whose key is `pod_name` and two whose key is `env`. It does not mean that ADP received three or two samples. + +Use the remainder row to account for hidden contexts. For example, `17 (other 4 metrics)` means the four hidden metric names contain 17 retained contexts in total. A tag remainder such as `12 values in 3 other tags` reports the summed unique-value counts for the hidden tag keys. + +## Protect and remove dump artifacts + +> [!WARNING] +> DogStatsD context dumps can contain metric names, hosts, client tags, and origin tags that reveal tenant or workload details. Treat every dump and copy as sensitive diagnostic data. + +On Unix, ADP creates the artifact with mode `0600`, owned by the ADP process owner. On Windows, the artifact inherits the configured `run_path` directory ACL, matching the Datadog Agent; secure a custom `run_path` for the ADP service identity and intended administrators before requesting a dump. Keep the resulting access restriction when you copy or store the file, use an encrypted and access-controlled transfer method, and avoid placing it in shared temporary directories. Remove the server-side artifact and every copied artifact manually when you finish the investigation. + +A later successful dump replaces the fixed server-side file but does not remove copies elsewhere. ADP does not apply a retention or cleanup policy to the completed file. + +## Understand snapshot consistency + +ADP currently has one `dsd_agg` owner. That owner constructs an internally consistent snapshot of its retained context map and then resumes processing. ADP serializes and compresses the snapshot after the owner resumes, so file creation does not hold the owner for the serialization work. + +Snapshot construction uses memory proportional to the number of retained contexts, or O(contexts), in addition to the aggregate state itself. Serialization also retains that snapshot until publication finishes. + +If ADP uses multiple aggregate owners in the future, each owner's snapshot remains internally consistent. The combined artifact would be a rolling snapshot because owners can respond at different moments; it would not represent one atomic instant across all owners. + +## Troubleshoot errors + +Use the status or file error in the CLI output to choose a response: + +- **Connection failure**: Verify that the CLI can reach the configured privileged API endpoint. +- **404 Not Found**: The context-dump route is absent when DogStatsD is disabled. Also verify that the CLI targets the intended ADP privileged API endpoint. +- **503 Service Unavailable**: The aggregate snapshot owner is not running or stopped before it responded. Wait for the topology and `dsd_agg` to become healthy, then retry. +- **504 Gateway Timeout**: The aggregate owner did not return a snapshot within the 30-second API snapshot deadline. Check ADP health and load before retrying. +- **500 Internal Server Error**: ADP could not publish the artifact. Verify that `run_path` is configured, exists as a directory, and is writable by the ADP process. An empty `run_path` fails; ADP does not fall back to the current working directory. +- **Missing or unreadable path**: An online request returns a path on ADP's filesystem. If the request succeeds but `top` cannot open that path, give the CLI access to the same filesystem or use `dump-contexts`, copy the completed file, and run offline with `--path`. +- **Malformed or truncated artifact**: A zstd initialization error or a JSON record error usually indicates an incomplete, corrupt, or unsupported file. Copy the artifact again after dump publication completes and verify the transfer before retrying. Pass the exact file path because offline mode has no current-working-directory fallback. + +## Artifact interoperability + +ADP writes a zstd-compressed stream of newline-delimited JSON (NDJSON) without a schema version. Every record uses the Agent diagnostic schema with these exact PascalCase fields: + +```json +{"Name":"request.count","Host":"node-a","Type":"Counter","TaggerTags":["pod_name:web-a"],"MetricTags":["env:prod","service:web"],"NoIndex":false,"Source":1} +``` + +`TaggerTags` contains origin tags and `MetricTags` contains client tags. The Datadog Agent and ADP `dogstatsd top --path` command-line clients can each read artifacts produced by either implementation. Saluki detects zstd compression from the file's magic bytes rather than its filename, and it also accepts plain NDJSON with the same records. + +This schema has no version field and serves as a diagnostic interoperability contract. It is not a supported metric ingestion API; do not send these records to a DogStatsD listener or depend on them as a general-purpose storage format. diff --git a/docs/agent-data-plane/index.md b/docs/agent-data-plane/index.md index 5dd73d700c..9ddc8c0605 100644 --- a/docs/agent-data-plane/index.md +++ b/docs/agent-data-plane/index.md @@ -9,3 +9,6 @@ This section of the documentation covers specific areas related to the developme code lives within the Saluki repository. It won't be relevant to those simply looking to contribute to or utilize Saluki for their own purposes, but is meant to live alongside the ADP codebase to ensure that all relevant information is quickly and easily accessible. + +Use [Inspect retained DogStatsD contexts](dogstatsd-top.md) to snapshot retained aggregation contexts, investigate +cardinality, and analyze sensitive dump artifacts offline. diff --git a/lib/saluki-components/Cargo.toml b/lib/saluki-components/Cargo.toml index 435166fb74..353e888a89 100644 --- a/lib/saluki-components/Cargo.toml +++ b/lib/saluki-components/Cargo.toml @@ -11,6 +11,7 @@ workspace = true [features] default = [] fips = ["saluki-io/fips"] +test-util = [] [dependencies] agent-data-plane-config = { workspace = true } diff --git a/lib/saluki-components/src/transforms/aggregate/mod.rs b/lib/saluki-components/src/transforms/aggregate/mod.rs index 0c56d500ec..0e162b9f50 100644 --- a/lib/saluki-components/src/transforms/aggregate/mod.rs +++ b/lib/saluki-components/src/transforms/aggregate/mod.rs @@ -1,5 +1,7 @@ use std::{ + future::pending, num::NonZeroU64, + sync::Mutex, time::{Duration, Instant}, }; @@ -17,13 +19,14 @@ use saluki_core::{ topology::{interconnect::BufferedDispatcher, OutputDefinition}, topology::{EventsBuffer, EventsDispatcher}, }; -use saluki_error::GenericError; +use saluki_error::{generic_error, GenericError}; use saluki_metrics::MetricsBuilder; use serde::Deserialize; use smallvec::SmallVec; use stringtheory::MetaString; use tokio::{ pin, select, + sync::{mpsc, oneshot}, time::{interval, interval_at}, }; use tracing::{debug, error, info, trace, warn}; @@ -35,6 +38,251 @@ mod config; use self::config::{HistogramConfiguration, HistogramStatistic}; const PASSTHROUGH_IDLE_FLUSH_CHECK_INTERVAL: Duration = Duration::from_secs(2); +const CONTEXT_SNAPSHOT_REQUEST_CHANNEL_CAPACITY: usize = 1; + +/// The shape of metric values retained by the aggregate transform. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AggregateMetricType { + /// Counter values. + Counter, + + /// Rate values. + Rate, + + /// Gauge values. + Gauge, + + /// Set values. + Set, + + /// Histogram values. + Histogram, + + /// Distribution values. + Distribution, +} + +impl From<&MetricValues> for AggregateMetricType { + fn from(values: &MetricValues) -> Self { + match values { + MetricValues::Counter(_) => Self::Counter, + MetricValues::Rate(_, _) => Self::Rate, + MetricValues::Gauge(_) => Self::Gauge, + MetricValues::Set(_) => Self::Set, + MetricValues::Histogram(_) => Self::Histogram, + MetricValues::Distribution(_) => Self::Distribution, + } + } +} + +/// A retained metric context and its small aggregation metadata. +/// +/// Cloning an entry shares the underlying context name and tags rather than copying their contents. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct AggregateContextSnapshotEntry { + context: Context, + metric_type: AggregateMetricType, + unit: MetaString, +} + +impl AggregateContextSnapshotEntry { + /// Returns the retained metric context. + pub fn context(&self) -> &Context { + &self.context + } + + /// Returns the shape of the retained metric values. + pub fn metric_type(&self) -> AggregateMetricType { + self.metric_type + } + + /// Returns the unit attached to the retained metric values, if one is set. + pub fn unit(&self) -> Option<&str> { + if self.unit.is_empty() { + None + } else { + Some(&self.unit) + } + } + + /// Creates a snapshot entry for downstream test and benchmark fixtures. + #[cfg(any(test, feature = "test-util"))] + pub fn for_test(context: Context, metric_type: AggregateMetricType, unit: MetaString) -> Self { + Self { + context, + metric_type, + unit, + } + } +} + +type AggregateContextSnapshot = Vec; +type AggregateContextSnapshotRequest = oneshot::Sender; +type AggregateContextSnapshotRequestReceiver = mpsc::Receiver; + +/// A handle for requesting retained-context snapshots from an aggregate transform. +/// +/// Snapshot construction runs on the aggregate owner task. The returned entries contain shared context handles and +/// small metric metadata, allowing callers to perform heavier processing after the owner resumes ingestion. +#[derive(Clone, Debug)] +pub struct AggregateContextSnapshotHandle { + requests: mpsc::Sender, +} + +impl AggregateContextSnapshotHandle { + /// Requests the aggregate transform's current retained contexts. + /// + /// The returned snapshot uses O(context count) memory. After delivery, the caller owns this memory, so callers that + /// retain snapshots must include their retained size in their own memory accounting. + /// + /// # Errors + /// + /// Returns an error if the aggregate owner is unavailable or stops before responding. + pub async fn snapshot(&self) -> Result, GenericError> { + let (response_tx, response_rx) = oneshot::channel(); + self.requests + .send(response_tx) + .await + .map_err(|_| generic_error!("aggregate context snapshot owner is unavailable"))?; + + response_rx + .await + .map_err(|_| generic_error!("aggregate context snapshot owner stopped before responding")) + } +} + +fn aggregate_context_snapshot_channel() -> (AggregateContextSnapshotHandle, AggregateContextSnapshotRequestReceiver) { + let (requests, receiver) = mpsc::channel(CONTEXT_SNAPSHOT_REQUEST_CHANNEL_CAPACITY); + (AggregateContextSnapshotHandle { requests }, receiver) +} + +#[inline] +fn send_context_snapshot_if_open( + response: AggregateContextSnapshotRequest, build_snapshot: impl FnOnce() -> AggregateContextSnapshot, +) { + if response.is_closed() { + return; + } + + let _ = response.send(build_snapshot()); +} + +/// An accepted retained-context snapshot request for test fixtures. +#[cfg(any(test, feature = "test-util"))] +pub struct AggregateContextSnapshotPendingResponse { + response: AggregateContextSnapshotRequest, +} + +#[cfg(any(test, feature = "test-util"))] +impl AggregateContextSnapshotPendingResponse { + /// Responds to the accepted snapshot request with the supplied entries. + /// + /// If the requester was canceled after the request was accepted, the response is discarded. + pub fn respond(self, snapshot: Vec) { + let _ = self.response.send(snapshot); + } +} + +/// A responder for retained-context snapshot test fixtures. +#[cfg(any(test, feature = "test-util"))] +pub struct AggregateContextSnapshotResponder { + receiver: AggregateContextSnapshotRequestReceiver, +} + +#[cfg(any(test, feature = "test-util"))] +impl AggregateContextSnapshotResponder { + /// Waits for one snapshot request and responds with the supplied entries. + /// + /// A canceled requester is treated as a successful no-op delivery. + /// + /// # Errors + /// + /// Returns an error if the request channel closes before a request arrives. + pub async fn respond(&mut self, snapshot: Vec) -> Result<(), GenericError> { + self.receive().await?.respond(snapshot); + Ok(()) + } + + /// Waits for one snapshot request and returns its pending response. + /// + /// The returned response lets tests deterministically control whether the owner responds, stops, or outlives a + /// canceled requester after accepting the request. + /// + /// # Errors + /// + /// Returns an error if the request channel closes before a request arrives. + pub async fn receive(&mut self) -> Result { + let response = self + .receiver + .recv() + .await + .ok_or_else(|| generic_error!("aggregate context snapshot request channel is closed"))?; + Ok(AggregateContextSnapshotPendingResponse { response }) + } + + /// Waits for one snapshot request and stops without responding. + /// + /// This accepts the request from the owner channel before dropping its one-shot response sender, allowing tests to + /// distinguish an owner that stops mid-request from an owner whose request channel is unavailable. + /// + /// # Errors + /// + /// Returns an error if the request channel closes before a request arrives. + pub async fn stop_after_receiving(&mut self) -> Result<(), GenericError> { + drop(self.receive().await?); + Ok(()) + } +} + +/// Creates a retained-context snapshot handle and owner-side responder for tests. +#[cfg(any(test, feature = "test-util"))] +pub fn aggregate_context_snapshot_channel_for_test( +) -> (AggregateContextSnapshotHandle, AggregateContextSnapshotResponder) { + let (handle, receiver) = aggregate_context_snapshot_channel(); + (handle, AggregateContextSnapshotResponder { receiver }) +} + +struct AggregateContextSnapshotChannel { + handle: AggregateContextSnapshotHandle, + receiver: Mutex>, +} + +impl AggregateContextSnapshotChannel { + fn take_receiver(&self) -> Result { + let mut receiver = self + .receiver + .lock() + .map_err(|_| generic_error!("aggregate context snapshot receiver lock is poisoned"))?; + receiver + .take() + .ok_or_else(|| generic_error!("aggregate context snapshot receiver has already been taken")) + } +} + +impl Default for AggregateContextSnapshotChannel { + fn default() -> Self { + let (handle, receiver) = aggregate_context_snapshot_channel(); + Self { + handle, + receiver: Mutex::new(Some(receiver)), + } + } +} + +#[cfg(test)] +impl std::fmt::Debug for AggregateContextSnapshotChannel { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AggregateContextSnapshotChannel") + .finish_non_exhaustive() + } +} + +#[cfg(test)] +impl PartialEq for AggregateContextSnapshotChannel { + fn eq(&self, _other: &Self) -> bool { + true + } +} const fn default_window_duration_seconds() -> NonZeroU64 { NonZeroU64::new(10).expect("not zero") @@ -185,6 +433,12 @@ pub struct AggregateConfiguration { /// distribution aggregation). #[serde(flatten)] hist_config: HistogramConfiguration, + + /// Owner-side channel used to coordinate retained-context snapshots. + /// + /// This runtime-only channel is never read from or written to configuration data. + #[serde(skip, default)] + context_snapshot_channel: AggregateContextSnapshotChannel, } impl AggregateConfiguration { @@ -204,13 +458,20 @@ impl AggregateConfiguration { passthrough_timestamped_metrics: default_passthrough_timestamped_metrics(), passthrough_idle_flush_timeout: default_passthrough_idle_flush_timeout(), hist_config: HistogramConfiguration::default(), + context_snapshot_channel: AggregateContextSnapshotChannel::default(), } } + + /// Returns a handle for requesting retained-context snapshots from the built aggregate transform. + pub fn context_snapshot_handle(&self) -> AggregateContextSnapshotHandle { + self.context_snapshot_channel.handle.clone() + } } #[async_trait] impl TransformBuilder for AggregateConfiguration { async fn build(&self, context: ComponentContext) -> Result, GenericError> { + let context_snapshot_requests = self.context_snapshot_channel.take_receiver()?; let metrics_builder = MetricsBuilder::from_component_context(&context); let telemetry = Telemetry::new(&metrics_builder); @@ -236,6 +497,7 @@ impl TransformBuilder for AggregateConfiguration { flush_open_windows: self.flush_open_windows, passthrough_batcher, passthrough_timestamped_metrics: self.passthrough_timestamped_metrics, + context_snapshot_requests: Some(context_snapshot_requests), })) } @@ -274,6 +536,12 @@ impl MemoryBounds for AggregateConfiguration { UsageExpr::struct_size::("aggregated metric"), ), UsageExpr::config("aggregate_context_limit", self.context_limit), + )) + // A snapshot is constructed while the aggregation state remains live, so its peak allocation is additive. + .with_expr(UsageExpr::product( + "retained context snapshot", + UsageExpr::struct_size::("snapshot entry"), + UsageExpr::config("aggregate_context_limit", self.context_limit), )); } } @@ -285,6 +553,7 @@ pub struct Aggregate { flush_open_windows: bool, passthrough_batcher: PassthroughBatcher, passthrough_timestamped_metrics: bool, + context_snapshot_requests: Option, } #[async_trait] @@ -345,6 +614,14 @@ impl Transform for Aggregate { } }, _ = passthrough_flush.tick() => self.passthrough_batcher.try_flush(context.dispatcher()).await, + snapshot_request = receive_context_snapshot_request(&mut self.context_snapshot_requests) => { + match snapshot_request { + Some(response) => { + send_context_snapshot_if_open(response, || self.state.snapshot_contexts()); + } + None => self.context_snapshot_requests = None, + } + }, maybe_events = context.events().next(), if !final_primary_flush => match maybe_events { Some(events) => { trace!(events_len = events.len(), "Received events."); @@ -415,6 +692,15 @@ impl Transform for Aggregate { } } +async fn receive_context_snapshot_request( + receiver: &mut Option, +) -> Option { + match receiver { + Some(receiver) => receiver.recv().await, + None => pending().await, + } +} + fn try_split_timestamped_values(mut metric: Metric) -> (Option, Option) { if metric.values().all_timestamped() { (Some(metric), None) @@ -568,6 +854,18 @@ impl AggregationState { self.contexts.is_empty() } + fn snapshot_contexts(&self) -> Vec { + let mut snapshot = Vec::with_capacity(self.contexts.len()); + for (context, aggregated) in &self.contexts { + snapshot.push(AggregateContextSnapshotEntry { + context: context.clone(), + metric_type: AggregateMetricType::from(&aggregated.values), + unit: aggregated.metadata.unit.clone(), + }); + } + snapshot + } + fn insert(&mut self, timestamp: u64, metric: Metric) -> bool { // If we haven't seen this context yet, and it would put us over the limit to insert it, then return early. if !self.contexts.contains_key(metric.context()) && self.contexts.len() >= self.context_limit { @@ -895,14 +1193,25 @@ const fn is_bucket_closed( // then we run it to make sure that we are always generating sequential timestamps for data points, etc. #[cfg(test)] mod tests { + use std::{cell::Cell, mem::size_of}; + use float_cmp::ApproxEqRatio as _; + use saluki_context::tags::{Tag, TagSet}; use saluki_core::{ - components::ComponentContext, - topology::{interconnect::Dispatcher, OutputName}, + accounting::{ComponentRegistry, MemoryLimiter}, + components::{ + destinations::{Destination, DestinationBuilder, DestinationContext}, + sources::{Source, SourceBuilder, SourceContext}, + ComponentContext, + }, + health::HealthRegistry, + runtime::Supervisor, + support::SubsystemIdentifier, + topology::{interconnect::Dispatcher, OutputDefinition, OutputName, TopologyBlueprint}, }; use saluki_metrics::test::TestRecorder; use stringtheory::MetaString; - use tokio::sync::mpsc; + use tokio::sync::{mpsc, oneshot}; use super::config::HistogramStatistic; use super::*; @@ -927,6 +1236,83 @@ mod tests { BUCKET_WIDTH_SECS.get() * (step + 1) } + struct ControlledMetricSource { + events: mpsc::Receiver, + } + + #[async_trait] + impl Source for ControlledMetricSource { + async fn run(mut self: Box, mut context: SourceContext) -> Result<(), GenericError> { + let shutdown = context.take_shutdown_handle(); + tokio::pin!(shutdown); + let mut events_open = true; + + loop { + select! { + _ = &mut shutdown => break, + maybe_event = self.events.recv(), if events_open => match maybe_event { + Some(event) => context.dispatcher().dispatch_one(event).await?, + None => events_open = false, + } + } + } + Ok(()) + } + } + + struct ControlledMetricSourceBuilder { + events: Mutex>>, + outputs: Vec>, + } + + #[async_trait] + impl SourceBuilder for ControlledMetricSourceBuilder { + fn outputs(&self) -> &[OutputDefinition] { + &self.outputs + } + + async fn build(&self, _context: ComponentContext) -> Result, GenericError> { + let events = self + .events + .lock() + .map_err(|_| generic_error!("controlled metric source receiver lock is poisoned"))? + .take() + .ok_or_else(|| generic_error!("controlled metric source receiver has already been taken"))?; + Ok(Box::new(ControlledMetricSource { events })) + } + } + + impl MemoryBounds for ControlledMetricSourceBuilder { + fn specify_bounds(&self, _builder: &mut MemoryBoundsBuilder) {} + } + + struct DrainingMetricDestination; + + #[async_trait] + impl Destination for DrainingMetricDestination { + async fn run(self: Box, mut context: DestinationContext) -> Result<(), GenericError> { + while context.events().next().await.is_some() {} + Ok(()) + } + } + + struct DrainingMetricDestinationBuilder; + + #[async_trait] + impl DestinationBuilder for DrainingMetricDestinationBuilder { + fn input_event_type(&self) -> EventType { + EventType::Metric + } + + async fn build(&self, _context: ComponentContext) -> Result, GenericError> { + Ok(Box::new(DrainingMetricDestination)) + } + } + + impl MemoryBounds for DrainingMetricDestinationBuilder { + fn specify_bounds(&self, _builder: &mut MemoryBoundsBuilder) {} + } + struct DispatcherReceiver { receiver: mpsc::Receiver, } @@ -1063,6 +1449,386 @@ mod tests { }; } + #[test] + fn aggregate_metric_type_matches_every_metric_shape() { + let cases = [ + (MetricValues::counter(1.0), AggregateMetricType::Counter), + ( + MetricValues::rate(1.0, Duration::from_secs(10)), + AggregateMetricType::Rate, + ), + (MetricValues::gauge(1.0), AggregateMetricType::Gauge), + (MetricValues::set("value"), AggregateMetricType::Set), + (MetricValues::histogram([1.0]), AggregateMetricType::Histogram), + ( + MetricValues::distribution(&[1.0][..]), + AggregateMetricType::Distribution, + ), + ]; + + for (values, expected) in cases { + assert_eq!(AggregateMetricType::from(&values), expected); + } + } + + #[test] + fn snapshot_contexts_preserves_full_context_shape_and_unit() { + let mut state = AggregationState::new( + BUCKET_WIDTH_SECS, + 10, + COUNTER_EXPIRE, + HistogramConfiguration::default(), + Telemetry::noop(), + ); + + let histogram_context = Context::from_static_parts("request.duration", &["env:prod"]) + .with_host(Some(MetaString::from_static("host-a"))) + .with_origin_tags(TagSet::from(Tag::from_static("container:one"))); + let gauge_context = Context::from_static_parts("request.duration", &["env:prod"]) + .with_host(Some(MetaString::from_static("host-b"))) + .with_origin_tags(TagSet::from(Tag::from_static("container:two"))); + let histogram = Metric::from_parts( + histogram_context.clone(), + MetricValues::histogram([12.0]), + MetricMetadata::default().with_unit(MetaString::from_static("millisecond")), + ); + let gauge = Metric::gauge(gauge_context.clone(), 2.0); + + assert!(state.insert(insert_ts(1), histogram)); + assert!(state.insert(insert_ts(1), gauge)); + + let mut snapshot = state.snapshot_contexts(); + assert_eq!(snapshot.len(), 2); + assert!(snapshot.capacity() >= state.contexts.len()); + snapshot.sort_by(|a, b| a.context().host().cmp(&b.context().host())); + + assert_eq!(snapshot[0].context(), &histogram_context); + assert_eq!(snapshot[0].context().tags().len(), 1); + assert_eq!(snapshot[0].context().origin_tags().len(), 1); + assert_eq!(snapshot[0].metric_type(), AggregateMetricType::Histogram); + assert_eq!(snapshot[0].unit(), Some("millisecond")); + + assert_eq!(snapshot[1].context(), &gauge_context); + assert_eq!(snapshot[1].metric_type(), AggregateMetricType::Gauge); + assert_eq!(snapshot[1].unit(), None); + } + + #[tokio::test] + async fn snapshot_contexts_follows_ordinary_context_lifecycle() { + let mut state = AggregationState::new( + BUCKET_WIDTH_SECS, + 10, + COUNTER_EXPIRE, + HistogramConfiguration::default(), + Telemetry::noop(), + ); + let context = Context::from_static_name("active.gauge"); + + assert!(state.insert(insert_ts(1), Metric::gauge(context.clone(), 1.0))); + let active_snapshot = state.snapshot_contexts(); + assert_eq!(active_snapshot.len(), 1); + assert_eq!(active_snapshot[0].context(), &context); + assert_eq!(active_snapshot[0].metric_type(), AggregateMetricType::Gauge); + + let _ = get_flushed_metrics(flush_ts(1), &mut state).await; + assert!(state.snapshot_contexts().is_empty()); + } + + #[tokio::test] + async fn snapshot_contexts_retains_idle_counter_until_expiry() { + let mut state = AggregationState::new( + BUCKET_WIDTH_SECS, + 10, + COUNTER_EXPIRE, + HistogramConfiguration::default(), + Telemetry::noop(), + ); + let context = Context::from_static_name("sparse.counter"); + + assert!(state.insert(insert_ts(1), Metric::counter(context.clone(), 1.0))); + let _ = get_flushed_metrics(flush_ts(1), &mut state).await; + let first_snapshot = state.snapshot_contexts(); + assert_eq!(first_snapshot.len(), 1); + assert_eq!(first_snapshot[0].context(), &context); + assert_eq!(first_snapshot[0].metric_type(), AggregateMetricType::Counter); + + let _ = get_flushed_metrics(flush_ts(2), &mut state).await; + let second_snapshot = state.snapshot_contexts(); + assert_eq!(second_snapshot.len(), 1); + assert_eq!(second_snapshot[0].context(), &context); + assert_eq!(second_snapshot[0].metric_type(), AggregateMetricType::Counter); + + let _ = get_flushed_metrics(flush_ts(3), &mut state).await; + assert!(state.snapshot_contexts().is_empty()); + } + + #[test] + fn canceled_snapshot_response_skips_snapshot_construction() { + let mut state = AggregationState::new( + BUCKET_WIDTH_SECS, + 10, + COUNTER_EXPIRE, + HistogramConfiguration::default(), + Telemetry::noop(), + ); + assert!(state.insert( + insert_ts(1), + Metric::gauge(Context::from_static_name("canceled.snapshot"), 1.0), + )); + let (response, receiver) = oneshot::channel(); + drop(receiver); + let snapshot_calls = Cell::new(0); + + send_context_snapshot_if_open(response, || { + snapshot_calls.set(snapshot_calls.get() + 1); + state.snapshot_contexts() + }); + + assert_eq!(snapshot_calls.get(), 0); + } + + #[test] + fn open_snapshot_response_constructs_once_and_returns_exact_entries() { + let mut state = AggregationState::new( + BUCKET_WIDTH_SECS, + 10, + COUNTER_EXPIRE, + HistogramConfiguration::default(), + Telemetry::noop(), + ); + let context = Context::from_static_name("open.snapshot"); + assert!(state.insert(insert_ts(1), Metric::gauge(context.clone(), 1.0))); + let expected = vec![AggregateContextSnapshotEntry { + context, + metric_type: AggregateMetricType::Gauge, + unit: MetaString::empty(), + }]; + let (response, mut receiver) = oneshot::channel(); + let snapshot_calls = Cell::new(0); + + send_context_snapshot_if_open(response, || { + snapshot_calls.set(snapshot_calls.get() + 1); + state.snapshot_contexts() + }); + + assert_eq!(snapshot_calls.get(), 1); + assert_eq!( + receiver.try_recv().expect("open requester should receive a snapshot"), + expected + ); + } + + #[test] + fn aggregate_memory_bounds_include_peak_context_snapshot() { + let mut config = AggregateConfiguration::with_defaults(); + config.context_limit = 17; + let registry = ComponentRegistry::default(); + config.specify_bounds(&mut registry.bounds_builder(&SubsystemIdentifier::from_dotted("test"))); + let bounds = registry.as_bounds(); + + let expected_minimum = size_of::(); + let aggregation_state_bytes = config.context_limit * (size_of::() + size_of::()); + let context_snapshot_bytes = config.context_limit * size_of::(); + + assert_eq!(bounds.total_minimum_required_bytes(), expected_minimum); + assert_eq!( + bounds.total_firm_limit_bytes(), + expected_minimum + aggregation_state_bytes + context_snapshot_bytes + ); + } + + #[tokio::test] + async fn production_owner_loop_serves_snapshots_and_stops_after_cancellation() { + tokio::time::timeout(Duration::from_secs(5), async { + let mut config = AggregateConfiguration::with_defaults(); + config.primary_flush_interval = Duration::from_secs(60); + let snapshot_handle = config.context_snapshot_handle(); + + let available_request_capacity = snapshot_handle.requests.capacity(); + let canceled_request = tokio::spawn({ + let snapshot_handle = snapshot_handle.clone(); + async move { snapshot_handle.snapshot().await } + }); + while snapshot_handle.requests.capacity() == available_request_capacity { + tokio::task::yield_now().await; + } + canceled_request.abort(); + assert!(canceled_request + .await + .expect_err("snapshot requester should be canceled") + .is_cancelled()); + + let (events_tx, events_rx) = mpsc::channel(1); + let source = ControlledMetricSourceBuilder { + events: Mutex::new(Some(events_rx)), + outputs: vec![OutputDefinition::default_output(EventType::Metric)], + }; + let component_registry = ComponentRegistry::default(); + let mut blueprint = TopologyBlueprint::new("aggregate_snapshot_owner", &component_registry); + blueprint + .add_source("source", source) + .expect("controlled source should be accepted") + .add_transform("aggregate", config) + .expect("aggregate transform should be accepted") + .add_destination("destination", DrainingMetricDestinationBuilder) + .expect("draining destination should be accepted"); + blueprint + .connect_components_in_order(["source", "aggregate", "destination"]) + .expect("test topology should connect"); + blueprint + .with_health_registry(HealthRegistry::new()) + .with_memory_limiter(MemoryLimiter::noop()) + .with_ambient_worker_pool(); + + let mut supervisor = + Supervisor::new("aggregate-snapshot-owner").expect("test supervisor should be created"); + supervisor.add_worker(blueprint); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let topology_task = tokio::spawn(async move { supervisor.run_with_shutdown(shutdown_rx).await }); + + let passthrough_context = Context::from_static_name("owner.loop.timestamped.gauge"); + events_tx + .send(Event::Metric(Metric::gauge( + passthrough_context.clone(), + (insert_ts(1), 1.0), + ))) + .await + .expect("controlled source should accept a timestamped event"); + + let retained_context = Context::from_static_name("owner.loop.mixed.counter"); + let mixed_values = ScalarPoints::from_iter([(None, 2.0), (NonZeroU64::new(insert_ts(1)), 3.0)]); + events_tx + .send(Event::Metric(Metric::counter(retained_context.clone(), mixed_values))) + .await + .expect("controlled source should accept a mixed event"); + + let snapshot = loop { + let snapshot = snapshot_handle + .snapshot() + .await + .expect("running aggregate should fulfill snapshots"); + if snapshot.iter().any(|entry| entry.context() == &retained_context) { + break snapshot; + } + tokio::task::yield_now().await; + }; + assert_eq!( + snapshot.len(), + 1, + "only the mixed metric's retained portion belongs in state" + ); + let entry = &snapshot[0]; + assert_eq!(entry.context(), &retained_context); + assert_eq!(entry.metric_type(), AggregateMetricType::Counter); + assert_eq!(entry.unit(), None); + assert!(snapshot.iter().all(|entry| entry.context() != &passthrough_context)); + + drop(events_tx); + drop(snapshot_handle); + shutdown_tx.send(()).expect("test topology should still be running"); + let topology_result = topology_task.await.expect("topology task should not panic"); + assert!( + topology_result.is_ok(), + "topology should stop cleanly: {topology_result:?}" + ); + }) + .await + .expect("production aggregate owner loop should complete without spinning or hanging"); + } + + #[tokio::test] + async fn snapshot_handle_round_trips_through_responder() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let expected = vec![AggregateContextSnapshotEntry::for_test( + Context::from_static_name("round.trip"), + AggregateMetricType::Gauge, + MetaString::from_static("widget"), + )]; + let snapshot_task = tokio::spawn(async move { handle.snapshot().await }); + + responder + .respond(expected.clone()) + .await + .expect("responder should receive and fulfill a snapshot request"); + + let actual = snapshot_task + .await + .expect("snapshot task should complete") + .expect("snapshot request should succeed"); + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn snapshot_handle_reports_dropped_receiver() { + let (handle, responder) = aggregate_context_snapshot_channel_for_test(); + drop(responder); + + let error = handle + .snapshot() + .await + .expect_err("snapshot should fail after its owner is dropped"); + assert!(error.to_string().contains("unavailable")); + } + + #[tokio::test] + async fn snapshot_responder_stops_after_accepting_request_and_cancels_response() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let snapshot_task = tokio::spawn(async move { handle.snapshot().await }); + + responder + .stop_after_receiving() + .await + .expect("responder should accept the snapshot request before stopping"); + + let error = snapshot_task + .await + .expect("snapshot task should complete") + .expect_err("snapshot should fail when the accepted response is dropped"); + assert!(error.to_string().contains("stopped before responding")); + } + + #[tokio::test] + async fn aggregate_configuration_receiver_can_only_be_taken_once() { + let config = AggregateConfiguration::with_defaults(); + let _handle = config.context_snapshot_handle(); + let first = config.build(ComponentContext::test_transform("aggregate_one")).await; + assert!(first.is_ok()); + + let second = config.build(ComponentContext::test_transform("aggregate_two")).await; + let error = match second { + Ok(_) => panic!("second build should not take the snapshot receiver again"), + Err(error) => error, + }; + assert!(error.to_string().contains("already been taken")); + } + + #[tokio::test] + async fn snapshot_responder_reports_closed_request_channel() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + drop(handle); + + let error = responder + .respond(Vec::new()) + .await + .expect_err("responder should fail when every request handle is dropped"); + assert!(error.to_string().contains("closed")); + } + + #[tokio::test] + async fn snapshot_responder_ignores_canceled_requester() { + let (handle, mut responder) = aggregate_context_snapshot_channel_for_test(); + let snapshot_task = tokio::spawn(async move { handle.snapshot().await }); + let pending_response = responder + .receive() + .await + .expect("responder should accept the snapshot request"); + + snapshot_task.abort(); + let _ = snapshot_task.await; + + pending_response.respond(Vec::new()); + } + #[test] fn bucket_is_closed() { // Cases are defined as: diff --git a/lib/saluki-components/src/transforms/mod.rs b/lib/saluki-components/src/transforms/mod.rs index 10c2c615bf..ebeac7c412 100644 --- a/lib/saluki-components/src/transforms/mod.rs +++ b/lib/saluki-components/src/transforms/mod.rs @@ -4,7 +4,14 @@ mod autoscaling_failover_gateway; pub use self::autoscaling_failover_gateway::AutoscalingFailoverGatewayConfiguration; mod aggregate; -pub use self::aggregate::AggregateConfiguration; +#[cfg(feature = "test-util")] +pub use self::aggregate::{ + aggregate_context_snapshot_channel_for_test, AggregateContextSnapshotPendingResponse, + AggregateContextSnapshotResponder, +}; +pub use self::aggregate::{ + AggregateConfiguration, AggregateContextSnapshotEntry, AggregateContextSnapshotHandle, AggregateMetricType, +}; mod chained; pub use self::chained::ChainedConfiguration; diff --git a/test/integration/cases/dogstatsd-top-windows/config.yaml b/test/integration/cases/dogstatsd-top-windows/config.yaml new file mode 100644 index 0000000000..d69d4618ec --- /dev/null +++ b/test/integration/cases/dogstatsd-top-windows/config.yaml @@ -0,0 +1,129 @@ +type: integration +name: "dogstatsd-top-windows" +description: "Verifies the privileged DogStatsD top workflow and context artifact replacement on Windows" +timeout: 180s +runtimes: [windows] + +env: + DD_API_KEY: "00000000000000000000000000000000" + DD_HOSTNAME: "integration-test-dogstatsd-top-windows" + DD_DATA_PLANE_ENABLED: "true" + DD_DATA_PLANE_STANDALONE_MODE: "false" + DD_DATA_PLANE_DOGSTATSD_ENABLED: "true" + +container: + exposed_ports: + - "58125/udp" + - "55101/tcp" + +procedure: + - parallel: + - assertion: port_listening + port: 58125 + protocol: udp + timeout: 60s + - assertion: port_listening + port: 55101 + protocol: tcp + timeout: 60s + # The route is POST-only. A GET returns 405 when registration is complete and 404 when absent. + - assertion: http_check + endpoint: "https://localhost:55101/agent/dogstatsd-contexts-dump" + status: + not_equal: 404 + insecure_skip_verify: true + timeout: 60s + - action: target_exec + command: + - "pwsh" + - "-NoProfile" + - "-NonInteractive" + - "-Command" + - | + $ErrorActionPreference = "Stop" + Set-StrictMode -Version 3.0 + + $Adp = "C:\adp\agent-data-plane.exe" + $Config = "C:\ProgramData\Datadog\datadog.yaml" + $ApiUrl = "https://127.0.0.1:55101/agent/dogstatsd-contexts-dump" + $Result = "C:\dogstatsd-top-result" + $CopiedDump = "C:\dogstatsd-top-copied.json.zstd" + $RequestRow = " 2`tintegration.windows.top.requests`t(2 instance, 1 env)" + $QueueRow = " 1`tintegration.windows.top.queue`t(1 queue)" + + function Invoke-AdpDogStatsD([string[]] $Arguments) { + $Output = @(& $Adp --config $Config dogstatsd @Arguments 2>&1) + if ($LASTEXITCODE -ne 0) { + throw "ADP command failed ($LASTEXITCODE): $($Arguments -join ' ')`n$($Output -join "`n")" + } + return $Output -join "`n" + } + + function Send-DogStatsDContexts { + $Payload = @( + "integration.windows.top.requests:1|c|#host:windows-a,env:prod,instance:a", + "integration.windows.top.requests:1|c|#host:windows-b,env:prod,instance:b", + "integration.windows.top.queue:1|c|#queue:alpha" + ) -join "`n" + $Bytes = [System.Text.Encoding]::UTF8.GetBytes($Payload) + $Client = [System.Net.Sockets.UdpClient]::new() + try { + 1..3 | ForEach-Object { [void] $Client.Send($Bytes, $Bytes.Length, "127.0.0.1", 58125) } + } finally { + $Client.Dispose() + } + } + + function Wait-ForOnlineReport { + $LastOutput = "" + foreach ($Attempt in 1..40) { + $LastOutput = Invoke-AdpDogStatsD @("top") + if ($LastOutput.Contains($RequestRow) -and $LastOutput.Contains($QueueRow)) { + return $LastOutput + } + Start-Sleep -Milliseconds 500 + } + throw "Retained contexts did not appear in online output:`n$LastOutput" + } + + function Get-DumpPath([string] $Output) { + $Lines = @($Output -split "`r?`n" | Where-Object { $_ -ne "" }) + if ($Lines.Count -ne 1 -or -not $Lines[0].StartsWith("Wrote ")) { + throw "Unexpected dump-contexts output: $Output" + } + $Path = $Lines[0].Substring("Wrote ".Length) + if (-not [System.IO.Path]::IsPathRooted($Path) -or [System.IO.Path]::GetFileName($Path) -ne "dogstatsd_contexts.json.zstd") { + throw "Unexpected context dump path: $Path" + } + if (-not (Test-Path -LiteralPath $Path -PathType Leaf)) { + throw "Context dump was not created: $Path" + } + return $Path + } + + Send-DogStatsDContexts + $Online = Wait-ForOnlineReport + if (-not $Online.StartsWith("Wrote ")) { + throw "Online top did not print the generated artifact path:`n$Online" + } + + $DumpPath = Get-DumpPath (Invoke-AdpDogStatsD @("dump-contexts")) + Copy-Item -LiteralPath $DumpPath -Destination $CopiedDump -Force + $Offline = Invoke-AdpDogStatsD @("top", "--path", $CopiedDump) + if (-not $Offline.Contains($RequestRow) -or -not $Offline.Contains($QueueRow)) { + throw "Offline report is missing retained contexts:`n$Offline" + } + + Set-Content -LiteralPath $Result -Value "passed" -Encoding ASCII + timeout: 90s + - assertion: file_contains + path: "C:\\dogstatsd-top-result" + pattern: "passed" + timeout: 10s + - parallel: + - assertion: process_stable_for + duration: 10s + - assertion: log_not_contains + pattern: "panic|PANIC" + regex: true + during: 10s diff --git a/test/integration/cases/dogstatsd-top/agent-7.80.3-contexts.json.zstd b/test/integration/cases/dogstatsd-top/agent-7.80.3-contexts.json.zstd new file mode 100644 index 0000000000..40d33c4e5e Binary files /dev/null and b/test/integration/cases/dogstatsd-top/agent-7.80.3-contexts.json.zstd differ diff --git a/test/integration/cases/dogstatsd-top/agent-7.80.3-contexts.ndjson b/test/integration/cases/dogstatsd-top/agent-7.80.3-contexts.ndjson new file mode 100644 index 0000000000..ff513907eb --- /dev/null +++ b/test/integration/cases/dogstatsd-top/agent-7.80.3-contexts.ndjson @@ -0,0 +1,4 @@ +{"Name":"interop.agent.requests","Host":"agent-b","Type":"Counter","TaggerTags":[],"MetricTags":["env:prod","instance:b"],"NoIndex":false,"Source":1} +{"Name":"interop.agent.queue","Host":"dogstatsd-interop-agent","Type":"Counter","TaggerTags":[],"MetricTags":["queue:alpha"],"NoIndex":false,"Source":1} +{"Name":"datadog.agent.started","Host":"dogstatsd-interop-agent","Type":"Count","TaggerTags":[],"MetricTags":["version:7.80.3","package_version:bin"],"NoIndex":false,"Source":0} +{"Name":"interop.agent.requests","Host":"agent-a","Type":"Counter","TaggerTags":[],"MetricTags":["env:prod","instance:a"],"NoIndex":false,"Source":1} diff --git a/test/integration/cases/dogstatsd-top/config.yaml b/test/integration/cases/dogstatsd-top/config.yaml new file mode 100644 index 0000000000..22b0ef6b0f --- /dev/null +++ b/test/integration/cases/dogstatsd-top/config.yaml @@ -0,0 +1,57 @@ +type: integration +name: "dogstatsd-top" +description: "Verifies retained DogStatsD contexts through the privileged API, online CLI, and offline Agent-compatible artifact" +timeout: 180s +runtimes: [linux] + +env: + DD_API_KEY: "00000000000000000000000000000000" + DD_HOSTNAME: "integration-test-dogstatsd-top" + DD_DATA_PLANE_ENABLED: "true" + DD_DATA_PLANE_STANDALONE_MODE: "false" + DD_DATA_PLANE_DOGSTATSD_ENABLED: "true" + +container: + files: + - "run_dogstatsd_top_test.py:/dogstatsd-top/run_dogstatsd_top_test.py" + # Produced by Datadog Agent 7.80.3 `dogstatsd dump-contexts`; the NDJSON file is the exact + # decompressed content of the zstd artifact. + - "agent-7.80.3-contexts.json.zstd:/dogstatsd-top/agent-7.80.3-contexts.json.zstd" + - "agent-7.80.3-contexts.ndjson:/dogstatsd-top/agent-7.80.3-contexts.ndjson" + exposed_ports: + - "58125/udp" + - "55101/tcp" + +procedure: + - parallel: + - assertion: port_listening + port: 58125 + protocol: udp + timeout: 60s + - assertion: port_listening + port: 55101 + protocol: tcp + timeout: 60s + # The route is POST-only. A GET returns 405 when registration is complete and 404 when absent. + - assertion: http_check + endpoint: "https://localhost:55101/agent/dogstatsd-contexts-dump" + status: + not_equal: 404 + insecure_skip_verify: true + timeout: 60s + - action: target_exec + command: + - "/opt/datadog-agent/embedded/bin/python3" + - "/dogstatsd-top/run_dogstatsd_top_test.py" + timeout: 90s + - assertion: file_contains + path: "/tmp/dogstatsd-top-integration-result" + pattern: "passed" + timeout: 10s + - parallel: + - assertion: process_stable_for + duration: 10s + - assertion: log_not_contains + pattern: "panic|PANIC" + regex: true + during: 10s diff --git a/test/integration/cases/dogstatsd-top/run_dogstatsd_top_test.py b/test/integration/cases/dogstatsd-top/run_dogstatsd_top_test.py new file mode 100644 index 0000000000..2e467c4e0b --- /dev/null +++ b/test/integration/cases/dogstatsd-top/run_dogstatsd_top_test.py @@ -0,0 +1,155 @@ +#!/opt/datadog-agent/embedded/bin/python3 + +import json +import shutil +import socket +import ssl +import stat +import subprocess +import time +import urllib.request +from pathlib import Path + +ADP = Path("/opt/datadog-agent/embedded/bin/agent-data-plane") +AGENT = Path("/opt/datadog-agent/bin/agent/agent") +CONFIG = Path("/etc/datadog-agent/datadog.yaml") +API_URL = "https://127.0.0.1:55101/agent/dogstatsd-contexts-dump" +DUMP_FILENAME = "dogstatsd_contexts.json.zstd" +COPIED_DUMP = Path("/tmp/dogstatsd-top-copied.json.zstd") +AGENT_FIXTURE_COMPRESSED = Path("/dogstatsd-top/agent-7.80.3-contexts.json.zstd") +AGENT_FIXTURE_PLAIN = Path("/dogstatsd-top/agent-7.80.3-contexts.ndjson") +ONLINE_OUTPUT = Path("/tmp/dogstatsd-top-online-output") +OFFLINE_OUTPUT = Path("/tmp/dogstatsd-top-offline-output") +AGENT_OUTPUT = Path("/tmp/dogstatsd-top-agent-output") +RESULT = Path("/tmp/dogstatsd-top-integration-result") + +REQUEST_ROW = " 2\tintegration.top.requests\t(2 instance, 1 env, 1 service)" +QUEUE_ROW = " 1\tintegration.top.queue\t(1 queue)" + + +def run(command): + completed = subprocess.run(command, text=True, capture_output=True, timeout=30) + if completed.returncode != 0: + raise RuntimeError( + f"command failed ({completed.returncode}): {' '.join(map(str, command))}\n" + f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + ) + return completed.stdout + + +def run_adp(*arguments): + return run([str(ADP), "--config", str(CONFIG), "dogstatsd", *arguments]) + + +def unverified_tls_context(): + context = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + return context + + +def post_dump(): + request = urllib.request.Request(API_URL, data=b"", method="POST") + return urllib.request.urlopen(request, context=unverified_tls_context(), timeout=10) + + +def send_dogstatsd_contexts(): + payload = b"\n".join( + [ + b"integration.top.requests:1|c|#host:host-a,env:prod,service:web,instance:a", + b"integration.top.requests:1|c|#host:host-b,env:prod,service:web,instance:b", + b"integration.top.queue:1|c|#queue:alpha", + ] + ) + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as client: + for _ in range(3): + client.sendto(payload, ("127.0.0.1", 58125)) + + +def wait_for_online_report(): + last_output = "" + for _ in range(40): + last_output = run_adp("top") + if REQUEST_ROW in last_output and QUEUE_ROW in last_output: + ONLINE_OUTPUT.write_text(last_output, encoding="utf-8") + return last_output + time.sleep(0.5) + raise AssertionError(f"retained contexts did not appear in online output:\n{last_output}") + + +def parse_written_path(output): + lines = output.splitlines() + if len(lines) != 1 or not lines[0].startswith("Wrote "): + raise AssertionError(f"unexpected dump-contexts output: {output!r}") + path = Path(lines[0].removeprefix("Wrote ")) + if not path.is_absolute() or path.name != DUMP_FILENAME: + raise AssertionError(f"unexpected context dump path: {path}") + if not path.is_file(): + raise AssertionError(f"context dump was not created: {path}") + return path + + +def normalize_report(output): + return output.replace("\r\n", "\n").rstrip("\n") + + +def validate_online_api(): + with post_dump() as response: + body = response.read().decode("utf-8") + if response.status != 200: + raise AssertionError(f"context dump API returned {response.status}: {body}") + if not response.headers.get_content_type() == "application/json": + raise AssertionError(f"unexpected API content type: {response.headers.get('Content-Type')}") + path = Path(json.loads(body)) + if path.name != DUMP_FILENAME or not path.is_file(): + raise AssertionError(f"context dump API returned an invalid artifact path: {path}") + + +def validate_agent_fixture_interoperability(): + reports = [] + for fixture in [AGENT_FIXTURE_COMPRESSED, AGENT_FIXTURE_PLAIN]: + reports.append(run_adp("top", "--path", str(fixture))) + reports.append(run([str(AGENT), "dogstatsd", "top", "--path", str(fixture)])) + + normalized = [normalize_report(report) for report in reports] + if any(report != normalized[0] for report in normalized[1:]): + raise AssertionError(f"Agent-produced fixture reports differ: {reports}") + if "interop.agent.requests" not in normalized[0]: + raise AssertionError(f"Agent-produced fixture did not contain expected contexts: {normalized[0]}") + + +def validate_dump_and_offline_interoperability(): + dump_path = parse_written_path(run_adp("dump-contexts")) + if stat.S_IMODE(dump_path.stat().st_mode) != 0o600: + raise AssertionError(f"context dump mode is not 0600: {oct(stat.S_IMODE(dump_path.stat().st_mode))}") + + shutil.copyfile(dump_path, COPIED_DUMP) + offline = run_adp("top", "--path", str(COPIED_DUMP)) + OFFLINE_OUTPUT.write_text(offline, encoding="utf-8") + if REQUEST_ROW not in offline or QUEUE_ROW not in offline: + raise AssertionError(f"offline report is missing retained contexts:\n{offline}") + + agent = run([str(AGENT), "dogstatsd", "top", "--path", str(COPIED_DUMP)]) + AGENT_OUTPUT.write_text(agent, encoding="utf-8") + if normalize_report(agent) != normalize_report(offline): + raise AssertionError(f"Agent and ADP reports differ:\nADP:\n{offline}\nAgent:\n{agent}") + + +def main(): + for required_path in [ADP, AGENT, CONFIG, AGENT_FIXTURE_COMPRESSED, AGENT_FIXTURE_PLAIN]: + if not required_path.exists(): + raise AssertionError(f"required test path does not exist: {required_path}") + + validate_agent_fixture_interoperability() + send_dogstatsd_contexts() + online = wait_for_online_report() + if not online.startswith("Wrote "): + raise AssertionError(f"online top did not print the generated artifact path:\n{online}") + + validate_online_api() + validate_dump_and_offline_interoperability() + RESULT.write_text("passed\n", encoding="utf-8") + + +if __name__ == "__main__": + main()