evictor.rs 5.35 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

16
17
use std::cmp::{Eq, Ordering};
use std::collections::{BTreeSet, HashMap};
18
use std::hash::Hash;
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

/// A wrapper for (T, counter) that implements Ord based only on counter
#[derive(Debug, Clone, Eq, PartialEq)]
struct PriorityItem<T> {
    item: T,
    counter: i64,
}

impl<T: Eq> Ord for PriorityItem<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.counter.cmp(&other.counter)
    }
}

impl<T: Eq> PartialOrd for PriorityItem<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}
38
39

/// An LRU evictor that maintains objects and evicts them based on their
40
/// priority counter. Lower counter values are evicted first.
41
42
#[derive(Debug)]
pub struct LRUEvictor<T: Clone + Eq + Hash> {
43
44
45
46
    free_table: HashMap<T, i64>,
    priority_queue: BTreeSet<PriorityItem<T>>,
    positive_counter: i64,
    negative_counter: i64,
47
48
49
50
51
52
}

impl<T: Clone + Eq + Hash> Default for LRUEvictor<T> {
    fn default() -> Self {
        Self {
            free_table: HashMap::new(),
53
54
55
            priority_queue: BTreeSet::new(),
            positive_counter: 0,
            negative_counter: 0,
56
57
58
59
60
        }
    }
}

impl<T: Clone + Eq + Hash> LRUEvictor<T> {
61
62
    pub fn new(_cleanup_threshold: usize) -> Self {
        Self::default()
63
64
    }

65
66
    pub fn keys(&self) -> std::collections::hash_map::Keys<'_, T, i64> {
        self.free_table.keys()
67
68
    }

69
70
71
72
73
74
    fn update(&mut self, object: T, counter: i64) {
        self.free_table.insert(object.clone(), counter);
        self.priority_queue.insert(PriorityItem {
            item: object,
            counter,
        });
75
76
77
    }

    pub fn insert(&mut self, object: T) {
78
79
80
81
82
83
84
        // Remove old entry if it exists
        if let Some(&old_counter) = self.free_table.get(&object) {
            self.priority_queue.remove(&PriorityItem {
                item: object.clone(),
                counter: old_counter,
            });
        }
85

86
87
88
89
90
        // Increment positive counter and insert
        self.positive_counter += 1;
        let counter = self.positive_counter;

        self.update(object, counter);
91
92
    }

93
94
95
96
97
98
99
100
    /// Push an object to the front with negative counter (highest priority for eviction)
    pub fn push_front(&mut self, object: T) {
        // Remove old entry if it exists
        if let Some(&old_counter) = self.free_table.get(&object) {
            self.priority_queue.remove(&PriorityItem {
                item: object.clone(),
                counter: old_counter,
            });
101
102
        }

103
104
105
        // Decrement negative counter and insert
        self.negative_counter -= 1;
        let counter = self.negative_counter;
106

107
108
        self.update(object, counter);
    }
109

110
111
    pub fn contains(&self, object: &T) -> bool {
        self.free_table.contains_key(object)
112
113
    }

114
115
116
117
118
119
120
    /// Evict an object based on LRU policy (lowest counter value)
    /// Returns the evicted object or None if no objects are available
    pub fn evict(&mut self) -> Option<T> {
        self.priority_queue.pop_first().map(|item| {
            self.free_table.remove(&item.item);
            item.item
        })
121
122
123
    }

    pub fn remove(&mut self, object: &T) -> bool {
124
125
126
127
128
129
130
131
132
133
        let Some(&counter) = self.free_table.get(object) else {
            return false;
        };

        self.free_table.remove(object);
        self.priority_queue.remove(&PriorityItem {
            item: object.clone(),
            counter,
        });
        true
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
    }

    pub fn len(&self) -> usize {
        self.free_table.len()
    }

    pub fn is_empty(&self) -> bool {
        self.free_table.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

149
150
151
152
    #[test]
    fn test_lru_evictor_eviction_order() {
        // Create a new LRUEvictor
        let mut evictor = LRUEvictor::<i32>::new(1); // threshold value doesn't matter anymore
153

154
        // Add items in the specified order
155
156
157
158
159
        evictor.insert(4);
        evictor.insert(3);
        evictor.insert(2);
        evictor.insert(1);
        evictor.insert(5);
160
161
162
163
        evictor.insert(1); // Updates counter for 1
        evictor.insert(4); // Updates counter for 4
        evictor.insert(2); // Updates counter for 2
        evictor.push_front(4);
164
165

        // Verify the eviction order
166
167
        let evicted = evictor.evict().unwrap();
        assert_eq!(evicted, 4);
168
169
170
171
172
173
174
175
176
177
178
179
        let evicted = evictor.evict().unwrap();
        assert_eq!(evicted, 3);
        let evicted = evictor.evict().unwrap();
        assert_eq!(evicted, 5);
        let evicted = evictor.evict().unwrap();
        assert_eq!(evicted, 1);
        let evicted = evictor.evict().unwrap();
        assert_eq!(evicted, 2);
        let evicted = evictor.evict();
        assert_eq!(evicted, None);
        assert_eq!(evictor.len(), 0);
    }
180
181

    // ... existing test_push_front test ...
182
}