opentelemetry_spanprocessor_any/sdk/trace/
evicted_hash_map.rs

1//! # Evicted Map
2
3use crate::{Key, KeyValue, Value};
4#[cfg(feature = "serialize")]
5use serde::{Deserialize, Serialize};
6use std::collections::hash_map::Entry;
7use std::collections::{HashMap, LinkedList};
8
9/// A hash map with a capped number of attributes that retains the most
10/// recently set entries.
11#[cfg_attr(feature = "serialize", derive(Deserialize, Serialize))]
12#[derive(Clone, Debug, PartialEq)]
13pub struct EvictedHashMap {
14    map: HashMap<Key, Value>,
15    evict_list: LinkedList<Key>,
16    max_len: u32,
17    dropped_count: u32,
18}
19
20impl EvictedHashMap {
21    /// Create a new `EvictedHashMap` with a given max length and capacity.
22    pub fn new(max_len: u32, capacity: usize) -> Self {
23        EvictedHashMap {
24            map: HashMap::with_capacity(capacity),
25            evict_list: LinkedList::new(),
26            max_len,
27            dropped_count: 0,
28        }
29    }
30
31    /// Inserts a key-value pair into the map.
32    pub fn insert(&mut self, item: KeyValue) {
33        let KeyValue { key, value } = item;
34        let mut already_exists = false;
35        // Check for existing item
36        match self.map.entry(key.clone()) {
37            Entry::Occupied(mut occupied) => {
38                occupied.insert(value);
39                already_exists = true;
40            }
41            Entry::Vacant(entry) => {
42                entry.insert(value);
43            }
44        }
45
46        if already_exists {
47            self.move_key_to_front(key);
48        } else {
49            // Add new item
50            self.evict_list.push_front(key);
51        }
52
53        // Verify size not exceeded
54        if self.evict_list.len() as u32 > self.max_len {
55            self.remove_oldest();
56            self.dropped_count += 1;
57        }
58    }
59
60    /// Returns the number of elements in the map.
61    pub fn len(&self) -> usize {
62        self.map.len()
63    }
64
65    /// Returns `true` if the map is empty.
66    pub fn is_empty(&self) -> bool {
67        self.map.is_empty()
68    }
69
70    /// Returns the dropped attribute count
71    pub fn dropped_count(&self) -> u32 {
72        self.dropped_count
73    }
74
75    /// Returns a front-to-back iterator.
76    pub fn iter(&self) -> Iter<'_> {
77        Iter(self.map.iter())
78    }
79
80    /// Returns a reference to the value corresponding to the key if it exists
81    pub fn get(&self, key: &Key) -> Option<&Value> {
82        self.map.get(key)
83    }
84
85    fn move_key_to_front(&mut self, key: Key) {
86        if self.evict_list.is_empty() {
87            // If empty, push front
88            self.evict_list.push_front(key);
89        } else if self.evict_list.front() == Some(&key) {
90            // Already the front, ignore
91        } else {
92            // Else split linked lists around key and combine
93            let key_idx = self
94                .evict_list
95                .iter()
96                .position(|k| k == &key)
97                .expect("key must exist in evicted hash map, this is a bug");
98            let mut tail = self.evict_list.split_off(key_idx);
99            let item = tail.pop_front().unwrap();
100            self.evict_list.push_front(item);
101            self.evict_list.append(&mut tail);
102        }
103    }
104
105    fn remove_oldest(&mut self) {
106        if let Some(oldest_item) = self.evict_list.pop_back() {
107            self.map.remove(&oldest_item);
108        }
109    }
110}
111
112/// An owned iterator over the entries of a `EvictedHashMap`.
113#[derive(Debug)]
114pub struct IntoIter(std::collections::hash_map::IntoIter<Key, Value>);
115
116impl Iterator for IntoIter {
117    type Item = (Key, Value);
118
119    fn next(&mut self) -> Option<Self::Item> {
120        self.0.next()
121    }
122}
123
124impl IntoIterator for EvictedHashMap {
125    type Item = (Key, Value);
126    type IntoIter = IntoIter;
127
128    fn into_iter(self) -> Self::IntoIter {
129        IntoIter(self.map.into_iter())
130    }
131}
132
133impl<'a> IntoIterator for &'a EvictedHashMap {
134    type Item = (&'a Key, &'a Value);
135    type IntoIter = Iter<'a>;
136
137    fn into_iter(self) -> Self::IntoIter {
138        Iter(self.map.iter())
139    }
140}
141
142/// An iterator over the entries of an `EvictedHashMap`.
143#[derive(Debug)]
144pub struct Iter<'a>(std::collections::hash_map::Iter<'a, Key, Value>);
145
146impl<'a> Iterator for Iter<'a> {
147    type Item = (&'a Key, &'a Value);
148
149    fn next(&mut self) -> Option<Self::Item> {
150        self.0.next()
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use std::collections::HashSet;
158
159    #[test]
160    fn insert_over_capacity_test() {
161        let max_len = 10;
162        let mut map = EvictedHashMap::new(max_len, max_len as usize);
163
164        for i in 0..=max_len {
165            map.insert(Key::new(i.to_string()).bool(true))
166        }
167
168        assert_eq!(map.dropped_count, 1);
169        assert_eq!(map.len(), max_len as usize);
170        assert_eq!(
171            map.map.keys().cloned().collect::<HashSet<_>>(),
172            (1..=max_len)
173                .map(|i| Key::new(i.to_string()))
174                .collect::<HashSet<_>>()
175        );
176    }
177}
opentelemetry_spanprocessor_any/sdk/trace/evicted_hash_map.rs

opentelemetry_spanprocessor_any/sdk/trace/
evicted_hash_map.rs
