Lru-k在Rust中的實現及原始碼解析

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

LRU-K 是一種快取淘汰演算法，旨在改進傳統的LRU（Least Recently Used，最近最少使用）演算法的效能。將其中高頻的資料達到K次訪問移入到另一個佇列進行保護。

演算法思想

  • LRU-K中的“K”代表最近使用的次數。因此，LRU可以認為是LRU-1的特例。
  • LRU-K的主要目的是為了解決LRU演算法“快取汙染”的問題。其核心思想是將“最近使用過1次”的判斷標準擴充套件為“最近使用過K次”。

工作原理

  • LRU-K需要維護兩個佇列：歷史佇列和快取佇列。
    1. 普通佇列：儲存著每次訪問的頁面。當頁面訪問次數達到K次時，該頁面從歷史佇列中移除，並新增到K次佇列中。
    2. K次佇列：儲存已經訪問K次的頁面。當快取佇列滿了之後，需要淘汰頁面時，會淘汰最後一個頁面，即“倒數第K次訪問離現在最久”的那個頁面。
  • 詳細說明：
    1. 頁面第一次被訪問時，新增到普通佇列中。
    2. 當普通佇列中的頁面滿了，根據一定的快取策略（如FIFO、LRU、LFU）進行淘汰。
    3. 當歷史佇列中的某個頁面第K次訪問時，該頁面從歷史佇列中移除，並新增到K次佇列中。
    4. K次佇列中的頁面再次被訪問時，會重新排序。

優缺點

  • 優點：
    1. LRU-K降低了“快取汙染”帶來的問題，因為只有當頁面被訪問K次後才會被加入快取佇列。
    2. LRU-K的命中率通常比LRU要高。
  • 缺點：
    1. LRU-K需要維護一個普通佇列，因此記憶體消耗會比LRU多。
    2. LRU-K需要基於次數進行排序（可以需要淘汰時再排序，也可以即時排序），因此CPU消耗比LRU要高。
    3. 當K值較大時（如LRU-3或更大的K值），雖然命中率會更高，但適應性較差，需要大量的資料訪問才能將歷史訪問記錄清除掉。

實際應用

  • 在實際應用中，LRU-2通常被認為是綜合各種因素後最優的選擇。

綜上所述，LRU-K透過引入“最近使用過K次”的判斷標準，有效地解決了LRU演算法中的“快取汙染”問題，提高了快取的命中率。然而，它也需要更多的記憶體和CPU資源來維護歷史佇列和進行排序操作。

結構設計

與Lru的結構類似，K與V均用指標方式儲存，避免在使用過程中出現Copy或者Clone的可能，提高效能。
注:該方法用了指標會相應的出現許多unsafe的程式碼，因為在Rsut中，訪問指標都被認為是unsafe。我們也可以使用陣列座標模擬指標的方式來模擬。

節點設計

相對普通的Lru節點，我們需要額外儲存次數資料。

/// LruK節點資料
struct LruKEntry<K, V> {
    pub key: mem::MaybeUninit<K>,
    pub val: mem::MaybeUninit<V>,
    pub times: usize,
    pub prev: *mut LruKEntry<K, V>,
    pub next: *mut LruKEntry<K, V>,
}

類設計

由於有K次的列表，所以需要維護兩個列表，在空間佔用上會比Lru多一些，主要多一個欄位訪問次數的維護

pub struct LruKCache<K, V, S> {
    map: HashMap<KeyRef<K>, NonNull<LruKEntry<K, V>>, S>,
    cap: usize,
    /// 觸發K次數，預設為2
    times: usize,
    /// K次的佇列
    head_times: *mut LruKEntry<K, V>,
    tail_times: *mut LruKEntry<K, V>,
    /// 普通佇列
    head: *mut LruKEntry<K, V>,
    tail: *mut LruKEntry<K, V>,
    /// 普通佇列的長度
    lru_count: usize,
}

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

impl<K, V, S> LruCache<K, V, S> {
    /// 提供hash函式
    pub fn with_hasher(cap: usize, times: usize, hash_builder: S) -> LruKCache<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(LruKEntry::new_empty()));
        let tail = Box::into_raw(Box::new(LruKEntry::new_empty()));
        unsafe {
            (*head).next = tail;
            (*tail).prev = head;
        }
        let head_times = Box::into_raw(Box::new(LruKEntry::new_empty()));
        let tail_times = Box::into_raw(Box::new(LruKEntry::new_empty()));
        unsafe {
            (*head_times).next = tail_times;
            (*tail_times).prev = head_times;
        }
        Self {
            map,
            cap,
            times,
            head_times,
            tail_times,
            head,
            tail,
            lru_count: 0,
        }
    }
}

