技术科普｜实现 RocksDB 的前缀迭代器有多难？

小彩蛋🎉🥚

文末有一个思考题，欢迎在评论区给出你的答案，随机抽一位小伙伴送出社区周边～

一、背景

RocksDB 的迭代器由于只能 Seek 到一个起始位置，然后不断往后迭代，直到所有数据都迭代完成。而一个非常常见的需求就是按特定前缀进行扫描，直到遇到一个不满足给定前缀的数据，就不再迭代。

如果我的 RocksDB 中有以下 key-value，那么我按照key-的前缀进行扫描，那么应该依次读到key-1，key-2，key-3，key-4和key-5，而最后一个数据不满足前缀，不应该读到。

key-1 -> value1
key-2 -> value2
key-3 -> value3
key-4 -> value4
key-5 -> value5
some-other-key -> some-other-value

注意，这篇文章中因为涉及 bloom filter，所以假设所有数据都在 sst 中。如果自己想验证的同学，记得在读之前 Flush 一下，保证数据不在 Memtable 中。

为了实现这个前缀扫描迭代器，我们就不能直接使用 RocksDB 中的Iterator，而需要在此基础上封装一层，比如下面的RocksPrefixIterator:

// Iterator that compares with prefix during iteration
class RocksPrefixIterator {
 public:
  explicit RocksPrefixIterator(rocksdb::Iterator* iter,
                               const std::string& prefix)
      : iter_(iter), prefix_(prefix) {}

  bool valid() const {
    // c++20 required
    return !!iter_ && iter_->Valid() && key().starts_with(prefix_);
  }

  void next() {
    iter_->Next();
  }

  std::string_view key() const {
    return iter_->key().ToStringView();
  }

  std::string_view value() const {
    return iter_->value().ToStringView();
  }

 protected:
  std::unique_ptr<rocksdb::Iterator> iter_;
  std::string prefix_;
};

其内部仍然是通过操作 RocksDB 的Iterator来进行迭代、读取 key 以及 value ，但在判断迭代器是否有效时，需要额外判断当前 key 是否匹配给定前缀。

它的使用方式也很简单，首先我们构造一个Iterator，通过Seek将其指向正确位置，然后将Iterator和prefix都保存到RocksPrefixIterator中，之后就正常迭代即可。

std::unique_ptr<RocksPrefixIterator> prefixScan(DB* db, const std::string& prefix) {
  ReadOptions read_options;
  Iterator* iter = db->NewIterator(read_options);
  iter->Seek(rocksdb::Slice(prefix));
  return std::make_unique<RocksPrefixIterator>(iter, prefix);
}

void foo() {
  auto iter = prefixScan(db, "key-");
  for (; iter->valid(); iter->next()) {
    cout << iter->key() << "\n";
  }
}

需要注意的是，RocksPrefixIterator中的prefix_，它不能是std::string_view、char*或者folly::StringPiece这样的指针类型，而必须是std::string这样有所有权的类型。

考虑上面的例子，我们在调用prefixScan函数时，传入的是一个临时的string对象，如果RocksPrefixIterator中的prefix_是一个string_view，那么在prefixScan函数返回之后，临时string对象已经析构，prefix_就指向了一块已经释放的内存。

到这我们已经有了一个基础版本的前缀扫描迭代器，我们接下来看看能否有一些优化的手段。

二、Prefix Seek API

RocksDB 在 3.0 这个版本退出了一直沿用到现在的 prefix seek api，如果不熟悉的，请务必先简单看下官方 wiki.

🔍https://github.com/facebook/rocksdb/wiki/Prefix-Seek

首先，prefix seek api 的核心功能是通过 prefix bloom filter 减少 IO 和 CPU 消耗，这里我们上原文：

The basic idea is that, if users know the iterating will be within one key prefix, the common prefix can be used to reduce costs.

The most commonly used prefix iterating technique is prefix bloom filter. If many sorted runs don’t contain any entry for this prefix, it can be filtered out by a bloom filter, and some I/Os and CPU for the sorted run can be ignored.

如果我们无法确保用户会迭代这个 Iterator 多少次（即调用多少次Next），就不能保证迭代器一直使用同一个前缀，这种情况下就不能使用 prefix seek 功能。

之后，使用 prefix bloom filter 的方式就是定义prefix_extractor，这部分可以参考 wiki 中的示例。当我们给定了prefix_extractor，Iterator的Seek和Next操作的确能够通过 prefix bloom filter 过滤掉一部分不需要的数据。

Options options;

// Set up bloom filter
rocksdb::BlockBasedTableOptions table_options;
table_options.filter_policy.reset(rocksdb::NewBloomFilterPolicy(10, false));
table_options.whole_key_filtering = false;  // If you also need Get() to use whole key filters, leave it to true.
options.table_factory.reset(
    rocksdb::NewBlockBasedTableFactory(table_options));  // For multiple column family setting, set up specific column family's ColumnFamilyOptions.table_factory instead.

