vector/enrichment_tables/memory/
table.rs

1#![allow(unsafe_op_in_unsafe_fn)] // TODO review ShallowCopy usage code and fix properly.
2
3use std::{
4    num::NonZeroU64,
5    pin::Pin,
6    sync::{Arc, Mutex, MutexGuard},
7    time::{Duration, Instant},
8};
9
10use async_trait::async_trait;
11use bytes::Bytes;
12use evmap::{
13    shallow_copy::CopyValue,
14    {self},
15};
16use evmap_derive::ShallowCopy;
17use futures::{
18    Stream, StreamExt,
19    stream::{self, BoxStream},
20};
21use thread_local::ThreadLocal;
22use tokio::{
23    sync::broadcast::{Receiver, Sender},
24    time::interval,
25};
26use tokio_stream::wrappers::IntervalStream;
27use vector_lib::{
28    ByteSizeOf, EstimatedJsonEncodedSizeOf,
29    config::LogNamespace,
30    enrichment::{Case, Condition, Error, IndexHandle, InternalError, Table},
31    event::{Event, EventStatus, Finalizable},
32    internal_event::{
33        ByteSize, BytesSent, CountByteSize, EventsSent, InternalEventHandle, Output, Protocol,
34    },
35    shutdown::ShutdownSignal,
36    sink::StreamSink,
37};
38use vrl::value::{KeyString, ObjectMap, Value};
39
40use super::source::MemorySource;
41use crate::{
42    SourceSender,
43    enrichment_tables::memory::{
44        MemoryConfig,
45        internal_events::{
46            MemoryEnrichmentTableFlushed, MemoryEnrichmentTableInsertFailed,
47            MemoryEnrichmentTableInserted, MemoryEnrichmentTableRead,
48            MemoryEnrichmentTableReadFailed, MemoryEnrichmentTableTtlExpired,
49        },
50    },
51};
52
53/// Single memory entry containing the value and TTL
54#[derive(Clone, Eq, PartialEq, Hash, ShallowCopy)]
55pub struct MemoryEntry {
56    value: String,
57    update_time: CopyValue<Instant>,
58    ttl: u64,
59}
60
61impl ByteSizeOf for MemoryEntry {
62    fn allocated_bytes(&self) -> usize {
63        self.value.size_of()
64    }
65}
66
67impl MemoryEntry {
68    pub(super) fn as_object_map(&self, now: Instant, key: &str) -> Result<ObjectMap, Error> {
69        let ttl = self
70            .ttl
71            .saturating_sub(now.duration_since(*self.update_time).as_secs());
72        Ok(ObjectMap::from([
73            (
74                KeyString::from("key"),
75                Value::Bytes(Bytes::copy_from_slice(key.as_bytes())),
76            ),
77            (
78                KeyString::from("value"),
79                // Unreachable in normal operation: `value` was serialized by `handle_value`.
80                serde_json::from_str::<Value>(&self.value).map_err(|source| Error::Internal {
81                    source: InternalError::FailedToDecode {
82                        details: source.to_string(),
83                    },
84                })?,
85            ),
86            (
87                KeyString::from("ttl"),
88                Value::Integer(ttl.try_into().unwrap_or(i64::MAX)),
89            ),
90        ]))
91    }
92
93    fn expired(&self, now: Instant) -> bool {
94        now.duration_since(*self.update_time).as_secs() > self.ttl
95    }
96}
97
98#[derive(Default)]
99struct MemoryMetadata {
100    byte_size: u64,
101}
102
103/// [`MemoryEntry`] combined with its key
104#[derive(Clone)]
105pub(super) struct MemoryEntryPair {
106    /// Key of this entry
107    pub(super) key: String,
108    /// The value of this entry
109    pub(super) entry: MemoryEntry,
110}
111
112// Used to ensure that these 2 are locked together
113pub(super) struct MemoryWriter {
114    pub(super) write_handle: evmap::WriteHandle<String, MemoryEntry>,
115    metadata: MemoryMetadata,
116}
117
118/// A struct that implements [vector_lib::enrichment::Table] to handle loading enrichment data from a memory structure.
119pub struct Memory {
120    read_handle_factory: evmap::ReadHandleFactory<String, MemoryEntry>,
121    read_handle: ThreadLocal<evmap::ReadHandle<String, MemoryEntry>>,
122    pub(super) write_handle: Arc<Mutex<MemoryWriter>>,
123    pub(super) config: MemoryConfig,
124    #[allow(dead_code)]
125    expired_items_receiver: Receiver<Vec<MemoryEntryPair>>,
126    expired_items_sender: Sender<Vec<MemoryEntryPair>>,
127}
128
129impl Memory {
130    /// Creates a new [Memory] based on the provided config.
131    pub fn new(config: MemoryConfig) -> Self {
132        let (read_handle, write_handle) = evmap::new();
133        // Buffer could only be used if source is stuck exporting available items, but in that case,
134        // publishing will not happen either, because the lock would be held, so this buffer is not
135        // that important
136        let (expired_tx, expired_rx) = tokio::sync::broadcast::channel(5);
137        Self {
138            config,
139            read_handle_factory: read_handle.factory(),
140            read_handle: ThreadLocal::new(),
141            write_handle: Arc::new(Mutex::new(MemoryWriter {
142                write_handle,
143                metadata: MemoryMetadata::default(),
144            })),
145            expired_items_sender: expired_tx,
146            expired_items_receiver: expired_rx,
147        }
148    }
149
150    /// Creates a new [Memory] based on the provided config and previous state.
151    pub fn from_previous_state(
152        config: MemoryConfig,
153        prev_state: Box<dyn std::any::Any + Send + Sync>,
154    ) -> Self {
155        if let Ok(prev_memory) = prev_state.downcast::<Memory>() {
156            Self {
157                config,
158                read_handle_factory: prev_memory.read_handle_factory,
159                read_handle: prev_memory.read_handle,
160                write_handle: prev_memory.write_handle,
161                expired_items_sender: prev_memory.expired_items_sender,
162                expired_items_receiver: prev_memory.expired_items_receiver,
163            }
164        } else {
165            Self::new(config)
166        }
167    }
168
169    pub(super) fn get_read_handle(&self) -> &evmap::ReadHandle<String, MemoryEntry> {
170        self.read_handle
171            .get_or(|| self.read_handle_factory.handle())
172    }
173
174    pub(super) fn subscribe_to_expired_items(&self) -> Receiver<Vec<MemoryEntryPair>> {
175        self.expired_items_sender.subscribe()
176    }
177
178    fn handle_value(&self, value: ObjectMap) {
179        let mut writer = self.write_handle.lock().expect("mutex poisoned");
180        let now = Instant::now();
181
182        for (k, value) in value.into_iter() {
183            let new_entry_key = String::from(k);
184            let Ok(v) = serde_json::to_string(&value) else {
185                emit!(MemoryEnrichmentTableInsertFailed {
186                    key: &new_entry_key,
187                    include_key_metric_tag: self.config.internal_metrics.include_key_tag
188                });
189                continue;
190            };
191            let new_entry = MemoryEntry {
192                value: v,
193                update_time: now.into(),
194                ttl: self
195                    .config
196                    .ttl_field
197                    .path
198                    .as_ref()
199                    .and_then(|p| value.get(p))
200                    .and_then(|v| v.as_integer())
201                    .map(|v| v as u64)
202                    .unwrap_or(self.config.ttl),
203            };
204            let new_entry_size = new_entry_key.size_of() + new_entry.size_of();
205            if let Some(max_byte_size) = self.config.max_byte_size
206                && writer
207                    .metadata
208                    .byte_size
209                    .saturating_add(new_entry_size as u64)
210                    > max_byte_size
211            {
212                // Reject new entries
213                emit!(MemoryEnrichmentTableInsertFailed {
214                    key: &new_entry_key,
215                    include_key_metric_tag: self.config.internal_metrics.include_key_tag
216                });
217                continue;
218            }
219            writer.metadata.byte_size = writer
220                .metadata
221                .byte_size
222                .saturating_add(new_entry_size as u64);
223            emit!(MemoryEnrichmentTableInserted {
224                key: &new_entry_key,
225                include_key_metric_tag: self.config.internal_metrics.include_key_tag
226            });
227            writer.write_handle.update(new_entry_key, new_entry);
228        }
229
230        if self.config.flush_interval.is_none() {
231            self.flush(writer);
232        }
233    }
234
235    fn scan_and_mark_for_deletion(&self, writer: &mut MutexGuard<'_, MemoryWriter>) -> bool {
236        let now = Instant::now();
237
238        let mut needs_flush = false;
239        // Since evmap holds 2 separate maps for the data, we are free to directly remove
240        // elements via the writer, while we are iterating the reader
241        // Refresh will happen only after we manually invoke it after iteration
242        if let Some(reader) = self.get_read_handle().read() {
243            for (k, v) in reader.iter() {
244                if let Some(entry) = v.get_one()
245                    && entry.expired(now)
246                {
247                    // Byte size is not reduced at this point, because the actual deletion
248                    // will only happen at refresh time
249                    writer.write_handle.empty(k.clone());
250                    emit!(MemoryEnrichmentTableTtlExpired {
251                        key: k,
252                        include_key_metric_tag: self.config.internal_metrics.include_key_tag
253                    });
254                    needs_flush = true;
255                }
256            }
257        };
258
259        needs_flush
260    }
261
262    fn scan(&self, mut writer: MutexGuard<'_, MemoryWriter>) {
263        let needs_flush = self.scan_and_mark_for_deletion(&mut writer);
264        if needs_flush {
265            self.flush(writer);
266        }
267    }
268
269    fn flush(&self, mut writer: MutexGuard<'_, MemoryWriter>) {
270        // First publish items to be removed, if needed
271        if self
272            .config
273            .source_config
274            .as_ref()
275            .map(|c| c.export_expired_items)
276            .unwrap_or_default()
277        {
278            let pending_removal = writer
279                .write_handle
280                .pending()
281                .iter()
282                // We only use empty operation to remove keys
283                .filter_map(|o| match o {
284                    evmap::Operation::Empty(k) => Some(k),
285                    _ => None,
286                })
287                .filter_map(|key| {
288                    writer.write_handle.get_one(key).map(|v| MemoryEntryPair {
289                        key: key.to_string(),
290                        entry: v.clone(),
291                    })
292                })
293                .collect::<Vec<_>>();
294            if let Err(error) = self.expired_items_sender.send(pending_removal) {
295                error!(
296                    message = "Error exporting expired items from memory enrichment table.",
297                    error = %error,
298                );
299            }
300        }
301
302        writer.write_handle.refresh();
303        if let Some(reader) = self.get_read_handle().read() {
304            let mut byte_size = 0;
305            for (k, v) in reader.iter() {
306                byte_size += k.size_of() + v.get_one().size_of();
307            }
308            writer.metadata.byte_size = byte_size as u64;
309            emit!(MemoryEnrichmentTableFlushed {
310                new_objects_count: reader.len(),
311                new_byte_size: byte_size
312            });
313        }
314    }
315
316    pub(crate) fn as_source(
317        &self,
318        shutdown: ShutdownSignal,
319        out: SourceSender,
320        log_namespace: LogNamespace,
321    ) -> MemorySource {
322        MemorySource {
323            memory: self.clone(),
324            shutdown,
325            out,
326            log_namespace,
327        }
328    }
329}
330
331impl Clone for Memory {
332    fn clone(&self) -> Self {
333        Self {
334            read_handle_factory: self.read_handle_factory.clone(),
335            read_handle: ThreadLocal::new(),
336            write_handle: Arc::clone(&self.write_handle),
337            config: self.config.clone(),
338            expired_items_sender: self.expired_items_sender.clone(),
339            expired_items_receiver: self.expired_items_sender.subscribe(),
340        }
341    }
342}
343
344impl Table for Memory {
345    fn find_table_row<'a>(
346        &self,
347        case: Case,
348        condition: &'a [Condition<'a>],
349        select: Option<&'a [String]>,
350        wildcard: Option<&Value>,
351        index: Option<IndexHandle>,
352    ) -> Result<ObjectMap, Error> {
353        let mut rows = self.find_table_rows(case, condition, select, wildcard, index)?;
354
355        match rows.pop() {
356            Some(row) if rows.is_empty() => Ok(row),
357            Some(_) => Err(Error::MoreThanOneRowFound),
358            None => Err(Error::NoRowsFound),
359        }
360    }
361
362    fn find_table_rows<'a>(
363        &self,
364        _case: Case,
365        condition: &'a [Condition<'a>],
366        _select: Option<&'a [String]>,
367        _wildcard: Option<&Value>,
368        _index: Option<IndexHandle>,
369    ) -> Result<Vec<ObjectMap>, Error> {
370        match condition.first() {
371            Some(_) if condition.len() > 1 => Err(Error::OnlyOneConditionAllowed),
372            Some(Condition::Equals { value, .. }) => {
373                let key = value.to_string_lossy();
374                match self.get_read_handle().get_one(key.as_ref()) {
375                    Some(row) => {
376                        emit!(MemoryEnrichmentTableRead {
377                            key: &key,
378                            include_key_metric_tag: self.config.internal_metrics.include_key_tag
379                        });
380                        row.as_object_map(Instant::now(), &key).map(|r| vec![r])
381                    }
382                    None => {
383                        emit!(MemoryEnrichmentTableReadFailed {
384                            key: &key,
385                            include_key_metric_tag: self.config.internal_metrics.include_key_tag
386                        });
387                        Ok(Default::default())
388                    }
389                }
390            }
391            Some(_) => Err(Error::OnlyEqualityConditionAllowed),
392            None => Err(Error::MissingCondition { kind: "Key" }),
393        }
394    }
395
396    fn add_index(&mut self, _case: Case, fields: &[&str]) -> Result<IndexHandle, Error> {
397        match fields.len() {
398            0 => Err(Error::MissingRequiredField { field: "Key" }),
399            1 => Ok(IndexHandle(0)),
400            _ => Err(Error::OnlyOneFieldAllowed),
401        }
402    }
403
404    /// Returns a list of the field names that are in each index
405    fn index_fields(&self) -> Vec<(Case, Vec<String>)> {
406        Vec::new()
407    }
408
409    /// Doesn't need reload, data is written directly
410    fn needs_reload(&self) -> bool {
411        false
412    }
413
414    fn extract_state(&self) -> Option<Box<dyn std::any::Any + Send + Sync>> {
415        let writer = self.write_handle.lock().expect("mutex poisoned");
416        self.flush(writer);
417        Some(Box::new(self.clone()))
418    }
419}
420
421impl std::fmt::Debug for Memory {
422    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
423        write!(f, "Memory {} row(s)", self.get_read_handle().len())
424    }
425}
426
427#[async_trait]
428impl StreamSink<Event> for Memory {
429    async fn run(mut self: Box<Self>, mut input: BoxStream<'_, Event>) -> Result<(), ()> {
430        let events_sent = register!(EventsSent::from(Output(None)));
431        let bytes_sent = register!(BytesSent::from(Protocol("memory_enrichment_table".into(),)));
432        let mut flush_interval: Pin<Box<dyn Stream<Item = tokio::time::Instant> + Send>> = self
433            .config
434            .flush_interval
435            .map(NonZeroU64::get)
436            .map(Duration::from_secs)
437            .map::<Pin<Box<dyn Stream<Item = tokio::time::Instant> + Send>>, _>(|d| {
438                Box::pin(IntervalStream::new(interval(d)))
439            })
440            .unwrap_or(Box::pin(stream::empty()));
441        let mut scan_interval = IntervalStream::new(interval(Duration::from_secs(
442            self.config.scan_interval.into(),
443        )));
444
445        loop {
446            tokio::select! {
447                event = input.next() => {
448                    let mut event = if let Some(event) = event {
449                        event
450                    } else {
451                        break;
452                    };
453                    let event_byte_size = event.estimated_json_encoded_size_of();
454
455                    let finalizers = event.take_finalizers();
456
457                    // Panic: This sink only accepts Logs, so this should never panic
458                    let log = event.into_log();
459
460                    if let (Value::Object(map), _) = log.into_parts() {
461                        self.handle_value(map)
462                    };
463
464                    finalizers.update_status(EventStatus::Delivered);
465                    events_sent.emit(CountByteSize(1, event_byte_size));
466                    bytes_sent.emit(ByteSize(event_byte_size.get()));
467                }
468
469                Some(_) = flush_interval.next() => {
470                    let writer = self.write_handle.lock().expect("mutex poisoned");
471                    self.flush(writer);
472                }
473
474                Some(_) = scan_interval.next() => {
475                    let writer = self.write_handle.lock().expect("mutex poisoned");
476                    self.scan(writer);
477                }
478            }
479        }
480        Ok(())
481    }
482}
483
484#[cfg(test)]
485mod tests {
486    use std::{num::NonZeroU64, slice::from_ref, time::Duration};
487
488    use futures::{StreamExt, future::ready};
489    use futures_util::stream;
490    use tokio::time;
491
492    use vector_lib::{
493        event::{EventContainer, MetricValue},
494        lookup::lookup_v2::OptionalValuePath,
495        metrics::Controller,
496        sink::VectorSink,
497    };
498
499    use super::*;
500    use crate::{
501        config::EnrichmentTableConfig,
502        enrichment_tables::memory::{
503            config::MemorySourceConfig, internal_events::InternalMetricsConfig,
504        },
505        event::{Event, LogEvent},
506        test_util::components::{
507            SINK_TAGS, SOURCE_TAGS, run_and_assert_sink_compliance,
508            run_and_assert_source_compliance,
509        },
510    };
511
512    fn build_memory_config(modfn: impl Fn(&mut MemoryConfig)) -> MemoryConfig {
513        let mut config = MemoryConfig::default();
514        modfn(&mut config);
515        config
516    }
517
518    #[test]
519    fn finds_row() {
520        let memory = Memory::new(Default::default());
521        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
522
523        let condition = Condition::Equals {
524            field: "key",
525            value: Value::from("test_key"),
526        };
527
528        assert_eq!(
529            Ok(ObjectMap::from([
530                ("key".into(), Value::from("test_key")),
531                ("ttl".into(), Value::from(memory.config.ttl)),
532                ("value".into(), Value::from(5)),
533            ])),
534            memory.find_table_row(Case::Sensitive, &[condition], None, None, None)
535        );
536    }
537
538    #[tokio::test]
539    async fn extract_state_preserves_data() {
540        let memory = Memory::new(Default::default());
541        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
542
543        let condition = Condition::Equals {
544            field: "key",
545            value: Value::from("test_key"),
546        };
547
548        let expected = ObjectMap::from([
549            ("key".into(), Value::from("test_key")),
550            ("ttl".into(), Value::from(memory.config.ttl)),
551            ("value".into(), Value::from(5)),
552        ]);
553        assert_eq!(
554            Ok(expected.clone()),
555            memory.find_table_row(
556                Case::Sensitive,
557                std::slice::from_ref(&condition),
558                None,
559                None,
560                None
561            )
562        );
563
564        // Now build a new table using old state
565        let new_memory = MemoryConfig::default()
566            .build(&Default::default(), memory.extract_state())
567            .await
568            .unwrap();
569        assert_eq!(
570            Ok(expected),
571            new_memory.find_table_row(Case::Sensitive, &[condition], None, None, None)
572        );
573    }
574
575    #[test]
576    fn calculates_ttl() {
577        let ttl = 100;
578        let secs_to_subtract = 10;
579        let memory = Memory::new(build_memory_config(|c| c.ttl = ttl));
580        {
581            let mut handle = memory.write_handle.lock().unwrap();
582            handle.write_handle.update(
583                "test_key".to_string(),
584                MemoryEntry {
585                    value: "5".to_string(),
586                    update_time: (Instant::now() - Duration::from_secs(secs_to_subtract)).into(),
587                    ttl,
588                },
589            );
590            handle.write_handle.refresh();
591        }
592
593        let condition = Condition::Equals {
594            field: "key",
595            value: Value::from("test_key"),
596        };
597
598        assert_eq!(
599            Ok(ObjectMap::from([
600                ("key".into(), Value::from("test_key")),
601                ("ttl".into(), Value::from(ttl - secs_to_subtract)),
602                ("value".into(), Value::from(5)),
603            ])),
604            memory.find_table_row(Case::Sensitive, &[condition], None, None, None)
605        );
606    }
607
608    #[test]
609    fn calculates_ttl_override() {
610        let global_ttl = 100;
611        let ttl_override = 10;
612        let memory = Memory::new(build_memory_config(|c| {
613            c.ttl = global_ttl;
614            c.ttl_field = OptionalValuePath::new("ttl");
615        }));
616        memory.handle_value(ObjectMap::from([
617            (
618                "ttl_override".into(),
619                Value::from(ObjectMap::from([
620                    ("val".into(), Value::from(5)),
621                    ("ttl".into(), Value::from(ttl_override)),
622                ])),
623            ),
624            (
625                "default_ttl".into(),
626                Value::from(ObjectMap::from([("val".into(), Value::from(5))])),
627            ),
628        ]));
629
630        let default_condition = Condition::Equals {
631            field: "key",
632            value: Value::from("default_ttl"),
633        };
634        let override_condition = Condition::Equals {
635            field: "key",
636            value: Value::from("ttl_override"),
637        };
638
639        assert_eq!(
640            Ok(ObjectMap::from([
641                ("key".into(), Value::from("default_ttl")),
642                ("ttl".into(), Value::from(global_ttl)),
643                (
644                    "value".into(),
645                    Value::from(ObjectMap::from([("val".into(), Value::from(5))]))
646                ),
647            ])),
648            memory.find_table_row(Case::Sensitive, &[default_condition], None, None, None)
649        );
650        assert_eq!(
651            Ok(ObjectMap::from([
652                ("key".into(), Value::from("ttl_override")),
653                ("ttl".into(), Value::from(ttl_override)),
654                (
655                    "value".into(),
656                    Value::from(ObjectMap::from([
657                        ("val".into(), Value::from(5)),
658                        ("ttl".into(), Value::from(ttl_override))
659                    ]))
660                ),
661            ])),
662            memory.find_table_row(Case::Sensitive, &[override_condition], None, None, None)
663        );
664    }
665
666    #[test]
667    fn removes_expired_records_on_scan_interval() {
668        let ttl = 100;
669        let memory = Memory::new(build_memory_config(|c| {
670            c.ttl = ttl;
671        }));
672        {
673            let mut handle = memory.write_handle.lock().unwrap();
674            handle.write_handle.update(
675                "test_key".to_string(),
676                MemoryEntry {
677                    value: "5".to_string(),
678                    update_time: (Instant::now() - Duration::from_secs(ttl + 10)).into(),
679                    ttl,
680                },
681            );
682            handle.write_handle.refresh();
683        }
684
685        // Finds the value before scan
686        let condition = Condition::Equals {
687            field: "key",
688            value: Value::from("test_key"),
689        };
690        assert_eq!(
691            Ok(ObjectMap::from([
692                ("key".into(), Value::from("test_key")),
693                ("ttl".into(), Value::from(0)),
694                ("value".into(), Value::from(5)),
695            ])),
696            memory.find_table_row(Case::Sensitive, from_ref(&condition), None, None, None)
697        );
698
699        // Force scan
700        let writer = memory.write_handle.lock().unwrap();
701        memory.scan(writer);
702
703        // The value is not present anymore
704        assert!(
705            memory
706                .find_table_rows(Case::Sensitive, &[condition], None, None, None)
707                .unwrap()
708                .pop()
709                .is_none()
710        );
711    }
712
713    #[test]
714    fn does_not_show_values_before_flush_interval() {
715        let ttl = 100;
716        let memory = Memory::new(build_memory_config(|c| {
717            c.ttl = ttl;
718            c.flush_interval = NonZeroU64::new(10);
719        }));
720        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
721
722        let condition = Condition::Equals {
723            field: "key",
724            value: Value::from("test_key"),
725        };
726
727        assert!(
728            memory
729                .find_table_rows(Case::Sensitive, &[condition], None, None, None)
730                .unwrap()
731                .pop()
732                .is_none()
733        );
734    }
735
736    #[test]
737    fn updates_ttl_on_value_replacement() {
738        let ttl = 100;
739        let memory = Memory::new(build_memory_config(|c| c.ttl = ttl));
740        {
741            let mut handle = memory.write_handle.lock().unwrap();
742            handle.write_handle.update(
743                "test_key".to_string(),
744                MemoryEntry {
745                    value: "5".to_string(),
746                    update_time: (Instant::now() - Duration::from_secs(ttl / 2)).into(),
747                    ttl,
748                },
749            );
750            handle.write_handle.refresh();
751        }
752        let condition = Condition::Equals {
753            field: "key",
754            value: Value::from("test_key"),
755        };
756
757        assert_eq!(
758            Ok(ObjectMap::from([
759                ("key".into(), Value::from("test_key")),
760                ("ttl".into(), Value::from(ttl / 2)),
761                ("value".into(), Value::from(5)),
762            ])),
763            memory.find_table_row(Case::Sensitive, from_ref(&condition), None, None, None)
764        );
765
766        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
767
768        assert_eq!(
769            Ok(ObjectMap::from([
770                ("key".into(), Value::from("test_key")),
771                ("ttl".into(), Value::from(ttl)),
772                ("value".into(), Value::from(5)),
773            ])),
774            memory.find_table_row(Case::Sensitive, &[condition], None, None, None)
775        );
776    }
777
778    #[test]
779    fn ignores_all_values_over_byte_size_limit() {
780        let memory = Memory::new(build_memory_config(|c| {
781            c.max_byte_size = Some(1);
782        }));
783        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
784
785        let condition = Condition::Equals {
786            field: "key",
787            value: Value::from("test_key"),
788        };
789
790        assert!(
791            memory
792                .find_table_rows(Case::Sensitive, &[condition], None, None, None)
793                .unwrap()
794                .pop()
795                .is_none()
796        );
797    }
798
799    #[test]
800    fn ignores_values_when_byte_size_limit_is_reached() {
801        let ttl = 100;
802        let memory = Memory::new(build_memory_config(|c| {
803            c.ttl = ttl;
804            c.max_byte_size = Some(150);
805        }));
806        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
807        memory.handle_value(ObjectMap::from([("rejected_key".into(), Value::from(5))]));
808
809        assert_eq!(
810            Ok(ObjectMap::from([
811                ("key".into(), Value::from("test_key")),
812                ("ttl".into(), Value::from(ttl)),
813                ("value".into(), Value::from(5)),
814            ])),
815            memory.find_table_row(
816                Case::Sensitive,
817                &[Condition::Equals {
818                    field: "key",
819                    value: Value::from("test_key")
820                }],
821                None,
822                None,
823                None
824            )
825        );
826
827        assert!(
828            memory
829                .find_table_rows(
830                    Case::Sensitive,
831                    &[Condition::Equals {
832                        field: "key",
833                        value: Value::from("rejected_key")
834                    }],
835                    None,
836                    None,
837                    None
838                )
839                .unwrap()
840                .pop()
841                .is_none()
842        );
843    }
844
845    #[test]
846    fn missing_key() {
847        let memory = Memory::new(Default::default());
848
849        let condition = Condition::Equals {
850            field: "key",
851            value: Value::from("test_key"),
852        };
853
854        assert!(
855            memory
856                .find_table_rows(Case::Sensitive, &[condition], None, None, None)
857                .unwrap()
858                .pop()
859                .is_none()
860        );
861    }
862
863    #[tokio::test]
864    async fn sink_spec_compliance() {
865        let event = Event::Log(LogEvent::from(ObjectMap::from([(
866            "test_key".into(),
867            Value::from(5),
868        )])));
869
870        let memory = Memory::new(Default::default());
871
872        run_and_assert_sink_compliance(
873            VectorSink::from_event_streamsink(memory),
874            stream::once(ready(event)),
875            &SINK_TAGS,
876        )
877        .await;
878    }
879
880    #[tokio::test]
881    async fn flush_metrics_without_interval() {
882        let event = Event::Log(LogEvent::from(ObjectMap::from([(
883            "test_key".into(),
884            Value::from(5),
885        )])));
886
887        let memory = Memory::new(Default::default());
888
889        run_and_assert_sink_compliance(
890            VectorSink::from_event_streamsink(memory),
891            stream::once(ready(event)),
892            &SINK_TAGS,
893        )
894        .await;
895
896        let metrics = Controller::get().unwrap().capture_metrics();
897        let insertions_counter = metrics
898            .iter()
899            .find(|m| {
900                matches!(m.value(), MetricValue::Counter { .. })
901                    && m.name() == "memory_enrichment_table_insertions_total"
902            })
903            .expect("Insertions metric is missing!");
904        let MetricValue::Counter {
905            value: insertions_count,
906        } = insertions_counter.value()
907        else {
908            unreachable!();
909        };
910        let flushes_counter = metrics
911            .iter()
912            .find(|m| {
913                matches!(m.value(), MetricValue::Counter { .. })
914                    && m.name() == "memory_enrichment_table_flushes_total"
915            })
916            .expect("Flushes metric is missing!");
917        let MetricValue::Counter {
918            value: flushes_count,
919        } = flushes_counter.value()
920        else {
921            unreachable!();
922        };
923        let object_count_gauge = metrics
924            .iter()
925            .find(|m| {
926                matches!(m.value(), MetricValue::Gauge { .. })
927                    && m.name() == "memory_enrichment_table_objects_count"
928            })
929            .expect("Object count metric is missing!");
930        let MetricValue::Gauge {
931            value: object_count,
932        } = object_count_gauge.value()
933        else {
934            unreachable!();
935        };
936        let byte_size_gauge = metrics
937            .iter()
938            .find(|m| {
939                matches!(m.value(), MetricValue::Gauge { .. })
940                    && m.name() == "memory_enrichment_table_byte_size"
941            })
942            .expect("Byte size metric is missing!");
943        assert_eq!(*insertions_count, 1.0);
944        assert_eq!(*flushes_count, 1.0);
945        assert_eq!(*object_count, 1.0);
946        assert!(!byte_size_gauge.is_empty());
947    }
948
949    #[tokio::test]
950    async fn flush_metrics_with_interval() {
951        let event = Event::Log(LogEvent::from(ObjectMap::from([(
952            "test_key".into(),
953            Value::from(5),
954        )])));
955
956        let memory = Memory::new(build_memory_config(|c| {
957            c.flush_interval = NonZeroU64::new(1);
958        }));
959
960        run_and_assert_sink_compliance(
961            VectorSink::from_event_streamsink(memory),
962            stream::iter(vec![event.clone(), event]).flat_map(|e| {
963                stream::once(async move {
964                    tokio::time::sleep(Duration::from_millis(600)).await;
965                    e
966                })
967            }),
968            &SINK_TAGS,
969        )
970        .await;
971
972        let metrics = Controller::get().unwrap().capture_metrics();
973        let insertions_counter = metrics
974            .iter()
975            .find(|m| {
976                matches!(m.value(), MetricValue::Counter { .. })
977                    && m.name() == "memory_enrichment_table_insertions_total"
978            })
979            .expect("Insertions metric is missing!");
980        let MetricValue::Counter {
981            value: insertions_count,
982        } = insertions_counter.value()
983        else {
984            unreachable!();
985        };
986        let flushes_counter = metrics
987            .iter()
988            .find(|m| {
989                matches!(m.value(), MetricValue::Counter { .. })
990                    && m.name() == "memory_enrichment_table_flushes_total"
991            })
992            .expect("Flushes metric is missing!");
993        let MetricValue::Counter {
994            value: flushes_count,
995        } = flushes_counter.value()
996        else {
997            unreachable!();
998        };
999        let object_count_gauge = metrics
1000            .iter()
1001            .find(|m| {
1002                matches!(m.value(), MetricValue::Gauge { .. })
1003                    && m.name() == "memory_enrichment_table_objects_count"
1004            })
1005            .expect("Object count metric is missing!");
1006        let MetricValue::Gauge {
1007            value: object_count,
1008        } = object_count_gauge.value()
1009        else {
1010            unreachable!();
1011        };
1012        let byte_size_gauge = metrics
1013            .iter()
1014            .find(|m| {
1015                matches!(m.value(), MetricValue::Gauge { .. })
1016                    && m.name() == "memory_enrichment_table_byte_size"
1017            })
1018            .expect("Byte size metric is missing!");
1019
1020        assert_eq!(*insertions_count, 2.0);
1021        // One is done right away and the next one after the interval
1022        assert_eq!(*flushes_count, 2.0);
1023        assert_eq!(*object_count, 1.0);
1024        assert!(!byte_size_gauge.is_empty());
1025    }
1026
1027    #[tokio::test]
1028    async fn flush_metrics_with_key() {
1029        let event = Event::Log(LogEvent::from(ObjectMap::from([(
1030            "test_key".into(),
1031            Value::from(5),
1032        )])));
1033
1034        let memory = Memory::new(build_memory_config(|c| {
1035            c.internal_metrics = InternalMetricsConfig {
1036                include_key_tag: true,
1037            };
1038        }));
1039
1040        run_and_assert_sink_compliance(
1041            VectorSink::from_event_streamsink(memory),
1042            stream::once(ready(event)),
1043            &SINK_TAGS,
1044        )
1045        .await;
1046
1047        let metrics = Controller::get().unwrap().capture_metrics();
1048        let insertions_counter = metrics
1049            .iter()
1050            .find(|m| {
1051                matches!(m.value(), MetricValue::Counter { .. })
1052                    && m.name() == "memory_enrichment_table_insertions_total"
1053            })
1054            .expect("Insertions metric is missing!");
1055
1056        assert!(insertions_counter.tag_matches("key", "test_key"));
1057    }
1058
1059    #[tokio::test]
1060    async fn flush_metrics_without_key() {
1061        let event = Event::Log(LogEvent::from(ObjectMap::from([(
1062            "test_key".into(),
1063            Value::from(5),
1064        )])));
1065
1066        let memory = Memory::new(Default::default());
1067
1068        run_and_assert_sink_compliance(
1069            VectorSink::from_event_streamsink(memory),
1070            stream::once(ready(event)),
1071            &SINK_TAGS,
1072        )
1073        .await;
1074
1075        let metrics = Controller::get().unwrap().capture_metrics();
1076        let insertions_counter = metrics
1077            .iter()
1078            .find(|m| {
1079                matches!(m.value(), MetricValue::Counter { .. })
1080                    && m.name() == "memory_enrichment_table_insertions_total"
1081            })
1082            .expect("Insertions metric is missing!");
1083
1084        assert!(insertions_counter.tag_value("key").is_none());
1085    }
1086
1087    #[tokio::test]
1088    async fn source_spec_compliance() {
1089        let mut memory_config = MemoryConfig::default();
1090        memory_config.source_config = Some(MemorySourceConfig {
1091            export_interval: Some(NonZeroU64::try_from(1).unwrap()),
1092            export_batch_size: None,
1093            remove_after_export: false,
1094            export_expired_items: false,
1095            source_key: "test".to_string(),
1096        });
1097        let memory = memory_config.get_or_build_memory(None).await;
1098        memory.handle_value(ObjectMap::from([("test_key".into(), Value::from(5))]));
1099
1100        let mut events: Vec<Event> = run_and_assert_source_compliance(
1101            memory_config,
1102            time::Duration::from_secs(5),
1103            &SOURCE_TAGS,
1104        )
1105        .await;
1106
1107        assert!(!events.is_empty());
1108        let event = events.remove(0);
1109        let log = event.as_log();
1110
1111        assert!(!log.value().is_empty());
1112    }
1113}
vector/enrichment_tables/memory/table.rs

vector/enrichment_tables/memory/
table.rs
