fix(enrichment tables): preserve memory enrichment table state on reload

changelog.d/25547_memory_table_state_loss.fix.md

@@ -0,0 +1,3 @@
+Fixed an issue where memory enrichment table state could be lost or detached from the sink component on configuration reload.
+
+authors: esensar Quad9DNS

lib/vector-vrl/enrichment/src/lib.rs

@@ -141,6 +141,11 @@ pub trait Table: DynClone {
 
     /// Returns true if the underlying data has changed and the table needs reloading.
     fn needs_reload(&self) -> bool;
+
+    /// Extracts state from this table
+    fn extract_state(&self) -> Option<Box<dyn std::any::Any + Send + Sync>> {
+        None
+    }
 }
 
 dyn_clone::clone_trait_object!(Table);

lib/vector-vrl/enrichment/src/tables.rs

@@ -196,6 +196,14 @@ impl TableRegistry {
             None => true,
         }
     }
+
+    /// Extracts state from the table if available.
+    pub fn extract_state(&self, table: &str) -> Option<Box<dyn std::any::Any + Send + Sync>> {
+        match &**self.tables.load() {
+            Some(tables) => tables.get(table).and_then(|t| t.extract_state()),
+            None => None,
+        }
+    }
 }
 
 impl std::fmt::Debug for TableRegistry {

src/app.rs

@@ -617,11 +617,19 @@ pub async fn load_configs(
         for (name, table) in config.enrichment_tables() {
             let files = table.inner.files_to_watch();
             let component_config = ComponentConfig::new(
-                files.into_iter().cloned().collect(),
+                files.clone().into_iter().cloned().collect(),
                 name.clone(),
                 ComponentType::EnrichmentTable,
             );
             watched_component_paths.push(component_config);
+            if table.as_sink(name).is_some() {
+                let sink_component_config = ComponentConfig::new(
+                    files.into_iter().cloned().collect(),
+                    name.clone(),
+                    ComponentType::Sink,
+                );
+                watched_component_paths.push(sink_component_config);
+            }
         }
 
         info!(

src/config/diff.rs

@@ -29,6 +29,7 @@ impl ConfigDiff {
             enrichment_tables: Difference::from_enrichment_tables(
                 &old.enrichment_tables,
                 &new.enrichment_tables,
+                &components_to_reload,
             ),
             components_to_reload,
         }
@@ -56,15 +57,15 @@ impl ConfigDiff {
         self.sources.is_changed(key)
             || self.transforms.is_changed(key)
             || self.sinks.is_changed(key)
-            || self.enrichment_tables.contains(key)
+            || self.enrichment_tables.is_changed(key)
     }
 
     /// Checks whether the given component is removed.
     pub fn is_removed(&self, key: &ComponentKey) -> bool {
         self.sources.is_removed(key)
             || self.transforms.is_removed(key)
             || self.sinks.is_removed(key)
-            || self.enrichment_tables.contains(key)
+            || self.enrichment_tables.is_removed(key)
     }
 }
 
@@ -116,6 +117,7 @@ impl Difference {
     fn from_enrichment_tables(
         old: &IndexMap<ComponentKey, EnrichmentTableOuter<OutputId>>,
         new: &IndexMap<ComponentKey, EnrichmentTableOuter<OutputId>>,
+        need_change: &HashSet<ComponentKey>,
     ) -> Self {
         let old_table_keys = extract_table_component_keys(old);
         let new_table_keys = extract_table_component_keys(new);
@@ -131,7 +133,7 @@ impl Difference {
                 // which can iterate in varied orders.
                 let old_value = serde_json::to_value(&old[*table_key]).unwrap();
                 let new_value = serde_json::to_value(&new[*table_key]).unwrap();
-                old_value != new_value
+                old_value != new_value || need_change.contains(*table_key)
             })
             .cloned()
             .map(|(_table_key, derived_component_key)| derived_component_key)
@@ -305,6 +307,7 @@ mod tests {
         let diff = Difference::from_enrichment_tables(
             &old_config.enrichment_tables,
             &new_config.enrichment_tables,
+            &Default::default(),
         );
 
         assert_eq!(diff.to_add, HashSet::from_iter(["memory_table_new".into()]));

src/config/enrichment_table.rs

@@ -112,7 +112,7 @@ where
 /// Generalized interface for describing and building enrichment table components.
 #[enum_dispatch]
 pub trait EnrichmentTableConfig: NamedComponent + core::fmt::Debug + Send + Sync {
-    /// Builds the enrichment table with the given globals.
+    /// Builds the enrichment table with the given globals and previous table state.
     ///
     /// If the enrichment table is built successfully, `Ok(...)` is returned containing the
     /// enrichment table.
@@ -124,8 +124,14 @@ pub trait EnrichmentTableConfig: NamedComponent + core::fmt::Debug + Send + Sync
     async fn build(
         &self,
         globals: &GlobalOptions,
+        prev_state: Option<Box<dyn std::any::Any + Send + Sync>>,
     ) -> crate::Result<Box<dyn vector_lib::enrichment::Table + Send + Sync>>;
 
+    /// Checks whether this table wants previous state, to try and restore it.
+    fn wants_previous_state(&self) -> bool {
+        false
+    }
+
     fn sink_config(
         &self,
         _default_key: &ComponentKey,

src/config/mod.rs

@@ -77,7 +77,7 @@ pub use vector_lib::{
 };
 
 #[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)]
-// // This is not a comprehensive set; variants are added as needed.
+// This is not a comprehensive set; variants are added as needed.
 pub enum ComponentType {
     Transform,
     Sink,

src/enrichment_tables/file.rs

@@ -239,6 +239,7 @@ impl EnrichmentTableConfig for FileConfig {
     async fn build(
         &self,
         globals: &crate::config::GlobalOptions,
+        _prev_state: Option<Box<dyn std::any::Any + Send + Sync>>,
     ) -> crate::Result<Box<dyn Table + Send + Sync>> {
         Ok(Box::new(File::new(
             self.clone(),

src/enrichment_tables/geoip.rs

@@ -101,6 +101,7 @@ impl EnrichmentTableConfig for GeoipConfig {
     async fn build(
         &self,
         _: &crate::config::GlobalOptions,
+        _: Option<Box<dyn std::any::Any + Send + Sync>>,
     ) -> crate::Result<Box<dyn Table + Send + Sync>> {
         Ok(Box::new(Geoip::new(self.clone())?))
     }

src/enrichment_tables/memory/config.rs

@@ -68,11 +68,27 @@ pub struct MemoryConfig {
     #[configurable(derived)]
     #[serde(default)]
     pub ttl_field: OptionalValuePath,
+    /// Behavior for memory table state on configuration reload.
+    #[configurable(derived)]
+    #[serde(default)]
+    pub reload_behavior: ReloadBehavior,
 
     #[serde(skip)]
     memory: Arc<Mutex<Option<Box<Memory>>>>,
 }
 
+/// Behavior for memory enrichment table state on configuration reload.
+#[configurable_component]
+#[derive(Clone, Default)]
+#[serde(rename_all = "kebab-case")]
+pub enum ReloadBehavior {
+    /// Always clear state on configuration reload.
+    #[default]
+    ClearState,
+    /// Try to preserve state when possible.
+    PreserveState,
+}
+
 /// Configuration for memory enrichment table source functionality.
 #[configurable_component]
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -123,6 +139,7 @@ impl Default for MemoryConfig {
             source_config: None,
             internal_metrics: InternalMetricsConfig::default(),
             ttl_field: OptionalValuePath::none(),
+            reload_behavior: Default::default(),
         }
     }
 }
@@ -136,10 +153,19 @@ const fn default_scan_interval() -> NonZeroU64 {
 }
 
 impl MemoryConfig {
-    pub(super) async fn get_or_build_memory(&self) -> Memory {
+    pub(super) async fn get_or_build_memory(
+        &self,
+        prev_state: Option<Box<dyn std::any::Any + Send + Sync>>,
+    ) -> Memory {
         let mut boxed_memory = self.memory.lock().await;
         *boxed_memory
-            .get_or_insert_with(|| Box::new(Memory::new(self.clone())))
+            .get_or_insert_with(|| {
+                if let Some(prev) = prev_state {
+                    Box::new(Memory::from_previous_state(self.clone(), prev))
+                } else {
+                    Box::new(Memory::new(self.clone()))
+                }
+            })
             .clone()
     }
 }
@@ -148,8 +174,13 @@ impl EnrichmentTableConfig for MemoryConfig {
     async fn build(
         &self,
         _globals: &crate::config::GlobalOptions,
+        prev_state: Option<Box<dyn std::any::Any + Send + Sync>>,
     ) -> crate::Result<Box<dyn Table + Send + Sync>> {
-        Ok(Box::new(self.get_or_build_memory().await))
+        Ok(Box::new(self.get_or_build_memory(prev_state).await))
+    }
+
+    fn wants_previous_state(&self) -> bool {
+        matches!(self.reload_behavior, ReloadBehavior::PreserveState)
     }
 
     fn sink_config(
@@ -177,7 +208,7 @@ impl EnrichmentTableConfig for MemoryConfig {
 #[typetag::serde(name = "memory_enrichment_table")]
 impl SinkConfig for MemoryConfig {
     async fn build(&self, _cx: SinkContext) -> crate::Result<(VectorSink, Healthcheck)> {
-        let sink = VectorSink::from_event_streamsink(self.get_or_build_memory().await);
+        let sink = VectorSink::from_event_streamsink(self.get_or_build_memory(None).await);
 
         Ok((sink, future::ok(()).boxed()))
     }
@@ -195,7 +226,7 @@ impl SinkConfig for MemoryConfig {
 #[typetag::serde(name = "memory_enrichment_table")]
 impl SourceConfig for MemoryConfig {
     async fn build(&self, cx: SourceContext) -> crate::Result<Source> {
-        let memory = self.get_or_build_memory().await;
+        let memory = self.get_or_build_memory(None).await;
 
         let log_namespace = cx.log_namespace(self.log_namespace);
 

src/enrichment_tables/memory/table.rs

@@ -142,6 +142,25 @@ impl Memory {
         }
     }
 
+    /// Creates a new [Memory] based on the provided config and previous state.
+    pub fn from_previous_state(
+        config: MemoryConfig,
+        prev_state: Box<dyn std::any::Any + Send + Sync>,
+    ) -> Self {
+        if let Ok(prev_memory) = prev_state.downcast::<Memory>() {
+            Self {
+                config,
+                read_handle_factory: prev_memory.read_handle_factory,
+                read_handle: prev_memory.read_handle,
+                write_handle: prev_memory.write_handle,
+                expired_items_sender: prev_memory.expired_items_sender,
+                expired_items_receiver: prev_memory.expired_items_receiver,
+            }
+        } else {
+            Self::new(config)
+        }
+    }
+
     pub(super) fn get_read_handle(&self) -> &evmap::ReadHandle<String, MemoryEntry> {
         self.read_handle
             .get_or(|| self.read_handle_factory.handle())
@@ -386,6 +405,12 @@ impl Table for Memory {
     fn needs_reload(&self) -> bool {
         false
     }
+
+    fn extract_state(&self) -> Option<Box<dyn std::any::Any + Send + Sync>> {
+        let writer = self.write_handle.lock().expect("mutex poisoned");
+        self.flush(writer);
+        Some(Box::new(self.clone()))
+    }
 }
 
 impl std::fmt::Debug for Memory {
@@ -465,6 +490,7 @@ mod tests {
 
     use super::*;
     use crate::{
+        config::EnrichmentTableConfig,
         enrichment_tables::memory::{
             config::MemorySourceConfig, internal_events::InternalMetricsConfig,
         },
@@ -501,6 +527,43 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn extract_state_preserves_data() {
+        let memory = Memory::new(Default::default());
+        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
+
+        let condition = Condition::Equals {
+            field: "key",
+            value: Value::from("test_key"),
+        };
+
+        let expected = ObjectMap::from([
+            ("key".into(), Value::from("test_key")),
+            ("ttl".into(), Value::from(memory.config.ttl)),
+            ("value".into(), Value::from(5)),
+        ]);
+        assert_eq!(
+            Ok(expected.clone()),
+            memory.find_table_row(
+                Case::Sensitive,
+                std::slice::from_ref(&condition),
+                None,
+                None,
+                None
+            )
+        );
+
+        // Now build a new table using old state
+        let new_memory = MemoryConfig::default()
+            .build(&Default::default(), memory.extract_state())
+            .await
+            .unwrap();
+        assert_eq!(
+            Ok(expected),
+            new_memory.find_table_row(Case::Sensitive, &[condition], None, None, None)
+        );
+    }
+
     #[test]
     fn calculates_ttl() {
         let ttl = 100;
@@ -1023,7 +1086,7 @@ mod tests {
             export_expired_items: false,
             source_key: "test".to_string(),
         });
-        let memory = memory_config.get_or_build_memory().await;
+        let memory = memory_config.get_or_build_memory(None).await;
         memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
 
         let mut events: Vec<Event> = run_and_assert_source_compliance(

src/enrichment_tables/mmdb.rs

@@ -36,6 +36,7 @@ impl EnrichmentTableConfig for MmdbConfig {
     async fn build(
         &self,
         _: &crate::config::GlobalOptions,
+        _: Option<Box<dyn std::any::Any + Send + Sync>>,
     ) -> crate::Result<Box<dyn Table + Send + Sync>> {
         Ok(Box::new(Mmdb::new(self.clone())?))
     }