// Define a prefix. In this way, a fixed length prefix extractor. A recommended one to use.
options.prefix_extractor.reset(NewCappedPrefixTransform(3));

DB* db;
Status s = DB::Open(options, "/tmp/rocksdb",  &db);

然后，官方 wiki 就说了，光配置好prefix_extractor还不行，它的使用要依赖于怎么设置的ReadOptions.这里 wiki 中主要描述了ReadOptions中的这三个参数：

• total_order_seek
• prefix_same_as_start
• auto_prefix_mode

我们将看看有没有办法能够借助于 prefix seek api 优化我们实现前缀扫描。

（一）Manual prefix iterating

在介绍这三个参数之前，这里我们说下容易被误导的一个地方：虽然 prefix seek api 可以在Seek和Next这两个操作中使用 prefix bloom filter，但是默认配置下是非常容易迭代到数据不满足给定的前缀的数据。比如我们以最开始的数据为例，当我们设置了长度为 4 的prefix_extractor，如果我们就像下面这样使用默认ReadOptions来进行迭代key-：

  // ...
  options.prefix_extractor.reset(NewCappedPrefixTransform(4));
  DB* db;
  Status s = DB::Open(options, "/tmp/rocksdb",  &db);

  // ...

  ReadOptions read_options;
  Iterator* iter = db->NewIterator(read_options);
  // 这两行有没有都无所谓 因为默认值都是false
  // read_options.total_order_seek = false;
  // read_options.auto_prefix_mode = false;

  // 会依次返回所有前缀为"key-"的数据 以及some-other-key
  iter->Seek(rocksdb::Slice("key-"));
  for (; iter->Valid(); iter->Next()) {
    cout << iter->Key() << "\n";
  }

尽管我们配置了prefix_extractor，但事实上是会迭代到some-other-key这条数据的。这是因为当我们把所有key-为前缀的数据都遍历完之后，迭代器会指向一条不满足key-为前缀的数据，也就超出了prefix_extractor所指向的合法范围（即key-）。

所以这种最常见的用法，反而会导致错误行为，这也就是 wiki 中的 Manual prefix iterating.用户需要在这种用法下自行确保迭代器正确使用，否则迭代结果是未定义的，可能会出现以下错误：

• 该读到的 key 没有读到
  
• 读到的 key 顺序错乱
• 读到不该读到的 key（比如被删掉的 key）
• 迭代非常慢

Manual prefix iterating 是为了保证兼容老版本的默认行为，RocksDB 并不会返回任何错误。

所谓用户自行确保正确使用迭代器，其实也就是检查迭代器当前指向的 key 经过prefix_extractor转换而得 prefix 是否发生了变化。这种前缀不匹配的行为，在下文中我们简单称其为越界。比如RocksPrefixIterator在迭代过程中，在valid函数内检查了前缀是否匹配，从而保证只读到 5 条数据。

When options.prefix_extractor is not nullptr and with default ReadOptions, iterators are not guaranteed a total order of all keys, but only keys for the same prefix.

准确来说，在默认行为下：

1. Iterator 在Seek时，会通过prefix_extractor获取到 SeekKey 的前缀 prefix.（后文中 SeekKey 都是指调用Seek时传入的 key）
   
2. 如果有匹配 prefix 的key，那么会将迭代器指向第一个大于等于它的 key，并且对于所有满足前缀的 key，能够保证全序关系。
3. 如果没有，或者是经过若干次Next之后，RocksDB 可能会返回valid()=false，也有可能返回任意一个大于前一个的 key. 这也是为什么默认行为下会出现上面提到的各种错误。

在默认行为下，所谓的 prefix seek 不等同于 prefix scan，不能保证只扫描满足prefix_extracto前缀的数据。

下面我们开始介绍 prefix seek api 相关的参数。

（二）total_order_seek

Users can only use prefix bloom filter to read inside a specific prefix. If iterator is outside a prefix, the feature needs to be disabled for specific iterators to prevent wrong results.

配置了prefix_extractor的前提下，如果无法保证迭代过程中不越界，该怎么使用迭代器呢。这时可以设置total_order_seek，它能完全忽略掉 prefix bloom filter，保证迭代器输出的所有key全局有序。

如果配置了prefix_extractor，又无法确保迭代器不会越界，那么请设置total_order_seek为true.

  // ...
  options.prefix_extractor.reset(NewCappedPrefixTransform(4));
  DB* db;
  Status s = DB::Open(options, "/tmp/rocksdb",  &db);

  // ...

  ReadOptions read_options;
  // 会忽略掉设置的prefix_extractor
  read_options.total_order_seek = true;
  Iterator* iter = db->NewIterator(read_options);
  iter->Seek(rocksdb::Slice("key-"));
  // 会依次返回所有前缀为"key-"的数据 以及some-other-key
  for (; iter->Valid(); iter->Next()) {
    cout << iter->Key() << "\n";
  }

