diff --git a/Cargo.toml b/Cargo.toml index 3deb3b14c8..bd6e0b9f0a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -695,6 +695,7 @@ sources-logs-mezmo = [ "sources-aws_s3", "sources-datadog_agent", "sources-demo_logs", + "sources-mezmo_analytics", "sources-mezmo_demo_logs", "sources-exec", "sources-fluent", @@ -754,6 +755,7 @@ sources-host_metrics = ["heim/cpu", "heim/host", "heim/memory", "heim/net"] sources-http_client = ["sources-utils-http-client"] sources-http_server = ["sources-utils-http", "sources-utils-http-headers", "sources-utils-http-query"] sources-internal_logs = [] +sources-mezmo_analytics = ["sources-internal_logs"] sources-mezmo_pipeline_state_variable_change = [] sources-mezmo_user_logs = ["sources-internal_logs"] sources-internal_metrics = [] diff --git a/lib/vector-core/Cargo.toml b/lib/vector-core/Cargo.toml index 3738a99ba6..851358daad 100644 --- a/lib/vector-core/Cargo.toml +++ b/lib/vector-core/Cargo.toml @@ -53,7 +53,7 @@ snafu.workspace = true socket2.workspace = true tokio = { workspace = true, features = ["net"] } tokio-openssl = { version = "0.6.5", default-features = false } -tokio-stream = { workspace = true, features = ["time"], optional = true } +tokio-stream = { workspace = true, features = ["sync", "time"] } tokio-util = { version = "0.7.0", default-features = false, features = ["time"] } tokio-postgres = { version = "0.7.7", default-features = false, features = ["runtime", "with-chrono-0_4", "with-uuid-1", "with-serde_json-1"] } toml.workspace = true @@ -104,7 +104,7 @@ vector-common = { path = "../vector-common", default-features = false, features [features] default = [] -lua = ["dep:mlua", "dep:tokio-stream", "vrl/lua"] +lua = ["dep:mlua", "vrl/lua"] vrl = [] test = ["vector-common/test", "proptest"] pgbouncer-integration-tests = ["vrl"] diff --git a/lib/vector-core/src/mezmo/analytics.rs b/lib/vector-core/src/mezmo/analytics.rs new file mode 100644 index 0000000000..da3741c1a1 --- /dev/null +++ b/lib/vector-core/src/mezmo/analytics.rs @@ -0,0 +1,146 @@ +//! In-process delivery for Mezmo analytics records. + +use std::sync::OnceLock; + +use futures::{Stream, StreamExt, future::ready}; +use tokio::sync::broadcast::{self, Receiver, Sender}; +use tokio_stream::wrappers::{BroadcastStream, errors::BroadcastStreamRecvError}; +use tracing::warn; +use vector_common::internal_event::{self, ComponentEventsDropped, UNINTENTIONAL}; + +use crate::event::LogEvent; + +const ANALYTICS_CHANNEL_CAPACITY: usize = 1_000; +const ANALYTICS_OUTPUT_COUNT: usize = 5; + +type AnalyticsSenders = [Sender; ANALYTICS_OUTPUT_COUNT]; + +static ANALYTICS: OnceLock = OnceLock::new(); + +/// A named output exposed by the `mezmo_analytics` source. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AnalyticsOutput { + UsageMetrics, + UsageMetricsByAnnotations, + LogClusters, + LogClusterSamples, + LogClusterUsage, +} + +impl AnalyticsOutput { + /// Returns the source output name. + pub const fn as_str(self) -> &'static str { + match self { + Self::UsageMetrics => "usage_metrics", + Self::UsageMetricsByAnnotations => "usage_metrics_by_annotations", + Self::LogClusters => "log_clusters", + Self::LogClusterSamples => "log_cluster_samples", + Self::LogClusterUsage => "log_cluster_usage", + } + } + + /// Returns every named output in declaration order. + pub const fn all() -> [Self; ANALYTICS_OUTPUT_COUNT] { + [ + Self::UsageMetrics, + Self::UsageMetricsByAnnotations, + Self::LogClusters, + Self::LogClusterSamples, + Self::LogClusterUsage, + ] + } +} + +/// A batch of analytics records destined for one named source output. +#[derive(Clone, Debug)] +pub struct AnalyticsEventBatch { + output: AnalyticsOutput, + events: Vec, +} + +impl AnalyticsEventBatch { + /// Creates a batch for one named output. + pub fn new(output: AnalyticsOutput, events: Vec) -> Self { + Self { output, events } + } + + /// Returns the named output associated with this batch. + pub const fn output(&self) -> AnalyticsOutput { + self.output + } + + /// Returns the records in this batch. + pub fn events(&self) -> &[LogEvent] { + &self.events + } + + /// Splits the batch into its named output and records. + pub fn into_parts(self) -> (AnalyticsOutput, Vec) { + (self.output, self.events) + } +} + +fn senders() -> &'static AnalyticsSenders { + ANALYTICS + .get_or_init(|| std::array::from_fn(|_| broadcast::channel(ANALYTICS_CHANNEL_CAPACITY).0)) +} + +/// Builds and publishes batches if a `mezmo_analytics` source is subscribed. +pub fn publish(build_batches: F) +where + F: FnOnce() -> B, + B: IntoIterator, +{ + let Some(senders) = ANALYTICS.get() else { + return; + }; + if senders.iter().all(|sender| sender.receiver_count() == 0) { + return; + } + + for batch in build_batches() { + let sender = &senders[batch.output() as usize]; + if sender.receiver_count() > 0 { + let _ = sender.send(batch); + } + } +} + +/// A subscription to analytics batches produced inside Vector. +pub struct AnalyticsSubscription { + receiver: Receiver, +} + +impl AnalyticsSubscription { + /// Subscribes to analytics batches for one output produced after this call. + pub fn subscribe(output: AnalyticsOutput) -> Self { + Self { + receiver: senders()[output as usize].subscribe(), + } + } + + /// Subscribes to analytics batches for every output. + pub fn subscribe_all() -> Vec { + AnalyticsOutput::all() + .into_iter() + .map(Self::subscribe) + .collect() + } + + /// Converts this subscription into its analytics batch stream. + pub fn into_stream(self) -> impl Stream + Unpin { + BroadcastStream::new(self.receiver).filter_map(|received| { + ready(match received { + Ok(batch) => Some(batch), + Err(error @ BroadcastStreamRecvError::Lagged(dropped_batches)) => { + warn!(message = "Mezmo analytics source lagged behind.", %error); + internal_event::emit(ComponentEventsDropped:: { + count: usize::try_from(dropped_batches).unwrap_or(usize::MAX), + reason: "Mezmo analytics source lagged behind.", + }); + None + } + }) + }) + } +} diff --git a/lib/vector-core/src/mezmo/mod.rs b/lib/vector-core/src/mezmo/mod.rs index 26e9103cd5..5cf046da62 100644 --- a/lib/vector-core/src/mezmo/mod.rs +++ b/lib/vector-core/src/mezmo/mod.rs @@ -1 +1,2 @@ +pub mod analytics; pub mod postgres; diff --git a/lib/vector-core/src/usage_metrics/flusher.rs b/lib/vector-core/src/usage_metrics/flusher.rs index ad0d97e218..e575540e69 100644 --- a/lib/vector-core/src/usage_metrics/flusher.rs +++ b/lib/vector-core/src/usage_metrics/flusher.rs @@ -1,4 +1,4 @@ -use super::{AnnotationMap, AnnotationSet, UsageMetricsKey, UsageMetricsValue}; +use super::{AnnotationMap, AnnotationSet, UsageMetricsKey, UsageMetricsValue, processor_name}; use crate::mezmo; use async_trait::async_trait; use chrono::Utc; @@ -24,8 +24,6 @@ const INSERT_BILLING_QUERY: &str = "INSERT INTO usage_metrics (event_ts, account const INSERT_PROFILES_QUERY: &str = "INSERT INTO usage_metrics_by_annotations (ts, account_id, component_id, count, size, annotations) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT DO NOTHING"; const DB_MAX_PARALLEL_EXECUTIONS: usize = 8; -const VECTOR_VERSION: &str = env!("CARGO_PKG_VERSION"); - #[async_trait] pub(crate) trait MetricsFlusher: Sync { async fn save_billing_metrics(&self, metrics: HashMap); @@ -67,7 +65,7 @@ impl DbFlusher { DbFlusherError::QueryError })?; - let processor_name = format!("app=vector,pod={pod_name},version={VECTOR_VERSION}"); + let processor_name = processor_name(pod_name); Ok(Self { conn_str, processor_name, @@ -249,7 +247,7 @@ impl HttpFlusher { headers: HashMap, max_delay: Duration, ) -> Self { - let processor_name = format!("app=vector,pod={pod_name},version={VECTOR_VERSION}"); + let processor_name = processor_name(pod_name); HttpFlusher { client: Client::new(), diff --git a/lib/vector-core/src/usage_metrics/mod.rs b/lib/vector-core/src/usage_metrics/mod.rs index 95a4f9ff97..091becdb93 100644 --- a/lib/vector-core/src/usage_metrics/mod.rs +++ b/lib/vector-core/src/usage_metrics/mod.rs @@ -1,3 +1,4 @@ +use chrono::Utc; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::collections::hash_map::Entry; @@ -18,8 +19,8 @@ use vrl::value::{KeyString, Value}; use crate::{ config::log_schema, - event::EventArray, - event::{MetricValue, array::EventContainer}, + event::{EventArray, LogEvent, MetricValue, array::EventContainer}, + mezmo::analytics::{self, AnalyticsEventBatch, AnalyticsOutput}, usage_metrics::flusher::HttpFlusher, }; use flusher::{DbFlusher, MetricsFlusher, StdErrFlusher}; @@ -30,6 +31,7 @@ const DEFAULT_FLUSH_INTERVAL_SECS: u64 = 20; const DEFAULT_PROFILE_FLUSH_INTERVAL: Duration = Duration::from_secs(60); const BASE_ARRAY_SIZE: usize = 8; // Add some overhead to the array and object size const BASE_BTREE_SIZE: usize = 8; +const VECTOR_VERSION: &str = env!("CARGO_PKG_VERSION"); static INTERNAL_TRANSFORM: ComponentKind = ComponentKind::Transform { internal: true }; mod flusher; @@ -65,6 +67,23 @@ pub enum ComponentKind { Transform { internal: bool }, } +impl ComponentKind { + fn as_str(&self) -> &'static str { + match self { + Self::Source { .. } => "source", + Self::Sink => "sink", + Self::Transform { .. } => "transform", + } + } + + fn is_internal(&self) -> bool { + match self { + Self::Source { internal } | Self::Transform { internal } => *internal, + Self::Sink => false, + } + } +} + impl FromStr for ComponentKind { type Err = ParseError; @@ -261,6 +280,129 @@ impl AnnotationSet { type AnnotationMap = HashMap; +fn usage_event( + key: &UsageMetricsKey, + timestamp: chrono::DateTime, + processor: &Value, +) -> LogEvent { + let pipeline_id = key + .pipeline_id + .as_ref() + .map_or(Value::Null, |value| Value::from(value.clone())); + + LogEvent::from(Value::Object(BTreeMap::from([ + ("timestamp".into(), Value::Timestamp(timestamp)), + ("account_id".into(), Value::from(key.account_id.clone())), + ("pipeline_id".into(), pipeline_id), + ("component_id".into(), Value::from(key.component_id.clone())), + ( + "component_type".into(), + Value::from(key.component_type.clone()), + ), + ( + "component_kind".into(), + Value::from(key.component_kind.as_str()), + ), + ( + "internal".into(), + Value::from(key.component_kind.is_internal()), + ), + ("processor".into(), processor.clone()), + ]))) +} + +fn processor_name(pod_name: &str) -> String { + format!("app=vector,pod={pod_name},version={VECTOR_VERSION}") +} + +fn usage_metrics_events(metrics: &HashMap) -> Vec { + let timestamp = Utc::now(); + let pod_name = env::var("POD_NAME").unwrap_or_else(|_| "not-set".to_owned()); + let processor = Value::from(processor_name(&pod_name)); + metrics + .iter() + .map(|(key, value)| { + let mut event = usage_event(key, timestamp, &processor); + event.insert("total_count", value.total_count as i64); + event.insert("total_size", value.total_size as i64); + event + }) + .collect() +} + +fn annotation_event( + mut event: LogEvent, + annotation: &AnnotationSet, + value: &UsageMetricsValue, +) -> LogEvent { + event.insert("count", value.total_count as i64); + event.insert("size", value.total_size as i64); + event.insert( + "annotations", + Value::from( + serde_json::to_value(annotation).expect("annotation sets should always serialize"), + ), + ); + event +} + +fn usage_metrics_by_annotations_events( + metrics: &HashMap, +) -> Vec { + let timestamp = Utc::now(); + let pod_name = env::var("POD_NAME").unwrap_or_else(|_| "not-set".to_owned()); + let processor = Value::from(processor_name(&pod_name)); + let event_count = metrics.values().map(HashMap::len).sum(); + let mut events = Vec::with_capacity(event_count); + + for (key, annotations) in metrics { + let mut annotations = annotations.iter(); + let Some((first_annotation, first_annotation_value)) = annotations.next() else { + continue; + }; + let base_event = usage_event(key, timestamp, &processor); + + for (annotation, value) in annotations { + events.push(annotation_event(base_event.clone(), annotation, value)); + } + events.push(annotation_event( + base_event, + first_annotation, + first_annotation_value, + )); + } + + events +} + +fn publish_usage_metrics(metrics: &HashMap) { + analytics::publish(|| { + let events = usage_metrics_events(metrics); + if events.is_empty() { + None + } else { + Some(AnalyticsEventBatch::new( + AnalyticsOutput::UsageMetrics, + events, + )) + } + }); +} + +fn publish_usage_metrics_by_annotations(metrics: &HashMap) { + analytics::publish(|| { + let events = usage_metrics_by_annotations_events(metrics); + if events.is_empty() { + None + } else { + Some(AnalyticsEventBatch::new( + AnalyticsOutput::UsageMetricsByAnnotations, + events, + )) + } + }); +} + /// Represents aggregated size and count information for events #[derive(Debug, Default)] pub struct UsageProfileValue { @@ -787,11 +929,15 @@ fn start_publishing_metrics_with_flusher( billing_events_count ); + publish_usage_metrics(&aggregated_billing); + // Flush billing metrics in the foreground flusher.save_billing_metrics(aggregated_billing).await; } if start_profile.elapsed() > profile_agg_window && !aggregated_profiles.is_empty() { + publish_usage_metrics_by_annotations(&aggregated_profiles); + // Flush aggregated profiles let flusher = Arc::clone(&flusher); @@ -1329,4 +1475,59 @@ mod tests { ]) ); } + + #[test] + fn creates_storage_neutral_analytics_events() { + let key = UsageMetricsKey { + account_id: "account".into(), + pipeline_id: Some("pipeline".into()), + component_id: "component".into(), + component_type: "http".into(), + component_kind: ComponentKind::Source { internal: false }, + }; + let metrics = HashMap::from([( + key.clone(), + UsageMetricsValue { + total_count: 2, + total_size: 20, + }, + )]); + + let event = usage_metrics_events(&metrics) + .pop() + .expect("one usage metrics event"); + assert_eq!(event.get("account_id"), Some(&Value::from("account"))); + assert_eq!(event.get("pipeline_id"), Some(&Value::from("pipeline"))); + assert_eq!(event.get("component_id"), Some(&Value::from("component"))); + assert_eq!(event.get("component_type"), Some(&Value::from("http"))); + assert_eq!(event.get("component_kind"), Some(&Value::from("source"))); + assert_eq!(event.get("internal"), Some(&Value::from(false))); + assert_eq!(event.get("total_count"), Some(&Value::from(2))); + assert_eq!(event.get("total_size"), Some(&Value::from(20))); + assert!(event.get("processor").is_some()); + assert!(event.get("timestamp").is_some()); + + let profiles = HashMap::from([( + key, + HashMap::from([( + AnnotationSet { + app: Some("api".into()), + host: None, + level: Some("error".into()), + log_type: None, + }, + UsageMetricsValue { + total_count: 1, + total_size: 10, + }, + )]), + )]); + let event = usage_metrics_by_annotations_events(&profiles) + .pop() + .expect("one annotated usage event"); + assert_eq!(event.get("annotations.app"), Some(&Value::from("api"))); + assert_eq!(event.get("annotations.level"), Some(&Value::from("error"))); + assert_eq!(event.get("count"), Some(&Value::from(1))); + assert_eq!(event.get("size"), Some(&Value::from(10))); + } } diff --git a/src/sources/mezmo_analytics/config.rs b/src/sources/mezmo_analytics/config.rs new file mode 100644 index 0000000000..7d63f9b549 --- /dev/null +++ b/src/sources/mezmo_analytics/config.rs @@ -0,0 +1,45 @@ +use vector_lib::{ + config::LogNamespace, + configurable::configurable_component, + mezmo::analytics::{AnalyticsOutput, AnalyticsSubscription}, + schema::Definition, +}; + +use crate::config::{DataType, SourceConfig, SourceContext, SourceOutput}; + +/// Configuration for the `mezmo_analytics` source. +#[configurable_component(source( + "mezmo_analytics", + "Expose Mezmo generated analytic records such as usage metrics." +))] +#[derive(Clone, Debug, Default)] +#[serde(deny_unknown_fields)] +pub struct MezmoAnalyticsConfig {} + +impl_generate_config_from_default!(MezmoAnalyticsConfig); + +#[async_trait::async_trait] +#[typetag::serde(name = "mezmo_analytics")] +impl SourceConfig for MezmoAnalyticsConfig { + async fn build(&self, cx: SourceContext) -> crate::Result { + Ok(Box::pin(super::run( + AnalyticsSubscription::subscribe_all(), + cx.out, + cx.shutdown, + ))) + } + + fn outputs(&self, _global_log_namespace: LogNamespace) -> Vec { + AnalyticsOutput::all() + .into_iter() + .map(|output| { + SourceOutput::new_maybe_logs(DataType::Log, Definition::any()) + .with_port(output.as_str()) + }) + .collect() + } + + fn can_acknowledge(&self) -> bool { + false + } +} diff --git a/src/sources/mezmo_analytics/mod.rs b/src/sources/mezmo_analytics/mod.rs new file mode 100644 index 0000000000..4e54bb7977 --- /dev/null +++ b/src/sources/mezmo_analytics/mod.rs @@ -0,0 +1,204 @@ +mod config; + +pub use config::MezmoAnalyticsConfig; + +use futures::{StreamExt, future::try_join_all}; +use vector_lib::{ + EstimatedJsonEncodedSizeOf, event::Event, mezmo::analytics::AnalyticsSubscription, +}; + +use crate::{ + SourceSender, + internal_events::{InternalLogsBytesReceived, InternalLogsEventsReceived, StreamClosedError}, + shutdown::ShutdownSignal, +}; + +async fn forward_output( + subscription: AnalyticsSubscription, + mut out: SourceSender, + shutdown: ShutdownSignal, +) -> Result<(), ()> { + let mut batches = subscription.into_stream().take_until(shutdown); + + while let Some(batch) = batches.next().await { + let (output, events) = batch.into_parts(); + let count = events.len(); + let byte_size = events.estimated_json_encoded_size_of().get(); + + emit!(InternalLogsBytesReceived { byte_size }); + emit!(InternalLogsEventsReceived { + count, + byte_size: byte_size.into(), + }); + + if out + .send_batch_named(output.as_str(), events.into_iter().map(Event::Log)) + .await + .is_err() + { + emit!(StreamClosedError { count }); + return Err(()); + } + } + + Ok(()) +} + +async fn run( + subscriptions: Vec, + out: SourceSender, + shutdown: ShutdownSignal, +) -> Result<(), ()> { + let forwarders = subscriptions + .into_iter() + .map(|subscription| forward_output(subscription, out.clone(), shutdown.clone())); + + try_join_all(forwarders).await.map(|_| ()) +} + +#[cfg(test)] +mod tests { + use futures::StreamExt; + use serial_test::serial; + use tokio::time::{Duration, timeout}; + use vector_lib::{ + config::LogNamespace, + event::{EventStatus, LogEvent, array::EventContainer}, + mezmo::analytics::{AnalyticsEventBatch, AnalyticsOutput, AnalyticsSubscription, publish}, + }; + + use super::*; + use crate::{ + config::SourceConfig, + test_util::components::{SOURCE_TAGS, SOURCE_TESTS, init_test}, + }; + + #[test] + fn generates_config() { + crate::test_util::test_generate_config::(); + } + + #[test] + fn exposes_named_outputs() { + let outputs = MezmoAnalyticsConfig {}.outputs(LogNamespace::Legacy); + let names = outputs + .into_iter() + .map(|output| output.port.expect("all outputs should be named")) + .collect::>(); + + assert_eq!( + names, + vec![ + "usage_metrics", + "usage_metrics_by_annotations", + "log_clusters", + "log_cluster_samples", + "log_cluster_usage", + ] + ); + } + + #[tokio::test] + #[serial] + async fn forwards_batches_to_their_named_output() { + init_test(); + let (mut out, _default_rx) = SourceSender::new_test(); + let mut usage_rx = out + .add_outputs(EventStatus::Delivered, "usage_metrics".to_owned()) + .flat_map(|item| futures::stream::iter(item.events.into_events())); + let subscriptions = AnalyticsSubscription::subscribe_all(); + let (trigger_shutdown, shutdown, shutdown_done) = ShutdownSignal::new_wired(); + let source = tokio::spawn(run(subscriptions, out, shutdown)); + + publish(|| { + [AnalyticsEventBatch::new( + AnalyticsOutput::UsageMetrics, + vec![LogEvent::from("usage")], + )] + }); + + let event = timeout(Duration::from_secs(1), usage_rx.next()) + .await + .expect("source should forward the batch") + .expect("named output should remain open"); + assert_eq!( + event + .into_log() + .get_message() + .expect("forwarded log should contain a message") + .to_string_lossy(), + "usage" + ); + + drop(trigger_shutdown); + shutdown_done.await; + assert_eq!(source.await.expect("source task should complete"), Ok(())); + SOURCE_TESTS.assert(&SOURCE_TAGS); + } + + #[tokio::test] + #[serial] + async fn blocked_output_does_not_block_other_outputs() { + const USAGE_BATCHES: usize = 200; + + let (mut out, _default_rx) = SourceSender::new_test(); + let mut usage_rx = out + .add_outputs( + EventStatus::Delivered, + AnalyticsOutput::UsageMetrics.as_str().to_owned(), + ) + .flat_map(|item| futures::stream::iter(item.events.into_events())); + let mut cluster_rx = out + .add_outputs( + EventStatus::Delivered, + AnalyticsOutput::LogClusters.as_str().to_owned(), + ) + .flat_map(|item| futures::stream::iter(item.events.into_events())); + let subscriptions = AnalyticsSubscription::subscribe_all(); + let (trigger_shutdown, shutdown, shutdown_done) = ShutdownSignal::new_wired(); + let source = tokio::spawn(run(subscriptions, out, shutdown)); + + publish(|| { + let mut batches = Vec::with_capacity(USAGE_BATCHES + 1); + for _ in 0..USAGE_BATCHES { + batches.push(AnalyticsEventBatch::new( + AnalyticsOutput::UsageMetrics, + vec![LogEvent::from("usage")], + )); + } + batches.push(AnalyticsEventBatch::new( + AnalyticsOutput::LogClusters, + vec![LogEvent::from("cluster")], + )); + batches + }); + + let cluster = timeout(Duration::from_secs(1), cluster_rx.next()) + .await + .expect("cluster output should not be blocked by usage metrics") + .expect("cluster output should remain open"); + assert_eq!( + cluster + .into_log() + .get_message() + .expect("cluster log should contain a message") + .to_string_lossy(), + "cluster" + ); + + timeout(Duration::from_secs(5), async { + for _ in 0..USAGE_BATCHES { + usage_rx + .next() + .await + .expect("usage metrics output should remain open"); + } + }) + .await + .expect("usage metrics should drain after it resumes"); + + drop(trigger_shutdown); + shutdown_done.await; + assert_eq!(source.await.expect("source task should complete"), Ok(())); + } +} diff --git a/src/sources/mod.rs b/src/sources/mod.rs index 617ea1de7f..9bac304313 100644 --- a/src/sources/mod.rs +++ b/src/sources/mod.rs @@ -56,6 +56,8 @@ pub mod kafka; pub mod kubernetes_logs; #[cfg(feature = "sources-logstash")] pub mod logstash; +#[cfg(feature = "sources-mezmo_analytics")] +pub mod mezmo_analytics; #[cfg(feature = "sources-mezmo_demo_logs")] pub mod mezmo_demo_logs; #[cfg(feature = "sources-mezmo_pipeline_state_variable_change")] diff --git a/src/transforms/mezmo_log_clustering/aggregate.rs b/src/transforms/mezmo_log_clustering/aggregate.rs new file mode 100644 index 0000000000..5b366916a0 --- /dev/null +++ b/src/transforms/mezmo_log_clustering/aggregate.rs @@ -0,0 +1,213 @@ +use std::{collections::HashMap, time::Duration}; + +use chrono::Utc; +use tokio::{sync::mpsc::UnboundedReceiver, time::sleep}; +use vector_lib::{ + event::{LogEvent, Value}, + mezmo::analytics::{self, AnalyticsEventBatch, AnalyticsOutput}, +}; + +use super::{ComponentInfo, LocalId, LogGroupAggregateInfo, LogGroupInfo, store}; + +const MAX_NEW_TEMPLATES_QUEUED: usize = 100; + +pub(crate) async fn aggregate_in_loop( + mut rx: UnboundedReceiver, + agg_window: Duration, +) { + let conn_str = match store::init_db_pool().await { + Ok(conn_str) => { + info!("Starting to store log clustering data in metrics db"); + Some(conn_str) + } + Err(err) => { + error!(message = "There was an error initializing the log clustering db client.", %err); + error!("No log clustering data will be stored in the db."); + None + } + }; + + let mut finished = false; + while !finished { + let mut aggregated: HashMap> = + HashMap::new(); + let timeout = sleep(agg_window); + tokio::pin!(timeout); + let mut new_templates = 0; + + loop { + tokio::select! { + _ = &mut timeout => { + // Break the inner loop, start a new timer + break; + }, + Some(info) = rx.recv() => { + let map = aggregated.entry(info.key).or_default(); + if info.template.is_some() { + new_templates += 1; + } + let aggregated_info = map.entry(info.local_id).or_default(); + aggregated_info.cluster_id = info.cluster_id; + aggregated_info.count += 1; + aggregated_info.size += info.size; + + // Template and annotations are conditionally sent + // Make sure we don't blindly overwrite the existing value + if info.template.is_some() { + aggregated_info.template = info.template; + } + if info.annotation_set.is_some() { + aggregated_info.annotation_set = info.annotation_set; + } + + info.samples.iter().for_each(|s| aggregated_info.samples.push(s.clone())); + + if new_templates > MAX_NEW_TEMPLATES_QUEUED { + break; + } + }, + else => { + // Channel closed + finished = true; + break; + } + } + } + + analytics::publish(|| analytics_batches(&aggregated)); + + if let Some(conn_str) = &conn_str { + store::save(conn_str, aggregated).await; + } + } +} + +fn analytics_batches( + aggregated: &HashMap>, +) -> Vec { + let timestamp = Utc::now(); + let mut clusters = Vec::new(); + let mut samples = Vec::new(); + let mut usage = Vec::new(); + + for (component, aggregates) in aggregated { + let account_id = Value::from(component.account_id.to_string()); + let component_id = Value::from(component.component_id.clone()); + + for aggregate in aggregates.values() { + let common_fields = vector_lib::btreemap!( + "timestamp" => Value::Timestamp(timestamp), + "account_id" => account_id.clone(), + "component_id" => component_id.clone(), + "log_cluster_id" => Value::from(aggregate.cluster_id.clone()) + ); + + if let Some(template) = &aggregate.template { + let mut cluster_fields = common_fields.clone(); + cluster_fields.insert("template".into(), Value::from(template.clone())); + cluster_fields.insert("first_seen_at".into(), Value::Timestamp(timestamp)); + cluster_fields.insert( + "annotations".into(), + aggregate + .annotation_set + .as_ref() + .map_or(Value::Null, |set| { + Value::from( + serde_json::to_value(set) + .expect("annotation sets should always serialize"), + ) + }), + ); + clusters.push(LogEvent::from(Value::Object(cluster_fields))); + + for sample in &aggregate.samples { + let mut sample_fields = common_fields.clone(); + sample_fields.insert("sample".into(), sample.clone()); + samples.push(LogEvent::from(Value::Object(sample_fields))); + } + } + + let mut usage_fields = common_fields; + usage_fields.insert("count".into(), Value::from(aggregate.count)); + usage_fields.insert("size".into(), Value::from(aggregate.size)); + usage.push(LogEvent::from(Value::Object(usage_fields))); + } + } + + [ + (AnalyticsOutput::LogClusters, clusters), + (AnalyticsOutput::LogClusterSamples, samples), + (AnalyticsOutput::LogClusterUsage, usage), + ] + .into_iter() + .filter_map(|(output, events)| { + if events.is_empty() { + None + } else { + Some(AnalyticsEventBatch::new(output, events)) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use uuid::Uuid; + use vector_lib::{event::Value, mezmo::analytics::AnalyticsOutput}; + + use super::*; + + #[test] + fn creates_all_log_cluster_analytics_batches() { + let component = ComponentInfo { + account_id: Uuid::nil(), + component_id: "analysis".into(), + }; + let aggregate = LogGroupAggregateInfo { + cluster_id: "cluster".into(), + count: 2, + size: 20, + template: Some("request <*>".into()), + annotation_set: None, + samples: vec![Value::from("request 42")], + }; + let aggregated = HashMap::from([(component, HashMap::from([(1, aggregate)]))]); + + let batches = analytics_batches(&aggregated); + assert_eq!(batches.len(), 3); + + let cluster = batches + .iter() + .find(|batch| batch.output() == AnalyticsOutput::LogClusters) + .expect("log cluster batch"); + assert_eq!(cluster.events().len(), 1); + assert_eq!( + cluster.events()[0].get("log_cluster_id"), + Some(&Value::from("cluster")) + ); + assert_eq!( + cluster.events()[0].get("template"), + Some(&Value::from("request <*>")) + ); + + let samples = batches + .iter() + .find(|batch| batch.output() == AnalyticsOutput::LogClusterSamples) + .expect("log cluster samples batch"); + assert_eq!(samples.events().len(), 1); + assert_eq!( + samples.events()[0].get("sample"), + Some(&Value::from("request 42")) + ); + + let usage = batches + .iter() + .find(|batch| batch.output() == AnalyticsOutput::LogClusterUsage) + .expect("log cluster usage batch"); + assert_eq!(usage.events().len(), 1); + assert_eq!(usage.events()[0].get("count"), Some(&Value::from(2))); + assert_eq!(usage.events()[0].get("size"), Some(&Value::from(20))); + } +} diff --git a/src/transforms/mezmo_log_clustering/mod.rs b/src/transforms/mezmo_log_clustering/mod.rs index fbc692d521..da58530a86 100644 --- a/src/transforms/mezmo_log_clustering/mod.rs +++ b/src/transforms/mezmo_log_clustering/mod.rs @@ -18,8 +18,8 @@ use uuid::Uuid; use vector_lib::config::{TransformOutput, log_schema}; use vector_lib::configurable::configurable_component; +use crate::transforms::mezmo_log_clustering::aggregate::aggregate_in_loop; use crate::transforms::mezmo_log_clustering::drain::{LocalId, LogClusterStatus}; -use crate::transforms::mezmo_log_clustering::store::save_in_loop; use mezmo::MezmoContext; use vector_lib::event::LogEvent; use vector_lib::usage_metrics::{ @@ -27,6 +27,7 @@ use vector_lib::usage_metrics::{ }; use vrl::value::Value; +mod aggregate; mod drain; mod store; @@ -100,17 +101,17 @@ const fn default_max_log_samples_amount() -> usize { impl_generate_config_from_default!(MezmoLogClusteringConfig); -type DbTransmitter = UnboundedSender; -static ONCE: OnceCell = OnceCell::const_new(); +type AggregateTransmitter = UnboundedSender; +static ONCE: OnceCell = OnceCell::const_new(); #[async_trait::async_trait] #[typetag::serde(name = "mezmo_log_clustering")] impl TransformConfig for MezmoLogClusteringConfig { async fn build(&self, context: &TransformContext) -> crate::Result { - // Create a channel with a db connection pool only once + // Create the aggregation channel only once. let mut account_id = None; let mut component_id = None; - let db_tx = if self.store_metrics { + let aggregate_tx = if self.store_metrics { let Some(mezmo_ctx) = context.mezmo_ctx.as_ref() else { return Err("Cannot store log clustering metrics without a component key".into()); }; @@ -126,10 +127,10 @@ impl TransformConfig for MezmoLogClusteringConfig { let tx = ONCE .get_or_init(move || async move { let (tx, rx) = mpsc::unbounded_channel(); - // Start saving in the background + // Start aggregating in the background. // This task will be running forever, topology changes should not affect it tokio::spawn(async move { - save_in_loop(rx, store_metrics_flush_interval).await; + aggregate_in_loop(rx, store_metrics_flush_interval).await; }); tx @@ -145,7 +146,7 @@ impl TransformConfig for MezmoLogClusteringConfig { self, account_id, component_id, - db_tx, + aggregate_tx, ))) } @@ -171,7 +172,7 @@ struct MezmoLogClustering { transform_status: Option, account_id: Option, component_id: Option, - db_tx: Option, + aggregate_tx: Option, } #[derive(Copy, Clone, PartialEq, Debug)] @@ -197,7 +198,7 @@ impl MezmoLogClustering { config: &MezmoLogClusteringConfig, account_id: Option, component_id: Option, - db_tx: Option, + aggregate_tx: Option, ) -> Self { let similarity_threshold = if config.similarity_threshold > 1.0 || config.similarity_threshold < 0.0 @@ -258,7 +259,7 @@ impl MezmoLogClustering { transform_status: None, account_id, component_id, - db_tx, + aggregate_tx, cluster_field: config.cluster_field.clone(), } } @@ -370,11 +371,11 @@ impl MezmoLogClustering { info.annotation_set = log.as_map().and_then(get_annotations); } - match self.db_tx.as_ref().expect("can't fail").send(info) { + match self.aggregate_tx.as_ref().expect("can't fail").send(info) { Ok(()) => { self.parser.mark_cluster_samples_as_stored(local_id); } - Err(_) => error!("Db channel closed"), + Err(_) => error!("Log clustering aggregation channel closed"), }; } else if status == TransformStatus::AnnotateEvent { let mut cluster = BTreeMap::new(); diff --git a/src/transforms/mezmo_log_clustering/store.rs b/src/transforms/mezmo_log_clustering/store.rs index 903c2ae184..45b6e32cb8 100644 --- a/src/transforms/mezmo_log_clustering/store.rs +++ b/src/transforms/mezmo_log_clustering/store.rs @@ -1,29 +1,24 @@ use crate::internal_events::mezmo_log_clustering::MezmoLogClusteringStore; -use crate::transforms::mezmo_log_clustering::{ - ComponentInfo, LocalId, LogGroupAggregateInfo, LogGroupInfo, -}; +use crate::transforms::mezmo_log_clustering::{ComponentInfo, LocalId, LogGroupAggregateInfo}; use chrono::Utc; use deadpool_postgres::Object; use futures_util::future::join_all; use std::collections::HashMap; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Instant; use std::vec::IntoIter; use tokio::sync::Mutex; -use tokio::sync::mpsc::UnboundedReceiver; -use tokio::time::sleep; use tokio_postgres::Statement; use tokio_postgres::types::{Json, ToSql}; use vector_lib::mezmo; -const MAX_NEW_TEMPLATES_QUEUED: usize = 100; const DB_MAX_PARALLEL_EXECUTIONS: usize = 8; const INSERT_USAGE_QUERY: &str = "INSERT INTO usage_metrics_log_cluster (ts, component_id, log_cluster_id, count, size) VALUES ($1, $2, $3, $4, $5)"; const INSERT_LOG_CLUSTER_QUERY: &str = "INSERT INTO log_clusters (ts, account_id, component_id, log_cluster_id, template, first_seen_at, annotations) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT DO NOTHING"; const INSERT_LOG_CLUSTER_SAMPLES_QUERY: &str = "INSERT INTO log_clusters_samples (ts, account_id, component_id, log_cluster_id, sample) VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING"; -async fn init_db_pool() -> crate::Result { +pub(super) async fn init_db_pool() -> crate::Result { let conn_str = match mezmo::postgres::get_connection_string("metrics") { Ok(conn_str) => conn_str, Err(err) => { @@ -55,74 +50,7 @@ async fn init_db_pool() -> crate::Result { Ok(conn_str) } -pub(crate) async fn save_in_loop(mut rx: UnboundedReceiver, agg_window: Duration) { - let conn_str = match init_db_pool().await { - Ok(conn_str) => conn_str, - Err(err) => { - error!(message = "There was error initializing log clustering db client", %err); - error!("No log clustering data will be stored in the db"); - // Dequeue and ignore - while let Some(_) = rx.recv().await { - // Do nothing - } - return; - } - }; - - info!("Starting to store log clustering data in metrics db"); - - let mut finished = false; - while !finished { - let mut aggregated: HashMap> = - HashMap::new(); - let timeout = sleep(agg_window); - tokio::pin!(timeout); - let mut new_templates = 0; - - loop { - tokio::select! { - _ = &mut timeout => { - // Break the inner loop, start a new timer - break; - }, - Some(info) = rx.recv() => { - let map = aggregated.entry(info.key).or_default(); - if info.template.is_some() { - new_templates += 1; - } - let aggregated_info = map.entry(info.local_id).or_default(); - aggregated_info.cluster_id = info.cluster_id; - aggregated_info.count += 1; - aggregated_info.size += info.size; - - // Template and annotations are conditionally sent - // Make sure we don't blindly overwrite the existing value - if info.template.is_some() { - aggregated_info.template = info.template; - } - if info.annotation_set.is_some() { - aggregated_info.annotation_set = info.annotation_set; - } - - info.samples.iter().for_each(|s| aggregated_info.samples.push(s.clone())); - - if new_templates > MAX_NEW_TEMPLATES_QUEUED { - break; - } - }, - else => { - // Channel closed - finished = true; - break; - } - } - } - - save(&conn_str, aggregated).await; - } -} - -async fn save( +pub(super) async fn save( conn_str: &str, aggregated: HashMap>, ) {