From ca44a3a4b95be14a70d28cab58707addf8c39471 Mon Sep 17 00:00:00 2001
From: Derek Zhang <161868481+GreyLilac09@users.noreply.github.com>
Date: Mon, 7 Jul 2025 11:01:11 -0400
Subject: [PATCH] feat(metrics): add TTL-based cache for metric sets and add
 expire_metrics_secs for Prometheus remote write sink (#23286)

* add ttl policy to metricset

* rust fmt

* rust

* update cue

* fix secs format

* fix float

* component docs

* update comments

* fix fmt

* update comment

* add cleanup and fix timestamp

* remove unnecessary function

* add MetricEntry as struct

* fix lint issues

* use is_some_and

* revert greptime_logs.cue
---
 .../fix_metric_set_memory_growth.feature.md   |   6 +
 src/sinks/prometheus/remote_write/config.rs   |   9 +
 src/sinks/prometheus/remote_write/sink.rs     |   4 +-
 src/sinks/util/buffer/metrics/normalize.rs    | 227 +++++++++++++++---
 src/sinks/util/builder.rs                     |  20 +-
 src/sinks/util/normalizer.rs                  |  10 +-
 .../sinks/base/prometheus_remote_write.cue    |  11 +
 7 files changed, 242 insertions(+), 45 deletions(-)
 create mode 100644 changelog.d/fix_metric_set_memory_growth.feature.md

diff --git a/changelog.d/fix_metric_set_memory_growth.feature.md b/changelog.d/fix_metric_set_memory_growth.feature.md
new file mode 100644
index 0000000000000..525648566668a
--- /dev/null
+++ b/changelog.d/fix_metric_set_memory_growth.feature.md
@@ -0,0 +1,6 @@
+- Add a TTL-based cache for metrics sets
+- Add `expire_metrics_secs` config for Prometheus remote write sink which uses the TTL-based cache
+- This fixes an issue where incremental metrics are preserved for the lifetime of Vector's runtime, which causes
+  indefinite memory growth
+
+authors: GreyLilac09
diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index 9a61d0be6baa0..03f6d2bee2905 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -93,6 +93,14 @@ pub struct RemoteWriteConfig {
     #[configurable(metadata(docs::advanced))]
     pub tenant_id: Option<Template>,
 
+    /// The amount of time, in seconds, that incremental metrics will persist in the internal metrics cache
+    /// after having not been updated before they expire and are removed.
+    ///
+    /// If unset, sending unique incremental metrics to this sink will cause indefinite memory growth.
+    #[serde(skip_serializing_if = "crate::serde::is_default")]
+    #[configurable(metadata(docs::common = false, docs::required = false))]
+    pub expire_metrics_secs: Option<f64>,
+
     #[configurable(derived)]
     pub tls: Option<TlsConfig>,
 
@@ -202,6 +210,7 @@ impl SinkConfig for RemoteWriteConfig {
             buckets,
             quantiles,
             default_namespace,
+            expire_metrics_secs: self.expire_metrics_secs,
             service,
         };
 
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 093cb7e0bc26b..8ce8dfa6a43f3 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -149,6 +149,7 @@ pub(super) struct RemoteWriteSink<S> {
     pub(super) default_namespace: Option<String>,
     pub(super) buckets: Vec<f64>,
     pub(super) quantiles: Vec<f64>,
+    pub(super) expire_metrics_secs: Option<f64>,
     pub(super) service: S,
 }
 
@@ -172,10 +173,11 @@ where
         let batch_settings = self.batch_settings;
         let tenant_id = self.tenant_id.clone();
         let service = self.service;
+        let expire_metrics_secs = self.expire_metrics_secs;
 
         input
             .filter_map(|event| future::ready(event.try_into_metric()))
-            .normalized_with_default::<PrometheusMetricNormalize>()
+            .normalized_with_ttl::<PrometheusMetricNormalize>(expire_metrics_secs)
             .filter_map(move |event| {
                 future::ready(make_remote_write_event(tenant_id.as_ref(), event))
             })
diff --git a/src/sinks/util/buffer/metrics/normalize.rs b/src/sinks/util/buffer/metrics/normalize.rs
index 57b77413b7f41..102ba67c2818f 100644
--- a/src/sinks/util/buffer/metrics/normalize.rs
+++ b/src/sinks/util/buffer/metrics/normalize.rs
@@ -1,5 +1,7 @@
 use indexmap::IndexMap;
 
+use std::time::{Duration, Instant};
+
 use vector_lib::event::{
     metric::{MetricData, MetricSeries},
     EventMetadata, Metric, MetricKind,
@@ -46,6 +48,14 @@ pub struct MetricNormalizer<N> {
 }
 
 impl<N> MetricNormalizer<N> {
+    /// Creates a new normalizer with TTL policy.
+    pub fn with_ttl(normalizer: N, ttl: TtlPolicy) -> Self {
+        Self {
+            state: MetricSet::with_ttl_policy(ttl),
+            normalizer,
+        }
+    }
+
     /// Gets a mutable reference to the current metric state for this normalizer.
     pub fn get_state_mut(&mut self) -> &mut MetricSet {
         &mut self.state
@@ -79,46 +89,178 @@ impl<N> From<N> for MetricNormalizer<N> {
     }
 }
 
-type MetricEntry = (MetricData, EventMetadata);
+/// Represents a stored metric entry with its data, metadata, and optional timestamp.
+#[derive(Clone, Debug)]
+pub struct MetricEntry {
+    /// The metric data containing the value and kind
+    pub data: MetricData,
+    /// Event metadata associated with this metric
+    pub metadata: EventMetadata,
+    /// Optional timestamp for TTL tracking
+    pub timestamp: Option<Instant>,
+}
+
+impl MetricEntry {
+    /// Creates a new MetricEntry with the given data, metadata, and timestamp.
+    pub const fn new(
+        data: MetricData,
+        metadata: EventMetadata,
+        timestamp: Option<Instant>,
+    ) -> Self {
+        Self {
+            data,
+            metadata,
+            timestamp,
+        }
+    }
+
+    /// Creates a new MetricEntry from a Metric and optional timestamp.
+    pub fn from_metric(metric: Metric, timestamp: Option<Instant>) -> (MetricSeries, Self) {
+        let (series, data, metadata) = metric.into_parts();
+        let entry = Self::new(data, metadata, timestamp);
+        (series, entry)
+    }
+
+    /// Converts this entry back to a Metric with the given series.
+    pub fn into_metric(self, series: MetricSeries) -> Metric {
+        Metric::from_parts(series, self.data, self.metadata)
+    }
+
+    /// Updates this entry's timestamp.
+    pub const fn update_timestamp(&mut self, timestamp: Option<Instant>) {
+        self.timestamp = timestamp;
+    }
+}
+
+/// Configuration for automatic cleanup of expired entries.
+#[derive(Clone, Debug)]
+pub struct TtlPolicy {
+    /// Time-to-live for entries
+    pub ttl: Duration,
+    /// How often to run cleanup
+    pub cleanup_interval: Duration,
+    /// Last time cleanup was performed
+    pub(crate) last_cleanup: Instant,
+}
+
+impl TtlPolicy {
+    /// Creates a new cleanup configuration with TTL.
+    /// Cleanup interval defaults to TTL/10 with a 10-second minimum.
+    pub fn new(ttl: Duration) -> Self {
+        Self {
+            ttl,
+            cleanup_interval: ttl.div_f32(10.0).max(Duration::from_secs(10)),
+            last_cleanup: Instant::now(),
+        }
+    }
+
+    /// Checks if it's time to run cleanup.
+    pub fn should_cleanup(&self) -> bool {
+        Instant::now().duration_since(self.last_cleanup) >= self.cleanup_interval
+    }
+
+    /// Marks cleanup as having been performed.
+    pub fn mark_cleanup_done(&mut self) {
+        self.last_cleanup = Instant::now();
+    }
+}
 
 /// Metric storage for use with normalization.
 ///
 /// This is primarily a wrapper around [`IndexMap`] (to ensure insertion order
 /// is maintained) with convenience methods to make it easier to perform
-/// normalization-specific operations.
-#[derive(Clone, Default, Debug)]
-pub struct MetricSet(IndexMap<MetricSeries, MetricEntry>);
+/// normalization-specific operations. It also includes an optional TTL policy
+/// to automatically expire old entries.
+#[derive(Clone, Debug, Default)]
+pub struct MetricSet {
+    inner: IndexMap<MetricSeries, MetricEntry>,
+    ttl_policy: Option<TtlPolicy>,
+}
 
 impl MetricSet {
-    /// Creates an empty `MetricSet` with the specified capacity.
-    ///
-    /// The metric set will be able to hold at least `capacity` elements without reallocating. If `capacity` is 0, the
-    /// metric set will not allocate.
+    /// Creates an empty MetricSet with the specified capacity.
     pub fn with_capacity(capacity: usize) -> Self {
-        Self(IndexMap::with_capacity(capacity))
+        Self {
+            inner: IndexMap::with_capacity(capacity),
+            ttl_policy: None,
+        }
+    }
+
+    /// Creates a MetricSet with custom cleanup configuration.
+    pub fn with_ttl_policy(ttl_policy: TtlPolicy) -> Self {
+        Self {
+            inner: IndexMap::default(),
+            ttl_policy: Some(ttl_policy),
+        }
+    }
+
+    /// Gets a reference to the TTL policy configuration.
+    pub const fn ttl_policy(&self) -> Option<&TtlPolicy> {
+        self.ttl_policy.as_ref()
+    }
+
+    /// Gets a mutable reference to the TTL policy configuration.
+    pub const fn ttl_policy_mut(&mut self) -> Option<&mut TtlPolicy> {
+        self.ttl_policy.as_mut()
+    }
+
+    /// Perform periodic cleanup if enough time has passed since the last cleanup
+    fn maybe_cleanup(&mut self) {
+        // Return early if no cleanup is needed
+        if !self
+            .ttl_policy()
+            .is_some_and(|config| config.should_cleanup())
+        {
+            return;
+        }
+        self.cleanup_expired();
+        if let Some(config) = self.ttl_policy_mut() {
+            config.mark_cleanup_done();
+        }
+    }
+
+    /// Removes expired entries based on TTL if configured.
+    fn cleanup_expired(&mut self) {
+        let now = Instant::now();
+        if let Some(config) = &self.ttl_policy {
+            self.inner.retain(|_, entry| match entry.timestamp {
+                Some(ts) => now.duration_since(ts) < config.ttl,
+                None => true,
+            });
+        }
     }
 
     /// Returns the number of elements in the set.
     pub fn len(&self) -> usize {
-        self.0.len()
+        self.inner.len()
+    }
+
+    fn create_timestamp(&self) -> Option<Instant> {
+        match self.ttl_policy() {
+            Some(_) => Some(Instant::now()),
+            _ => None,
+        }
     }
 
-    /// Returns `true` if the set contains no elements.
+    /// Returns true if the set contains no elements.
     pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
+        self.inner.is_empty()
     }
 
-    /// Consumes this `MetricSet` and returns a vector of `Metric`.
-    pub fn into_metrics(self) -> Vec<Metric> {
-        self.0
+    /// Consumes this MetricSet and returns a vector of Metric.
+    pub fn into_metrics(mut self) -> Vec<Metric> {
+        // Always cleanup on final consumption
+        self.cleanup_expired();
+        self.inner
             .into_iter()
-            .map(|(series, (data, metadata))| Metric::from_parts(series, data, metadata))
+            .map(|(series, entry)| entry.into_metric(series))
             .collect()
     }
 
     /// Either pass the metric through as-is if absolute, or convert it
     /// to absolute if incremental.
     pub fn make_absolute(&mut self, metric: Metric) -> Option<Metric> {
+        self.maybe_cleanup();
         match metric.kind() {
             MetricKind::Absolute => Some(metric),
             MetricKind::Incremental => Some(self.incremental_to_absolute(metric)),
@@ -128,6 +270,7 @@ impl MetricSet {
     /// Either convert the metric to incremental if absolute, or
     /// aggregate it with any previous value if already incremental.
     pub fn make_incremental(&mut self, metric: Metric) -> Option<Metric> {
+        self.maybe_cleanup();
         match metric.kind() {
             MetricKind::Absolute => self.absolute_to_incremental(metric),
             MetricKind::Incremental => Some(metric),
@@ -138,23 +281,21 @@ impl MetricSet {
     /// state buffer to keep track of the value throughout the entire
     /// application uptime.
     fn incremental_to_absolute(&mut self, mut metric: Metric) -> Metric {
-        match self.0.get_mut(metric.series()) {
+        let timestamp = self.create_timestamp();
+        match self.inner.get_mut(metric.series()) {
             Some(existing) => {
-                if existing.0.value.add(metric.value()) {
-                    metric = metric.with_value(existing.0.value.clone());
+                if existing.data.value.add(metric.value()) {
+                    metric = metric.with_value(existing.data.value.clone());
+                    existing.update_timestamp(timestamp);
                 } else {
                     // Metric changed type, store this as the new reference value
-                    self.0.insert(
-                        metric.series().clone(),
-                        (metric.data().clone(), EventMetadata::default()),
-                    );
+                    let (series, entry) = MetricEntry::from_metric(metric.clone(), timestamp);
+                    self.inner.insert(series, entry);
                 }
             }
             None => {
-                self.0.insert(
-                    metric.series().clone(),
-                    (metric.data().clone(), EventMetadata::default()),
-                );
+                let (series, entry) = MetricEntry::from_metric(metric.clone(), timestamp);
+                self.inner.insert(series, entry);
             }
         }
         metric.into_absolute()
@@ -182,42 +323,47 @@ impl MetricSet {
         // introducing a small amount of lag before a metric is emitted by having to wait to see it
         // again, but this is a behavior we have to observe for sinks that can only handle
         // incremental updates.
-        match self.0.get_mut(metric.series()) {
+        let timestamp = self.create_timestamp();
+        match self.inner.get_mut(metric.series()) {
             Some(reference) => {
                 let new_value = metric.value().clone();
                 // From the stored reference value, emit an increment
-                if metric.subtract(&reference.0) {
-                    reference.0.value = new_value;
+                if metric.subtract(&reference.data) {
+                    reference.data.value = new_value;
+                    reference.update_timestamp(timestamp);
                     Some(metric.into_incremental())
                 } else {
                     // Metric changed type, store this and emit nothing
-                    self.insert(metric);
+                    self.insert(metric, timestamp);
                     None
                 }
             }
             None => {
                 // No reference so store this and emit nothing
-                self.insert(metric);
+                self.insert(metric, timestamp);
                 None
             }
         }
     }
 
-    fn insert(&mut self, metric: Metric) {
-        let (series, data, metadata) = metric.into_parts();
-        self.0.insert(series, (data, metadata));
+    fn insert(&mut self, metric: Metric, timestamp: Option<Instant>) {
+        let (series, entry) = MetricEntry::from_metric(metric, timestamp);
+        self.inner.insert(series, entry);
     }
 
     pub fn insert_update(&mut self, metric: Metric) {
+        self.maybe_cleanup();
+        let timestamp = self.create_timestamp();
         let update = match metric.kind() {
             MetricKind::Absolute => Some(metric),
             MetricKind::Incremental => {
                 // Incremental metrics update existing entries, if present
-                match self.0.get_mut(metric.series()) {
+                match self.inner.get_mut(metric.series()) {
                     Some(existing) => {
                         let (series, data, metadata) = metric.into_parts();
-                        if existing.0.update(&data) {
-                            existing.1.merge(metadata);
+                        if existing.data.update(&data) {
+                            existing.metadata.merge(metadata);
+                            existing.update_timestamp(timestamp);
                             None
                         } else {
                             warn!(message = "Metric changed type, dropping old value.", %series);
@@ -229,7 +375,7 @@ impl MetricSet {
             }
         };
         if let Some(metric) = update {
-            self.insert(metric);
+            self.insert(metric, timestamp);
         }
     }
 
@@ -237,6 +383,7 @@ impl MetricSet {
     ///
     /// If the series existed and was removed, returns `true`.  Otherwise, `false`.
     pub fn remove(&mut self, series: &MetricSeries) -> bool {
-        self.0.shift_remove(series).is_some()
+        self.maybe_cleanup();
+        self.inner.shift_remove(series).is_some()
     }
 }
diff --git a/src/sinks/util/builder.rs b/src/sinks/util/builder.rs
index 45c01ce41dcdd..8750e2ae862a9 100644
--- a/src/sinks/util/builder.rs
+++ b/src/sinks/util/builder.rs
@@ -1,3 +1,6 @@
+use futures_util::{stream::Map, Stream, StreamExt};
+use pin_project::pin_project;
+use std::time::Duration;
 use std::{
     convert::Infallible,
     fmt,
@@ -8,9 +11,6 @@ use std::{
     sync::Arc,
     task::{Context, Poll},
 };
-
-use futures_util::{stream::Map, Stream, StreamExt};
-use pin_project::pin_project;
 use tower::Service;
 use tracing::Span;
 use vector_lib::stream::{
@@ -220,6 +220,20 @@ pub trait SinkBuilderExt: Stream {
         Normalizer::new(self, N::default())
     }
 
+    /// Normalizes a stream of [`Metric`] events with a normalizer and an optional TTL.
+    fn normalized_with_ttl<N>(self, maybe_ttl_secs: Option<f64>) -> Normalizer<Self, N>
+    where
+        Self: Stream<Item = Metric> + Unpin + Sized,
+        N: MetricNormalize + Default,
+    {
+        match maybe_ttl_secs {
+            None => Normalizer::new(self, N::default()),
+            Some(ttl) => {
+                Normalizer::new_with_ttl(self, N::default(), Duration::from_secs(ttl as u64))
+            }
+        }
+    }
+
     /// Creates a [`Driver`] that uses the configured event stream as the input to the given
     /// service.
     ///
diff --git a/src/sinks/util/normalizer.rs b/src/sinks/util/normalizer.rs
index c1ab6cd047913..a8560264e5e3b 100644
--- a/src/sinks/util/normalizer.rs
+++ b/src/sinks/util/normalizer.rs
@@ -1,13 +1,14 @@
 use std::{
     pin::Pin,
     task::{ready, Context, Poll},
+    time::Duration,
 };
 
 use futures_util::{stream::Fuse, Stream, StreamExt};
 use pin_project::pin_project;
 use vector_lib::event::Metric;
 
-use super::buffer::metrics::{MetricNormalize, MetricNormalizer};
+use super::buffer::metrics::{MetricNormalize, MetricNormalizer, TtlPolicy};
 
 #[pin_project]
 pub struct Normalizer<St, N>
@@ -29,6 +30,13 @@ where
             normalizer: MetricNormalizer::from(normalizer),
         }
     }
+
+    pub fn new_with_ttl(stream: St, normalizer: N, ttl: Duration) -> Self {
+        Self {
+            stream: stream.fuse(),
+            normalizer: MetricNormalizer::with_ttl(normalizer, TtlPolicy::new(ttl)),
+        }
+    }
 }
 
 impl<St, N> Stream for Normalizer<St, N>
diff --git a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
index 4ca726cbd445b..1a0c3fa7c8533 100644
--- a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
+++ b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
@@ -308,6 +308,17 @@ base: components: sinks: prometheus_remote_write: configuration: {
 		required: true
 		type: string: examples: ["https://localhost:8087/api/v1/write"]
 	}
+	expire_metrics_secs: {
+		common: false
+		description: """
+			The amount of time, in seconds, that incremental metrics will persist in the internal metrics cache
+			after having not been updated before they expire and are removed.
+
+			If unset, sending unique incremental metrics to this sink will cause indefinite memory growth.
+			"""
+		required: false
+		type: float: {}
+	}
 	quantiles: {
 		description: """
 			Quantiles to use for aggregating [distribution][dist_metric_docs] metrics into a summary.