From 8572e2a6fdf96f908487f4b84ab034bc7435f568 Mon Sep 17 00:00:00 2001 From: Emilia Hane Date: Mon, 15 Jan 2024 11:44:35 +0100 Subject: [PATCH] improve usability of schnellru::LruMap and LruCache (LinkedHashSet cache) (#6016) --- Cargo.lock | 2 + Cargo.toml | 1 + crates/net/network/Cargo.toml | 2 + crates/net/network/src/cache.rs | 117 ++++++++++++++++++++++++++++++-- 4 files changed, 116 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d475e5d6fd..62aebd6b02 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6264,6 +6264,7 @@ dependencies = [ "async-trait", "auto_impl", "criterion", + "derive_more", "enr", "ethers-core", "ethers-middleware", @@ -6295,6 +6296,7 @@ dependencies = [ "reth-tokio-util", "reth-tracing", "reth-transaction-pool", + "schnellru", "secp256k1 0.27.0", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 3cfe7e49be..6153220440 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -188,6 +188,7 @@ aquamarine = "0.5" bytes = "1.5" bitflags = "2.4" clap = "4" +derive_more = "0.99.17" eyre = "0.6" tracing = "0.1.0" tracing-appender = "0.2" diff --git a/crates/net/network/Cargo.toml b/crates/net/network/Cargo.toml index 0e363eebf2..3812c7350e 100644 --- a/crates/net/network/Cargo.toml +++ b/crates/net/network/Cargo.toml @@ -63,6 +63,8 @@ linked_hash_set = "0.1" linked-hash-map = "0.5.6" rand.workspace = true secp256k1 = { workspace = true, features = ["global-context", "rand-std", "recovery"] } +derive_more.workspace = true +schnellru.workspace = true enr = { workspace = true, features = ["rust-secp256k1"], optional = true } tempfile = { workspace = true, optional = true } diff --git a/crates/net/network/src/cache.rs b/crates/net/network/src/cache.rs index 45860024ec..520c125168 100644 --- a/crates/net/network/src/cache.rs +++ b/crates/net/network/src/cache.rs @@ -1,5 +1,13 @@ +use core::hash::BuildHasher; +use derive_more::{Deref, DerefMut}; use linked_hash_set::LinkedHashSet; -use std::{borrow::Borrow, hash::Hash, num::NonZeroUsize}; +use schnellru::{self, ByLength, Limiter, RandomState, Unlimited}; +use std::{ + borrow::Borrow, + fmt::{self, Write}, + hash::Hash, + num::NonZeroUsize, +}; /// A minimal LRU cache based on a `LinkedHashSet` with limited capacity. /// @@ -12,7 +20,7 @@ pub struct LruCache { } impl LruCache { - /// Creates a new `LruCache` using the given limit + /// Creates a new [`LruCache`] using the given limit pub fn new(limit: NonZeroUsize) -> Self { Self { inner: LinkedHashSet::new(), limit } } @@ -26,14 +34,21 @@ impl LruCache { /// If the set did not have this value present, true is returned. /// If the set did have this value present, false is returned. pub fn insert(&mut self, entry: T) -> bool { + let (new_entry, _evicted_val) = self.insert_and_get_evicted(entry); + new_entry + } + + /// Same as [`Self::insert`] but returns a tuple, where the second index is the evicted value, + /// if one was evicted. + pub fn insert_and_get_evicted(&mut self, entry: T) -> (bool, Option) { if self.inner.insert(entry) { if self.limit.get() == self.inner.len() { // remove the oldest element in the set - _ = self.remove_lru(); + return (true, self.remove_lru()) } - return true + return (true, None) } - false + (false, None) } /// Remove the least recently used entry and return it. @@ -44,6 +59,12 @@ impl LruCache { self.inner.pop_front() } + #[allow(dead_code)] + /// Expels the given value. Returns true if the value existed. + pub fn remove(&mut self, value: &T) -> bool { + self.inner.remove(value) + } + /// Returns `true` if the set contains a value. pub fn contains(&self, value: &Q) -> bool where @@ -57,6 +78,20 @@ impl LruCache { pub fn iter(&self) -> impl Iterator + '_ { self.inner.iter() } + + /// Returns number of elements currently in cache. + #[cfg(test)] + #[allow(dead_code)] + pub fn len(&self) -> usize { + self.inner.len() + } + + /// Returns `true` if there are currently no elements in the cache. + #[cfg(test)] + #[allow(dead_code)] + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } } impl Extend for LruCache @@ -70,8 +105,56 @@ where } } +/// Wrapper of [`schnellru::LruMap`] that implements [`fmt::Debug`]. +#[derive(Deref, DerefMut)] +pub struct LruMap(schnellru::LruMap) +where + K: Hash + PartialEq, + L: Limiter, + S: BuildHasher; + +impl fmt::Debug for LruMap +where + K: Hash + PartialEq + fmt::Display, + V: fmt::Debug, + L: Limiter, + S: BuildHasher, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut debug_struct = f.debug_struct("LruMap"); + for (k, v) in self.0.iter() { + let mut key_str = String::new(); + write!(&mut key_str, "{k}")?; + debug_struct.field(&key_str, &v); + } + debug_struct.finish() + } +} + +impl LruMap +where + K: Hash + PartialEq, +{ + /// Returns a new cache with default limiter and hash builder. + #[allow(dead_code)] + pub fn new(max_length: u32) -> Self { + LruMap(schnellru::LruMap::new(ByLength::new(max_length))) + } +} + +impl LruMap +where + K: Hash + PartialEq, +{ + /// Returns a new cache with [`Unlimited`] limiter and default hash builder. + #[allow(dead_code)] + pub fn new_unlimited() -> Self { + LruMap(schnellru::LruMap::new(Unlimited)) + } +} + #[cfg(test)] -mod tests { +mod test { use super::*; #[test] @@ -115,4 +198,26 @@ mod tests { assert!(cache.contains(e)); } } + + #[test] + #[allow(dead_code)] + fn test_debug_impl_lru_map() { + use derive_more::Display; + + #[derive(Debug, Hash, PartialEq, Eq, Display)] + struct Key(i8); + + #[derive(Debug)] + struct Value(i8); + + let mut cache = LruMap::new(2); + let key_1 = Key(1); + let value_1 = Value(11); + cache.insert(key_1, value_1); + let key_2 = Key(2); + let value_2 = Value(22); + cache.insert(key_2, value_2); + + assert_eq!("LruMap { 2: Value(22), 1: Value(11) }", format!("{cache:?}")) + } }