From 70eca028cdd70c23f614e253e657e3d3129e262d Mon Sep 17 00:00:00 2001 From: Gaius Date: Fri, 14 Mar 2025 15:49:56 +0800 Subject: [PATCH] refactor(dragonfly-client-storage): optimize lru_cache for storage (#1040) Signed-off-by: Gaius --- .../src/cache/lru_cache.rs | 148 ++++++++++-------- dragonfly-client-storage/src/cache/mod.rs | 1 - 2 files changed, 84 insertions(+), 65 deletions(-) diff --git a/dragonfly-client-storage/src/cache/lru_cache.rs b/dragonfly-client-storage/src/cache/lru_cache.rs index 77a14a87..414c14e3 100644 --- a/dragonfly-client-storage/src/cache/lru_cache.rs +++ b/dragonfly-client-storage/src/cache/lru_cache.rs @@ -1,11 +1,28 @@ -use std::{borrow::Borrow, collections::HashMap, hash::Hash, hash::Hasher, num::NonZeroUsize, ptr}; +/* + * Copyright 2025 The Dragonfly Authors + * + * 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. + */ + +use std::{borrow::Borrow, collections::HashMap, hash::Hash, hash::Hasher}; -#[derive(Debug, Clone, Copy)] /// KeyRef is a reference to the key. +#[derive(Debug, Clone, Copy)] struct KeyRef { k: *const K, } +/// KeyRef implements Hash for KeyRef. impl Hash for KeyRef { fn hash(&self, state: &mut H) { unsafe { @@ -15,6 +32,7 @@ impl Hash for KeyRef { } } +/// KeyRef implements PartialEq for KeyRef. impl PartialEq for KeyRef { fn eq(&self, other: &Self) -> bool { unsafe { @@ -25,12 +43,14 @@ impl PartialEq for KeyRef { } } +/// KeyRef implements Eq for KeyRef. impl Eq for KeyRef {} -#[repr(transparent)] /// KeyWrapper is a wrapper for the key. +#[repr(transparent)] struct KeyWrapper(K); +/// KeyWrapper implements reference conversion. impl KeyWrapper { /// from_ref creates a new KeyWrapper from a reference to the key. fn from_ref(key: &K) -> &Self { @@ -38,12 +58,14 @@ impl KeyWrapper { } } +/// KeyWrapper implements Hash for KeyWrapper. impl Hash for KeyWrapper { fn hash(&self, state: &mut H) { self.0.hash(state) } } +/// KeyWrapper implements PartialEq for KeyWrapper. impl PartialEq for KeyWrapper { #![allow(unknown_lints)] #[allow(clippy::unconditional_recursion)] @@ -52,13 +74,16 @@ impl PartialEq for KeyWrapper { } } +/// KeyWrapper implements Eq for KeyWrapper. impl Eq for KeyWrapper {} +/// KeyWrapper implements Borrow for KeyWrapper. impl Borrow> for KeyRef where K: Borrow, Q: ?Sized, { + /// borrow borrows the key. fn borrow(&self) -> &KeyWrapper { unsafe { let key = &*self.k; @@ -67,6 +92,7 @@ where } } +/// Entry is a cache entry. struct Entry { key: K, value: V, @@ -74,7 +100,9 @@ struct Entry { next: Option<*mut Entry>, } +/// Entry implements Drop for Entry. impl Entry { + /// new creates a new Entry. fn new(key: K, value: V) -> Self { Self { key, @@ -85,25 +113,28 @@ impl Entry { } } +/// LruCache is a least recently used cache. pub struct LruCache { - cap: NonZeroUsize, + capacity: usize, map: HashMap, Box>>, head: Option<*mut Entry>, tail: Option<*mut Entry>, _marker: std::marker::PhantomData, } +/// LruCache implements LruCache. impl LruCache { /// new creates a new LruCache. - pub fn new(cap: NonZeroUsize) -> Self { + pub fn new(capacity: usize) -> Self { Self { - cap, - map: HashMap::with_capacity(cap.get()), + capacity, + map: HashMap::with_capacity(capacity), head: None, tail: None, _marker: std::marker::PhantomData, } } + /// get gets the value of the key. pub fn get<'a, Q>(&'a mut self, k: &Q) -> Option<&'a V> where @@ -113,11 +144,8 @@ impl LruCache { if let Some(entry) = self.map.get_mut(KeyWrapper::from_ref(k)) { let entry_ptr: *mut Entry = &mut **entry; - unsafe { - self.detach(entry_ptr); - self.attach(entry_ptr); - } - + self.detach(entry_ptr); + self.attach(entry_ptr); Some(&unsafe { &*entry_ptr }.value) } else { None @@ -125,35 +153,24 @@ impl LruCache { } /// put puts the key and value into the cache. - pub fn put(&mut self, key: K, value: V) -> Option { - let key_ref = KeyRef { k: &key }; + pub fn put(&mut self, key: K, mut value: V) -> Option { + if let Some(existing_entry) = self.map.get_mut(&KeyRef { k: &key }) { + let entry = existing_entry.as_mut(); + std::mem::swap(&mut entry.value, &mut value); - if let Some(existing_entry) = self.map.get_mut(&key_ref) { - let entry_ptr = existing_entry.as_mut() as *mut Entry; - - let old_value = unsafe { - let old_value = ptr::read(&(*entry_ptr).value); - ptr::write(&mut (*entry_ptr).value, value); - old_value - }; - - unsafe { - self.detach(entry_ptr); - self.attach(entry_ptr); - } - - return Some(old_value); + let entry_ptr: *mut Entry = entry; + self.detach(entry_ptr); + self.attach(entry_ptr); + return Some(value); } let mut evicted_value = None; - if self.map.len() >= self.cap.get() { + if self.map.len() >= self.capacity { if let Some(tail) = self.tail { + self.detach(tail); + unsafe { - let tail_key_ref = KeyRef { k: &(*tail).key }; - - self.detach(tail); - - if let Some(entry) = self.map.remove(&tail_key_ref) { + if let Some(entry) = self.map.remove(&KeyRef { k: &(*tail).key }) { evicted_value = Some(entry.value); } } @@ -161,7 +178,7 @@ impl LruCache { } let new_entry = Box::new(Entry::new(key, value)); - let key_ptr = &new_entry.key as *const K; + let key_ptr: *const K = &new_entry.key; let entry_ptr = Box::into_raw(new_entry); unsafe { @@ -174,30 +191,35 @@ impl LruCache { } /// detach detaches the entry from the cache. - unsafe fn detach(&mut self, entry: *mut Entry) { - let prev = (*entry).prev; - let next = (*entry).next; + fn detach(&mut self, entry: *mut Entry) { + unsafe { + let prev = (*entry).prev; + let next = (*entry).next; - match prev { - Some(prev) => (*prev).next = next, - None => self.head = next, + match prev { + Some(prev) => (*prev).next = next, + None => self.head = next, + } + + match next { + Some(next) => (*next).prev = prev, + None => self.tail = prev, + } + + (*entry).prev = None; + (*entry).next = None; } - - match next { - Some(next) => (*next).prev = prev, - None => self.tail = prev, - } - - (*entry).prev = None; - (*entry).next = None; } /// attach attaches the entry to the cache. - unsafe fn attach(&mut self, entry: *mut Entry) { + fn attach(&mut self, entry: *mut Entry) { match self.head { - Some(old_head) => { - (*entry).next = Some(old_head); - (*old_head).prev = Some(entry); + Some(head) => { + unsafe { + (*entry).next = Some(head); + (*head).prev = Some(entry); + } + self.head = Some(entry); } None => { @@ -234,13 +256,11 @@ impl LruCache { } let tail = self.tail?; + self.detach(tail); unsafe { - self.detach(tail); - - let tail_key_ref = KeyRef { k: &(*tail).key }; self.map - .remove(&tail_key_ref) + .remove(&KeyRef { k: &(*tail).key }) .map(|entry| (entry.key, entry.value)) } } @@ -269,7 +289,7 @@ mod tests { #[test] fn test_put() { - let mut cache = LruCache::new(NonZeroUsize::new(3).unwrap()); + let mut cache = LruCache::new(3); assert_eq!(cache.put("key1".to_string(), "value1".to_string()), None); assert_eq!(cache.put("key2".to_string(), "value2".to_string()), None); @@ -296,7 +316,7 @@ mod tests { #[test] fn test_get() { - let mut cache = LruCache::new(NonZeroUsize::new(3).unwrap()); + let mut cache = LruCache::new(3); cache.put("key1".to_string(), "value1".to_string()); cache.put("key2".to_string(), "value2".to_string()); @@ -318,7 +338,7 @@ mod tests { #[test] fn test_peek() { - let mut cache = LruCache::new(NonZeroUsize::new(3).unwrap()); + let mut cache = LruCache::new(3); cache.put("key1".to_string(), "value1".to_string()); cache.put("key2".to_string(), "value2".to_string()); @@ -337,7 +357,7 @@ mod tests { assert_eq!(cache.peek("key3"), Some(&"value3".to_string())); assert_eq!(cache.peek("key4"), Some(&"value4".to_string())); - let mut cache = LruCache::new(NonZeroUsize::new(2).unwrap()); + let mut cache = LruCache::new(2); cache.put("key1".to_string(), "value1".to_string()); cache.put("key2".to_string(), "value2".to_string()); @@ -345,7 +365,7 @@ mod tests { cache.put("key3".to_string(), "value3".to_string()); assert_eq!(cache.peek("key1"), None); - let mut cache = LruCache::new(NonZeroUsize::new(2).unwrap()); + let mut cache = LruCache::new(2); cache.put("key1".to_string(), "value1".to_string()); cache.put("key2".to_string(), "value2".to_string()); @@ -357,7 +377,7 @@ mod tests { #[test] fn test_contains() { - let mut cache = LruCache::new(NonZeroUsize::new(5).unwrap()); + let mut cache = LruCache::new(5); let test_cases = vec![ ("piece_1", Bytes::from("data 1"), false), diff --git a/dragonfly-client-storage/src/cache/mod.rs b/dragonfly-client-storage/src/cache/mod.rs index 2c9ca894..aff4f975 100644 --- a/dragonfly-client-storage/src/cache/mod.rs +++ b/dragonfly-client-storage/src/cache/mod.rs @@ -121,7 +121,6 @@ impl Cache { /// new creates a new cache with the specified capacity. pub fn new(capacity: usize, tasks_capacity: usize) -> Result { let capacity = NonZeroUsize::new(capacity).ok_or(Error::InvalidParameter)?; - let tasks_capacity = NonZeroUsize::new(tasks_capacity).ok_or(Error::InvalidParameter)?; let tasks = Arc::new(RwLock::new(LruCache::new(tasks_capacity))); Ok(Cache {