优化 RisingWave 中 LSM-Tree Iterator 的 rust 代码

作者：温一鸣，RisingWave Labs 内核开发工程师

​本文内容源于我们在 Global Tour of Rust @上海活动的回顾，点击链接可查看活动现场完整分享视频。

1. Hummock读写路径简介

在 RisingWave 中，我们用 rust 自研了云原生的 LSM 存储引擎 Hummock ，并将之用来存储流计算中有状态算子的状态。

与一般的 LSM 存储引擎类似，新写入 Hummock 的数据将会写入 mutable mem-table 中，在特定条件下，mutable mem-table 会被 freeze 成 immutable mem-table。最终immutable mem-table 会写成 SST(sorted string table)文件落入持久化的存储中，同时SST会在 LSM 的元数据中被加到 overlapping L0 中。经过 compaction，overlapping L0中的数据会进入 non-overlapping 的 LSM 下层。

从最上层 mutable mem-table 到最下层 non-overlapping level 的整体数据组织形式可见下图。

​

有状态算子会对状态存储进行 get (点查)或者 iter (范围查询)。

在处理 get 请求时，在经过 min-max 以及 bloom filter 的过滤后，Hummock 会从最上层的 mutable mem-table 到最下层的 non-overlapping levels 逐一查找，当查找到对应的 key 时，停止查找并返回对应的 value。

在处理 iter 请求时，与处理 get 请求不同，经过前期的过滤后，给定范围内的数据可能存在于任意一层数据中。因此，我们需要将每一层的数据进行合并。每一层数据包含若干路有序的数据。mutable mem-table和immutable mem-table都是内存中的一路有序数据，而 overlapping L0 中，每一个 SST 本身是一路有序数据，最后 non-overlapping levels 每一层中的 SST 相互不重叠，所以每一层是一路有序数据。从而我们可以对各层数据进行多路归并来处理范围查询。

在 RisingWave 中，每一路有序的数据都会抽象成 HummockIterator。HummockIterator 是一个 rust 的 trait，每一个有序的结构都会实现这个 trait。 经过简化后的 HummockIterator 的定义如下

#[async_trait]
pub trait HummockIterator: Send + 'static {
    async fn next(&mut self) -> HummockResult<()>;

    fn key(&self) -> &[u8];
    fn value(&self) -> Option<&[u8]>;
}

#[async_trait]
impl HummockIterator for MemtableIterator {...}

#[async_trait]
impl HummockIterator for SstIterator {...}

我们通过一个 MergeIterator 对多个 HummockIterator 使用 heap 进行多路归并。

由于 HummockIterator 是一个 trait，而具体实现这个 trait 的有多个类型的有序结构(如 mem-table iterator, SST iterator)，同时 rust 是静态类型的语言，无法直接将多个类型放到同一个容器中，因此我们使用 Box 来将多个类型的 HummockIterator 统一在一起，得到以下的 MergeIterator 的实现。

pub struct MergeIterator {
    heap: BinaryHeap<Box<dyn HummockIterator>>,
}

#[async_trait]
impl HummockIterator for MergeIterator {
    async fn next(&mut self) -> HummockResult<()> {
        if let Some(top) = self.heap.peek_mut() {
            top.next().await?
        }
        Ok(())
    }

  ...
}

2. 代码中的动态分发

在以上代码中，一共有两处用到了动态分发，即 Box<dyn ...>。其中一处是为了统一多种实现了 ·HummockIterator trait的类型而使用了 Box ，另一处则为 #[async_trait] 宏的使用。

由于在 next 方法中可能涉及 IO，例如拉取 SST 中的下一个 block，在 HummockIterator 的定义中，next 被设计成了一个异步方法，使得在调用栈底层进行 IO 时可以被用户态的调度器挂起。rust 中的异步方法并不会立刻返回其返回值，而是根据具体实现方法的代码，返回一个实现 Future trait的匿名类型，最终通过对这个匿名类型调用其实现的 Future trait中的 poll 方法，获得最终的返回值。在 rust 中，对于返回类型相同的两个异步方法，由于其具体实现的代码不同，因此他们所返回的实现了 Future trait 的类型也不同。

在不同的对 HummockIterator的实现中，由于他们对 next 这个异步方法的实现代码不同，因此在调用 next 方法时返回的实现了 Future trait 的类型也不同。然而，返回类型不确定的 trait 不是 object safe 的，无法使用 Box<dyn ...>。而 async_trait 这个宏，将实现的异步方法时的实现 Future trait 的返回值，全部转成了通过动态分发的 BoxFuture，得到了一个统一的返回类型。

尽管动态分发为代码带来了便利，但是在多路归并这种 CPU 密集的场景中，动态分发将会带来不小的开销。因此我们尝试将代码中的动态分发改成静态分发，以减小运行时动态分发带来的开销。

3. 优化动态分发

首先，我们尝试将 async_trait 这个宏去掉。在去掉宏以后，在不同的 HummockIterator 的实现中，他们不再返回统一的 BoxFuture，而是返回其代码所对应的实现了 Future trait 的返回类型。因此，我们可以看成在不同的 HummockIterator 的实现中，都有一个实现了 Future trait 的类型，作为 HummockIterator 的这个实现中关联类型。于是我们的 trait 可以先修改成以下形式，其中 NextFuture 为具体实现 next 方法时所产生的关联类型。