元素插入及刪除

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

pub fn capture_insert(&mut self, k: K, mut v: V) -> Option<(K, V, bool)> {
    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, true))
        }
        None => {
            let (val, 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);
            }
            val.map(|(k, v)| (k, v, false))
        }
    }
}

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,
        }
    }

與Lru的操作方式類似，但是主要集中在attach及detach因為有兩個佇列，需要正確的附著在正確的佇列之上。

attach
/// 加到佇列中
fn attach(&mut self, entry: *mut LruKEntry<K, V>) {
    unsafe {
        (*entry).times = (*entry).times.saturating_add(1);
        if (*entry).times < self.times {
            self.lru_count += 1;
            (*entry).next = (*self.head).next;
            (*(*entry).next).prev = entry;
            (*entry).prev = self.head;
            (*self.head).next = entry;
        } else {
            (*entry).next = (*self.head_times).next;
            (*(*entry).next).prev = entry;
            (*entry).prev = self.head_times;
            (*self.head_times).next = entry;
        }
    }
}

在加入到佇列的時候，需將訪問次數+1，然後判斷是否達到K次的次數，如果達到將其加入到head_times佇列中。
其中使用了saturating_add，這裡說個Rust與其它語言的差別。
因為在Rust中不像c語言，如果在c語言中，定義一個uchar型別

uchar times = 255;
times += 1; //此時times為0，不會有任何異常

但是在Rust中

let mut times: u8 = 255;
times = times.overflowing_add(1); // 此時times為0，因為上溢位了
times = times.saturating_add(1); // 此時times為255，因為達到了最大值
times += 1; // 此時將會發生panic

此時這函式的效率基本上等同於Lru的，相對僅僅是多維護times欄位及lru_count欄位。

detach
fn detach(&mut self, entry: *mut LruKEntry<K, V>) {
    unsafe {
        (*(*entry).prev).next = (*entry).next;
        (*(*entry).next).prev = (*entry).prev;

        if (*entry).times < self.times {
            self.lru_count -= 1;
        }
    }
}

與Lru中的類似，僅僅如果次數在k次以下的時候維護lru_count，效率基本一致。

replace_or_create_node
fn replace_or_create_node(&mut self, k: K, v: V) -> (Option<(K, V)>, NonNull<LruKEntry<K, V>>) {
    if self.len() == self.cap {
        let old_key = if self.lru_count > 0 {
            KeyRef {
                k: unsafe { &(*(*(*self.tail).prev).key.as_ptr()) },
            }
        } else {
            KeyRef {
                k: unsafe { &(*(*(*self.tail_times).prev).key.as_ptr()) },
            }
        };
        let old_node = self.map.remove(&old_key).unwrap();
        let node_ptr: *mut LruKEntry<K, V> = old_node.as_ptr();
        unsafe  {
            (*node_ptr).times = 0;
        }
        let replaced = unsafe {
            (
                mem::replace(&mut (*node_ptr).key, mem::MaybeUninit::new(k)).assume_init(),
                mem::replace(&mut (*node_ptr).val, mem::MaybeUninit::new(v)).assume_init(),
            )
        };

        self.detach(node_ptr);

        (Some(replaced), old_node)
    } else {
        (None, unsafe {
            NonNull::new_unchecked(Box::into_raw(Box::new(LruKEntry::new(k, v))))
        })
    }

淘汰資料，優先淘汰普通佇列的資料，如果普通佇列為空，將進入淘汰K次佇列。區別就是在於淘汰時多選擇一次資料。效率上也基本上可以忽略不計。

其它操作

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

如何使用

在cargo.toml中新增

[dependencies]
algorithm = "0.1"
示例
use algorithm::LruKCache;
fn main() {
    let mut lru = LruKCache::with_times(3, 3);
    lru.insert("this", "lru");
    for _ in 0..3 {
        let _ = lru.get("this");
    }
    lru.insert("hello", "algorithm");
    lru.insert("auth", "tickbh");
    assert!(lru.len() == 3);
    lru.insert("auth1", "tickbh");
    assert_eq!(lru.get("this"), Some(&"lru"));
    assert_eq!(lru.get("hello"), None);
    assert!(lru.len() == 3);
}

完整專案地址

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

結語

Lru-k與lru的區別在於多維護一個佇列，及每個元素多維護一個次數選項，對於效能的影響不大，僅僅多耗一點cpu，但是可以相應的提高命中率，下一章將介紹LFU按頻次的淘汰機制。

相關文章