Lru在Rust中的實現, 原始碼解析

问蒙服务框架發表於2024-06-07
Rust原始碼

LRU（Least Recently Used）是一種常用的頁面置換演算法，其核心思想是選擇最近最久未使用的頁面予以淘汰。

LRU演算法原理

  • 基本思想：LRU演算法基於一個假設，即如果一個資料在最近一段時間沒有被訪問到，那麼在將來它被訪問的可能性也很低。因此，當快取空間不足時，演算法會選擇最久未使用的資料進行淘汰。

  • 實現方式：LRU演算法通常透過維護一個佇列或連結串列來實現。當訪問一個頁面時，如果該頁面已經在佇列中，則將其移動到佇列的頭部（最近使用）；如果該頁面不在佇列中，則將其新增到佇列的頭部，並檢查佇列長度是否超過預設的閾值。如果佇列長度超過閾值，則淘汰佇列尾部的頁面（最久未使用）。

LRU演算法的優缺點

  • 優點：
    • LRU演算法能夠利用時間區域性性原理，保留最近使用過的頁面，提高快取命中率。
    • 演算法簡單，易於實現。
  • 缺點：
    • 需要維護一個佇列或陣列，佔用額外的空間。
    • 當頁面訪問模式具有迴圈週期時，LRU演算法可能會淘汰掉正在使用的頁面，導致快取命中率下降。
    • 對於隨機訪問的頁面輸入序列，LRU演算法的表現可能不如其他演算法。

結構設計

在Lru的結構中，我們要避免key或者val的複製。
因為key此時需要在雙向列表中儲存也需要在HashMap中儲存，所以我們要以下方案：

  • Rc<K>引用計數

透過引用計數來控制生命週期
優點：不用處理不安全的程式碼
缺點：因為Val可能在遍歷中被更改，所以不能儲存在雙向列表裡，取得值的時候需要進行一次Hash

  • *mut K 裸指標

透過unsafe編碼來實現
優點：在雙向列表及HashMap中均儲存一份數值，遍歷或者根據key取值均只需一次操作
缺點：需要引入ptr，即用指標的方式來進行生命週期管理

節點設計

此時我們用的是裸指標的方式，讓我們先來定義節點資料，資料將儲存在該節點裡面，key及val的生命週期隨節點管理，在刪除的時候需同時在列表及在HashMap中進行刪除

/// Lru節點資料
struct LruEntry<K, V> {
    /// 頭部節點及尾部結點均未初始化值
    pub key: mem::MaybeUninit<K>,
    /// 頭部節點及尾部結點均未初始化值
    pub val: mem::MaybeUninit<V>,
    pub prev: *mut LruEntry<K, V>,
    pub next: *mut LruEntry<K, V>,
}

類設計

接下來需要設計LruCache結構，將由一個map儲存資料結構，一個雙向連結串列儲存訪問優先順序，cap表示快取的容量。

pub struct LruCache<K, V, S> {
    /// 儲存資料結構
    map: HashMap<KeyRef<K>, NonNull<LruEntry<K, V>>, S>,
    /// 快取的總容量
    cap: usize,
    /// 雙向列表的頭
    head: *mut LruEntry<K, V>,
    /// 雙向列表的尾
    tail: *mut LruEntry<K, V>,
}

其中KeyRef僅僅只是表示裸指標的一層包裝，方便重新實現Hash，Eq等trait

#[derive(Clone)]
struct KeyRef<K> {
    pub k: *const K,
}

首先初始化物件，初始化map及空的雙向連結串列：

impl<K, V, S> LruCache<K, V, S> {
    /// 提供hash函式
    pub fn with_hasher(cap: usize, hash_builder: S) -> LruCache<K, V, S> {
        let cap = cap.max(1);
        let map = HashMap::with_capacity_and_hasher(cap, hash_builder);
        let head = Box::into_raw(Box::new(LruEntry::new_empty()));
        let tail = Box::into_raw(Box::new(LruEntry::new_empty()));
        unsafe {
            (*head).next = tail;
            (*tail).prev = head;
        }
        Self {
            map,
            cap,
            head,
            tail,
        }
    }
}

元素插入

插入物件，分已在快取內和不在快取內：

