作者：張金鵬

背景知識

TiKV 是一個強一致的支援事務的分散式 KV 儲存。TiKV 通過 raft 來保證多副本之間的強一致，事務這塊 TiKV 參考了 Google 的 Percolator 事務模型，並進行了一些優化。

當 TiKV 的 Service 層收到請求之後，會根據請求的型別把這些請求轉發到不同的模組進行處理。對於從 TiDB 下推的讀請求，比如 sum，avg 操作，會轉發到 Coprocessor 模組進行處理，對於 KV 請求會直接轉發到 Storage 進行處理。

KV 操作根據功能可以被劃分為 Raw KV 操作以及 Txn KV 操作兩大類。Raw KV 操作包括 raw put、raw get、raw delete、raw batch get、raw batch put、raw batch delete、raw scan 等普通 KV 操作。 Txn KV 操作是為了實現事務機制而設計的一系列操作，如 prewrite 和 commit 分別對應於 2PC 中的 prepare 和 commit 階段的操作。

本文將為大家介紹 TiKV 原始碼中的 Storage 模組，它位於 Service 與底層 KV 儲存引擎之間，主要負責事務的併發控制。TiKV 端事務相關的實現都在 Storage 模組中。

原始碼解析

接下來我們將從 Engine、Latches、Scheduler 和 MVCC 等幾個方面來講解 Storage 相關的原始碼。

1. Engine trait

TiKV 把底層 KV 儲存引擎抽象成一個 Engine trait（trait 類似其他語言的 interface），定義見 storage/kv/mod.rs。Engint trait 主要提供了讀和寫兩個介面，分別為 async_snapshot

和 async_write。呼叫者把要寫的內容交給 async_write，async_write 通過回撥的方式告訴呼叫者寫操作成功完成了或者遇到錯誤了。同樣的，async_snapshot 通過回撥的方式把資料庫的快照返回給呼叫者，供呼叫者讀，或者把遇到的錯誤返回給呼叫者。

pub trait Engine: Send + Clone + 'static {
    type Snap: Snapshot;
    fn async_write(&self, ctx: &Contect, batch: Vec<Modify>, callback: Callback<()>) -> Result<()>;
    fn async_snapshot(&self, ctx: &Context, callback: Callback<Self::Snap>) -> Result<()>;
}
 

只要實現了以上兩個介面，都可以作為 TiKV 的底層 KV 儲存引擎。在 3.0 版本中，TiKV 支援了三種不同的 KV 儲存引擎，包括單機 RocksDB 引擎、記憶體 B 樹引擎和 RaftKV 引擎，分別位於 storage/kv 資料夾下面的 rocksdb_engine.rs、btree_engine.rs 和 raftkv.rs。其中單機 RocksDB 引擎和記憶體紅黑樹引擎主要用於單元測試和分層 benchmark，TiKV 真正使用的是 RaftKV 引擎。當呼叫 RaftKV 的 async_write 進行寫入操作時，如果 async_write 通過回撥方式成功返回了，說明寫入操作已經通過 raft 複製給了大多數副本，並且在 leader 節點（呼叫者所在 TiKV）完成寫入了，後續 leader 節點上的讀就能夠看到之前寫入的內容。

2. Raw KV 執行流程

Raw KV 系列介面是繞過事務直接操縱底層資料的介面，沒有事務控制，比較簡單，所以在介紹更復雜的事務 KV 的執行流程前，我們先介紹 Raw KV 的執行流程。

Raw put

raw put 操作不需要 Storage 模組做額外的工作，直接把要寫的內容通過 engine 的 async_write 介面傳送給底層的 KV 儲存引擎就好了。呼叫堆疊為 service/kv.rs: raw_put -> storage/mod.rs: async_raw_put。

impl<E: Engine> Storage<E> {
    pub fn async_raw_put(
        &self,
        ctx: Context,
        cf: String,
        key: Vec<u8>,
        value: Vec<u8>,
        callback: Callback<()>,
    ) -> Result<()> {
        // Omit some limit checks about key and value here...
        self.engine.async_write(
            &ctx,
            vec![Modify::Put(
                Self::rawkv_cf(&cf),
                Key::from_encoded(key),
                value,
            )],
            Box::new(|(_, res)| callback(res.map_err(Error::from))),
        )?;
        Ok(())
    }
}

Raw get

同樣的，raw get 只需要呼叫 engine 的 async_snapshot 拿到資料庫快照，然後直接讀取就可以了。當然對於 RaftKV 引擎，async_snapshot 在返回資料庫快照之前會做一些檢查工作，比如會檢查當前訪問的副本是否是 leader（3.0.0 版本只支援從 leader 進行讀操作，follower read 目前仍然在開發中），另外也會檢查請求中攜帶的 region 版本資訊是否足夠新。