注意，total_order_seek只能保证全局顺序。它和默认行为一样是无法确保迭代器不会越界，因此也会读到some-other-key这条数据。

（三）prefix_same_as_start

当设置了prefix_same_as_start时，RocksDB 会进入所谓的 prefix mode，它的行为如下：

prefix mode 和默认行为有一个最大的区别，它能够保证当迭代器指向的 key 对应的前缀不匹配 SeekKey 的前缀时，一定会返回Valid()=false


举例来说，按下面的方式去遍历时，SeekKey 的前缀是key-，因此只会读到 5 条数据，不会读到some-other-key

  // ...
  options.prefix_extractor.reset(NewCappedPrefixTransform(4));
  DB* db;
  Status s = DB::Open(options, "/tmp/rocksdb",  &db);

  // ...

  ReadOptions read_options;
  // 设置为prefix mode
  read_options.prefix_same_as_start = true;

  Iterator* iter = db->NewIterator(read_options);
  // 注意SeekKey为"key-" 它经过prefix_extractor转换得到的prefix也是"key-"
  // 在prefix mode下就只会依次返回所有前缀为"key-"的数据
  iter->Seek(rocksdb::Slice("key-"));
  for (; iter->Valid(); iter->Next()) {
    cout << iter->Key() << "\n";
  }

乍一看之下，前缀扫描功能就可以借助prefix_same_as_start来完成，然而这其中隐藏了一个很大的隐患。

如果 SeekKey 长度大于prefix_extractor长度时，比如下面例子中我们 SeekKey 设置成key-3，按照前缀扫描的期望，我们只应该读到key-3这一条数据。然而在 RocksDB 的 prefix mode 下，实际会迭代到key-3，key-4和key-5。

  // ...
  options.prefix_extractor.reset(NewCappedPrefixTransform(4));
  DB* db;
  Status s = DB::Open(options, "/tmp/rocksdb",  &db);

  // ...

  ReadOptions read_options;
  // 设置为prefix mode
  read_options.prefix_same_as_start = true;

  Iterator* iter = db->NewIterator(read_options);
  iter->Seek(rocksdb::Slice("key-3"));
  // 会依次返回key-3, key-4和key-5
  for (; iter->Valid(); iter->Next()) {
    cout << iter->Key() << "\n";
  }

这是怎么回事呢？注意 SeekKey 为”key-3” 它经过prefix_extractor转换得到的 prefix 是”key-”，并且由于设置了prefix_same_as_start，RockDB 会处于 prefix mode 下，因此就只会返回大于等于 SeekKey，且所有前缀为”key-“的数据，也就是key-3，key-4和key-5。

在 prefix mode 下，所谓的 prefix seek 也不等同于 prefix scan，只能保证只扫描满足prefix_extractor前缀的数据，而不是只扫描所有满足 SeekKey 作为前缀的数据。

我们总结一下，在 prefix mode下 ，实际上 RocksDB 的行为会被 SeekKey 的长度，以及 SeekKey 经过prefix_extractor转换得到的 prefix 长度所影响，我们假设prefix_extractor都是CappedPrefixTransform(4)：

这也是 wiki 中完全没有提及的，而只是说 prone to misuse.在我看来，这都不是容易用错，恐怕是根本不可能用对吧。

（四）auto_prefix_mode

由于这个 prefix mode 实在是太难以正确使用，RocksDB 在 6.8 推出了一个新功能，也就是 auto prefix mode.它就像total_order_seek和prefix_same_as_start的缝合怪一样：

• 大多数情况下，整体行为和total_order_seek
一样，保证所有 key 的有序输出。有一种情况二者行为不一样，后面就解释
  
• 但过程中可能会利用一些能用的 prefix bloom filter，减少 IO 和 CPU 消耗。

  // ...
  options.prefix_extractor.reset(NewCappedPrefixTransform(4));
  DB* db;
  Status s = DB::Open(options, "/tmp/rocksdb",  &db);

  // ...

  ReadOptions read_options;
  // 默认为false
  read_options.auto_prefix_mode = true;

  Iterator* iter = db->NewIterator(read_options);
  iter->Seek(rocksdb::Slice("key-"));
  // 会依次返回所有前缀为"key-"的数据 以及some-other-key
  for (; iter->Valid(); iter->Next()) {
    cout << iter->Key() << "\n";
  }