pub fn capture_insert(&mut self, k: K, mut v: V) -> Option<(K, V)> {
    let key = KeyRef::new(&k);
    match self.map.get_mut(&key) {
        Some(entry) => {
            let entry_ptr = entry.as_ptr();
            unsafe {
                mem::swap(&mut *(*entry_ptr).val.as_mut_ptr(), &mut v);
            }
            self.detach(entry_ptr);
            self.attach(entry_ptr);

            Some((k, v))
        }
        None => {
            let (_, entry) = self.replace_or_create_node(k, v);
            let entry_ptr = entry.as_ptr();
            self.attach(entry_ptr);
            unsafe {
                self.map
                    .insert(KeyRef::new((*entry_ptr).key.as_ptr()), entry);
            }
            None
        }
    }
}

存在該元素時，將進行替換

unsafe {
    mem::swap(&mut *(*entry_ptr).val.as_mut_ptr(), &mut v);
}

並且重新維護訪問佇列，需要detach然後重新attach使其在佇列的最前面，保證最近訪問最晚淘汰，從而實現Lru。
如果元素不存在，那麼將判斷是否快取佇列為滿，如果滿了將要淘汰的資料進行替換，如果未滿建立新的元素，即replace_or_create_node。

元素刪除

在將元素刪除時，需要維護好我們的佇列，防止出現佇列錯誤及野指標等

pub fn remove<Q>(&mut self, k: &Q) -> Option<(K, V)>
    where
        K: Borrow<Q>,
        Q: Hash + Eq + ?Sized,
{
    match self.map.remove(KeyWrapper::from_ref(k)) {
        Some(l) => unsafe {
            self.detach(l.as_ptr());
            let node = *Box::from_raw(l.as_ptr());
            Some((node.key.assume_init(), node.val.assume_init()))
        },
        None => None,
    }
}

這裡因為移除時，僅僅需要一個可以轉化成K的引用即可以，並不需要嚴格的K，利於程式設計。例如:

let mut map = LruCache::new(2);
map.insert("aaaa".to_string(), "bbb");
map.remove("aaaa");
assert!(map.len() == 0);

在此處我們就不需要嚴格的構建String物件。由於map中的key我們用的是KeyRef，在這裡，我們構建一個KeyWrapper進行轉化。

// 確保新型別與其內部型別的記憶體佈局完全相同
#[repr(transparent)]
struct KeyWrapper<Q: ?Sized>(Q);

impl<K, Q> Borrow<KeyWrapper<Q>> for KeyRef<K>
where
    K: Borrow<Q>,
    Q: ?Sized,
{
    fn borrow(&self) -> &KeyWrapper<Q> {
        let key = unsafe { &*self.k }.borrow();
        KeyWrapper::from_ref(key)
    }
}

如果移除成功，那麼將從雙向連結串列中同步移除，並且將指標中的值重新變成Rust中的物件，以便可以及時被drop，避免記憶體洩漏。

self.detach(l.as_ptr());
let node = *Box::from_raw(l.as_ptr());
Some((node.key.assume_init(), node.val.assume_init()))

其它操作

  • pop移除棧頂上的資料，最近使用的
  • pop_last移除棧尾上的資料，最久未被使用的
  • contains_key判斷是否包含key值
  • raw_get直接獲取key的值，不會觸發雙向連結串列的維護
  • get獲取key的值，並維護雙向連結串列
  • get_mut獲取key的值，並可以根據需要改變val的值
  • retain 根據函式保留符合條件的元素

如何使用

在cargo.toml中新增

[dependencies]
algorithm = "0.1"
示例
use algorithm::LruCache;
fn main() {
    let mut lru = LruCache::new(3);
    lru.insert("now", "ok");
    lru.insert("hello", "algorithm");
    lru.insert("this", "lru");
    lru.insert("auth", "tickbh");
    assert!(lru.len() == 3);
    assert_eq!(lru.get("hello"), Some(&"algorithm"));
    assert_eq!(lru.get("this"), Some(&"lru"));
    assert_eq!(lru.get("now"), None);
}

完整專案地址

https://github.com/tickbh/algorithm-rs

結語

Lru在快取場景中也是極其重要的一環，但是普通的Lru容易將熱點資料進行移除，如果短時間內大量的資料進入則會將需要快取的資料全部清空，後續將介紹改進演算法Lru-k及Lfu演算法。

相關文章