3. Latches

在事務模式下，為了防止多個請求同時對同一個 key 進行寫操作，請求在寫這個 key 之前必須先獲取這個 key 的記憶體鎖。為了和事務中的鎖進行區分，我們稱這個記憶體鎖為 latch，對應的是 storage/txn/latch.rs 檔案中的 Latch 結構體。每個 Latch 內部包含一個等待佇列，沒有拿到 latch 的請求按先後順序插入到等待佇列中，隊首的請求被認為拿到了該 latch。

#[derive(Clone)]
struct Latch {
    pub waiting: VecDeque<u64>,
}

Latches 是一個包含多個 Latch 的結構體，內部包含一個固定長度的 Vector，Vector 的每個 slot 對應一個 Latch。預設配置下 Latches 內部 Vector 的長度為 2048000。每個 TiKV 有且僅有一個 Latches 例項，位於 Storage.Scheduler 中。

pub struct Latches {
    slots: Vec<Mutex<Latch>>,
    size: usize,
}

Latches 的 gen_lock 介面用於計算寫入請求執行前所需要獲取的所有 latch。gen_lock 通過計算所有 key 的 hash，然後用這些 hash 對 Vector 的長度進行取模得到多個 slots，對這些 slots 經過排序去重得到該命令需要的所有 latch。這個過程中的排序是為了保證獲取 latch 的順序性防止出現死鎖情況。

impl Latches {
    pub fn gen_lock<H: Hash>(&self, keys: &[H]) -> Lock {
        // prevent from deadlock, so we sort and deduplicate the index.
        let mut slots: Vec<usize> = keys.iter().map(|x|
        self.calc_slot(x)).collect();
        slots.sort();
        slots.dedup();
        Lock::new(slots)
    }
}

4. Storage 和事務排程器 Scheduler

Storage

Storage 定義在 storage/mod.rs 檔案中，下面我們介紹下 Storage 幾個重要的成員：

engine：代表的是底層的 KV 儲存引擎。

sched：事務排程器，負責併發事務請求的排程工作。

read_pool：讀取執行緒池，所有隻讀 KV 請求，包括事務的非事務的，如 raw get、txn kv get 等最終都會在這個執行緒池內執行。由於只讀請求不需要獲取 latches，所以為其分配一個獨立的執行緒池直接執行，而不是與非只讀事務共用事務排程器。

gc_worker：從 3.0 版本開始，TiKV 支援分散式 GC，每個 TiKV 有一個 gc_worker 執行緒負責定期從 PD 更新 safepoint，然後進行 GC 工作。

pessimistic_txn_enabled： 另外 3.0 版本也支援悲觀事務，pessimistic_txn_enabled 為 true 表示 TiKV 以支援悲觀事務的模式啟動，關於悲觀事務後續會有一篇原始碼閱讀文章專門介紹，這裡我們先跳過。

pub struct Storage<E: Engine> {
    engine: E,
    sched: Scheduler<E>,
    read_pool: ReadPool,
    gc_worker: GCWorker<E>,
    pessimistic_txn_enabled: bool,
    // Other fields...
}

對於只讀請求，包括 txn get 和 txn scan，Storage 呼叫 engine 的 async_snapshot 獲取資料庫快照之後交給 read_pool 執行緒池進行處理。寫入請求，包括 prewrite、commit、rollback 等，直接交給 Scheduler 進行處理。Scheduler 的定義在 storage/txn/scheduler.rs 中。

Scheduler

pub struct Scheduler<E: Engine> {
    engine: Option<E>,
    inner: Arc<SchedulerInner>,
}

struct SchedulerInner {
    id_alloc, AtomicU64,
    task_contexts: Vec<Mutex<HashMap<u64, TaskContext>>>,
    lathes: Latches,
    sched_pending_write_threshold: usize,
    worker_pool: SchedPool,
    high_priority_pool: SchedPool,
    // Some other fields...
}

接下來簡單介紹下 Scheduler 幾個重要的成員：

id_alloc：到達 Scheduler 的請求都會被分配一個唯一的 command id。

latches：寫請求到達 Scheduler 之後會嘗試獲取所需要的 latch，如果暫時獲取不到所需要的 latch，其對應的 command id 會被插入到 latch 的 waiting list 裡，當前面的請求執行結束後會喚醒 waiting list 裡的請求繼續執行，這部分邏輯我們將會在下一節 prewrite 請求在 scheduler 中的執行流程中介紹。

task_contexts：用於儲存 Scheduler 中所有請求的上下文，比如暫時未能獲取所需 latch 的請求都會被暫存在 task_contexts 中。

sched_pending_write_threshold：用於統計 Scheduler 內所有寫入請求的寫入流量，可以通過該指標對 Scheduler 的寫入操作進行流控。