pub trait HummockIterator: Send + 'static {
    type NextFuture:  
        Future<Output = HummockResult<()>> + Send;
    
    fn next(&mut self) -> Self::NextFuture;
    fn key(&self) -> &[u8];
    fn value(&self) -> Option<&[u8]>;
}

而在实现 HummockIterator 的代码中，我们可以通过 TAIT (trait alias impl trait)，将我们在实现 next 方法时所产生的实现 Future trait 的匿名类，指定为该类型在实现 HummockIterator 时的 NextFuture 关联类型。

impl HummockIterator for MergeIterator {
    type NextFuture =
        impl Future<Output = HummockResult<()>>;
    
    fn next(&mut self) -> Self::NextFuture {
        async move {
            if let Some(top) = self.heap.peek_mut() {
                top.next().await?
            }
            Ok(())
        }
    }

    ...
}

但上述代码在实现时会遇到以下报错：

fn next(&mut self) -> Self::NextFuture {
        |--------- hidden type `[async block@src/lib.rs:87:9: 92:10]` captures the anonymous lifetime defined here

以上问题的原因是，在 next 的实现的 Future 中使用了 self，因此捕获了 self 的生命周期。在其返回值中没有表达出其捕获了生命周期，因此报错。要解决这个问题，我们需要给 NextFuture 带上生命周期。此时我们可以使用 rust 的 GAT(generic associated type) 来为关联类型带上生命周期。

pub trait HummockIterator: Send + 'static {
    type NextFuture<'a>:  
        Future<Output = HummockResult<()>> + Send + 'a
    where Self: 'a;
    
    fn next(&mut self) -> Self::NextFuture<'_>;
    fn key(&self) -> &[u8];
    fn value(&self) -> Option<&[u8]>;
}

在上述修改下，我们最终可以在不使用 async_trait 的情况下，定义并实现包含异步方法的 HummockIterator。而在我们进行多路归并的 MergeIterator 中，我们可以使用 HummockIterator 的泛型来取代之前的 Box。

pub struct MergeIterator<I: HummockIterator> {
    heap: BinaryHeap<I>,
}

此时，MergeIterator 只能接受一个实现了 HummockIterator 的类型，但在实际应用中，MergeIterator 需要接受多种类型的 HummockIterator。此时，我们可以通过 enum，对不同类型的 HummockIterator 进行手动转发，并合并成一个类型作为 MergeIterator的泛型参数。

pub enum HummockIteratorUnion<
    I1: HummockIterator,
    I2: HummockIterator,
    I3: HummockIterator,
> {
    First(I1),
    Second(I2),
    Third(I3),
}

impl<
    I1: HummockIterator<Direction = D>,
    I2: HummockIterator<Direction = D>,
    I3: HummockIterator<Direction = D>,
> HummockIterator for HummockIteratorUnion<I1, I2, I3>
{
    type NextFuture<'a> = impl Future<Output = HummockResult<()>> + 'a;

    fn next(&mut self) -> Self::NextFuture<'_> {
        async move {
            match self {
                First(iter) => iter.next().await,
                Second(iter) => iter.next().await,
                Third(iter) => iter.next().await,
            }
        }
    }
    
    ...
}

最终，一个静态类型的 MergeIterator 的具体类型是

type HummockMergeIterator = MergeIterator<
  HummockIteratorUnion<
    // For mem-table
    MemtableIterator,
    // For overlapping level SST
    SstIterator,
    // For non-overlapping level sorted runs
    ConcatIterator<SstIterator>,
  >
>;

至此，我们完成了对代码中动态分发的优化，我们对优化前后的代码跑 benchmark 对比优化前后的性能，发现得到了很可观的性能提升。

用时

用时减少

box dyn

309.58 ms

0%

单类型 MergeIterator

198.94 ms

-35.7%

多类型 MergeIterator

237.88 ms

-23.2%

4. 简化代码

上述代码中，HummockIterator 在定义和实现中，都需要对关联类型做精细的处理，导致代码十分复杂。

在最新的 rust nightly 版本中，rust 提供了 feature impl_trait_in_assoc_type，使得我们可以在 trait 的定义中不通过关联类型，直接去定义返回的 Future。而如果使用 feature async_fn_in_trait，在实现 trait 中的异步方法时，我们无需自己将代码包在一个 async block 中，而是可以直接将代码当成一个 async 方法来实现。

最终，我们可以将上述代码简化成以下代码:

pub trait HummockIterator: Send + 'static {
    fn next(&mut self) ->
        impl Future<Output = HummockResult<()>> + Send + '_;
    fn key(&self) -> &[u8];
    fn value(&self) -> Option<&[u8]>;
}

impl HummockIterator for MergeIterator {
    async fn next(&mut self) -> HummockResult<()> {
        if let Some(top) = self.heap.peek_mut() {
            top.next().await?
        }
        Ok(())
    }
    ...
}

注: 如果不是因为 tokio 对 Future 有 Send 的要求，在上述 trait 的定义中，可以直接将 next 定义成[1]

async fn next(&mut self) -> HummockResult<()>;

关于 RisingWave

RisingWave 是一款分布式 SQL 流处理数据库，旨在帮助用户降低实时应用的的开发成本。作为专为云上分布式流处理而设计的系统，RisingWave 为用户提供了与 PostgreSQL 类似的使用体验，并且具备比 Flink 高出 10 倍的性能以及更低的成本。了解更多：

GitHub: risingwave.com/github

官网: risingwave.com

公众号： RisingWave 中文开源社区