TiKV 源码解析系列文章（十一）Storage - 事务控制层

背景知识

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。Engine 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 接口来进行的事务约束检验。