worker_pool，high_priority_pool：兩個執行緒池，寫請求在呼叫 engine 的 async_write 之前需要進行事務約束的檢驗工作，這些工作都是在這個兩個執行緒池中執行的。

prewrite 請求在 Scheduler 中的執行流程

下面我們以 prewrite 請求為例子來講解下寫請求在 Scheduler 中是如何處理的：

1）Scheduler 收到 prewrite 請求的時候首先會進行流控判斷，如果 Scheduler 裡的請求過多，會直接返回 SchedTooBusy 錯誤，提示等一會再發送，否則進入下一步。

2）接著會嘗試獲取所需要的 latch，如果獲取 latch 成功那麼直接進入下一步。如果獲取 latch 失敗，說明有其他請求佔住了 latch，這種情況說明其他請求可能也正在對相同的 key 進行操作，那麼當前 prewrite 請求會被暫時掛起來，請求的上下文會暫存在 Scheduler 的 task_contexts 裡面。當前面的請求執行結束之後會將該 prewrite 請求重新喚醒繼續執行。

impl<E: Engine> Scheduler<E> {
    fn try_to_wake_up(&self, cid: u64) {
        if self.inner.acquire_lock(cid) {
            self.get_snapshot(cid);
        }
    }
    fn release_lock(&self, lock: &Lock, cid: u64) {
        let wakeup_list = self.inner.latches.release(lock, cid);
        for wcid in wakeup_list {
            self.try_to_wake_up(wcid);
        }
    }
}

3）獲取 latch 成功之後會呼叫 Scheduler 的 get_snapshot 介面從 engine 獲取資料庫的快照。get_snapshot 內部實際上就是呼叫 engine 的 async_snapshot 介面。然後把 prewrite 請求以及剛剛獲取到的資料庫快照交給 worker_pool 進行處理。如果該 prewrite 請求優先順序欄位是 high 就會被分發到 high_priority_pool 進行處理。high_priority_pool 是為了那些高優先順序請求而設計的，比如 TiDB 系統內部的一些請求要求 TiKV 快速返回，不能由於 worker_pool 繁忙而被卡住。需要注意的是，目前 high_priority_pool 與 worker_pool 僅僅是語義上不同的兩個執行緒池，它們內部具有相同的作業系統排程優先順序。

4）worker_pool 收到 prewrite 請求之後，主要工作是從拿到的資料庫快照裡確認當前 prewrite 請求是否能夠執行，比如是否已經有更大 ts 的事務已經對資料進行了修改，具體的細節可以參考 Percolator 論文，或者參考我們的官方部落格 《TiKV 事務模型概覽》。當判斷 prewrite 是可以執行的，會呼叫 engine 的 async_write 介面執行真正的寫入操作。這部分的具體的程式碼見 storage/txn/process.rs 中的 process_write_impl 函式。

5）當 async_write 執行成功或失敗之後，會呼叫 Scheduler 的 release_lock 函式來釋放 latch 並且喚醒等待在這些 latch 上的請求繼續執行。

5. MVCC

TiKV MVCC 相關的程式碼位於 storage/mvcc 資料夾下，強烈建議大家在閱讀這部分程式碼之前先閱讀 Percolator 論文，或者我們的官方部落格 《TiKV 事務模型概覽》。

MVCC 下面有兩個比較關鍵的結構體，分別為 MvccReader 和 MvccTxn。MvccReader 位於 storage/mvcc/reader/reader.rs 檔案中，它主要提供讀功能，將多版本的處理細節隱藏在內部。比如 MvccReader 的 get 介面，傳入需要讀的 key 以及 ts，返回這個 ts 可以看到的版本或者返回 key is lock 錯誤等。

impl<S: Snapshot> MvccReader<S> {
    pub fn get(&mut self, key: &Key, mut ts: u64) -> Result<Option<Value>>;
}

MvccTxn 位於 storage/mvcc/txn.rs 檔案中，它主要提供寫之前的事務約束檢驗功能，上一節 prewrite 請求的處理流程中第四步就是通過呼叫 MvccTxn 的 prewrite 介面來進行的事務約束檢驗。

小結

TiKV 端事務相關的實現都位於 Storage 模組中，該文帶大家簡單概覽了下這部分幾個關鍵的點，想了解更多細節的讀者可以自行閱讀這部分的原始碼（code talks XD）。另外從 3.0 版本開始，TiDB 和 TiKV 支援悲觀事務，TiKV 端對應的程式碼主要位於 storage/lock_manager 以及上面提到的 MVCC 模組中。

原文閱讀：https://pingcap.com/blog-cn/tikv-soucre-code-reading-11/

更多 TiKV 原始碼閱讀：https://pingcap.com/blog-cn/#TiKV-%E6%BA%90%E7%A0%81%E8%A7%A3%E6%9E%90

相關推薦