所以 auto prefix mode 是官方 wiki 推荐使用 prefix bloom filter 的方法，并且说一旦稳定就会将其默认设置为 true.为什么至今默认值仍是 false 呢，当然是因为有 Bug 了…… 出现 Bug 的前提是，RocksDB 中写入了比prefix_extractor
指定长度还要小的数据。

我们可以看下 RocksDB 中的单测

🔍https://github.com/facebook/rocksdb/blob/v7.8.3/db/db_test2.cc#L6352

假设目前有以下这些 key：

a
b
b\x00XYZ
c

RocksDB 的配置为：

    Options options = CurrentOptions();
    BlockBasedTableOptions table_options =
        *options.table_factory->GetOptions<BlockBasedTableOptions>();
    table_options.filter_policy.reset(NewBloomFilterPolicy(10, false));
    options.table_factory.reset(NewBlockBasedTableFactory(table_options));
    options.statistics = CreateDBStatistics();
    options.comparator = BytewiseComparator();

    // 这两个自带的prefix_extractor在下面的例子中都有问题
    options.prefix_extractor.reset(NewCappedPrefixTransform(2));
    // options.prefix_extractor.reset(NewFixedPrefixTransform(2));

当处于 auto prefix mode 时，Slice("a\xff")时，虽然设置了iterate_upper_bound为b\x00，但此时应该读到b才对，而实际并没有读到。

    {
      ReadOptions ro;
      ro.total_order_seek = false;
      ro.auto_prefix_mode = true;
      ro.iterate_upper_bound = "b\x00";

      iterator->Seek(Slice("a\xff"));
      // !!! BUG !!! See "BUG" section of auto_prefix_mode.
      ASSERT_FALSE(iterator->Valid());
      EXPECT_EQ(1, TestGetAndResetTickerCount(options, stat));
      ASSERT_OK(iterator->status());
    }

而如果设置total_order_seek，则能正确读到b.

    {
      // To prove that is the wrong result, now use total order seek
      ReadOptions tos_ro = ro;
      tos_ro.total_order_seek = true;
      tos_ro.auto_prefix_mode = false;

      iterator.reset(db_->NewIterator(tos_ro));
      iterator->Seek(Slice(a_end_stuff, 2));
      ASSERT_TRUE(iterator->Valid());
      ASSERT_EQ("b", iterator->key().ToString());
      EXPECT_EQ(0, TestGetAndResetTickerCount(options, stat));
      ASSERT_OK(iterator->status());
    }

最后，目前根据代码实现来看，如果没有写入了比prefix_extractor指定长度还要短的数据时，auto_prefix_mode是不会有问题的。

四、前缀扫描迭代器

现在我们基本知道了 prefix seek api 的所有行为。回到如何实现前缀扫描迭代器，一开始我们实现的前缀扫描迭代器有没有什么可以优化的地方呢？那注意到一开始的示例中我们其实没有使用prefix_extractor，也就自然不能利用 prefix bloom filter.

那如果我们设置了prefix_extractor，比如下面的用法呢？

options.prefix_extractor.reset(NewCappedPrefixTransform(4));

那可以看到我们在开头实现的prefixScan函数使用的是默认行为即 manual prefix iterating，这种情况下不会利用 prefix bloom filter，但因为我们在RocksPrefixIterator中手动检查了前缀是否匹配，所以是能保证有序迭代输出的。

class RocksPrefixIterator {
  // ...
  bool valid() const {
    return !!iter_ && iter_->Valid() && key().starts_with(prefix_);
  }
  // ...
};

std::unique_ptr<RocksPrefixIterator> prefixScan(DB* db, const std::string& prefix) {
  ReadOptions read_options;
  Iterator* iter = db->NewIterator(read_options);
  iter->Seek(rocksdb::Slice(prefix));
  return std::make_unique<RocksPrefixIterator>(iter, prefix);
}

如果想利用 prefix bloom filter 来提升性能呢，最简单的办法就是使用 auto prefix mode（如果不会触发前面提到的 bug 的话）。

std::unique_ptr<RocksPrefixIterator> prefixScan(DB* db, const std::string& prefix) {
  ReadOptions read_options;
  Iterator* iter = db->NewIterator(read_options);
  read_options.auto_prefix_mode = true;
  iter->Seek(rocksdb::Slice(prefix));
  return std::make_unique<RocksPrefixIterator>(iter, prefix);
}

然而，在这样的模式下，RocksPrefixIterator中仍然需要每次在valid函数中手动检查前缀是否匹配。

那还有没有优化空间呢？什么情况能够不需要检查前缀是否匹配，即valid函数可以像下面这样实现呢？

  bool valid() const {
    return !!iter_ && iter_->Valid();
  }

✨欢迎在评论区给出你的思考与答案～随机抽一位小伙伴送出社区周边～

Reference

https://github.com/facebook/rocksdb/wiki/Prefix-Seek#manual-prefix-iterating