几年之前,我以为我已经基本掌握了RocksDB prefix bloom filter,没想到这两天经过一阵研究,又有了更深入的认知。

背景

RocksDB的迭代器由于只能Seek到一个起始位置,然后不断往后迭代,直到所有数据都迭代完成。而一个非常常见的需求就是按特定前缀进行扫描,直到遇到一个不满足给定前缀的数据,就不再迭代。如果我的RocksDB中有以下key-value,那么我按照key-的前缀进行扫描,那么应该依次读到key-1key-2key-3key-4key-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将其指向正确位置,然后将Iteratorprefix都保存到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_viewchar*或者folly::StringPiece这样的指针类型,而必须是std::string这样有所有权的类型。考虑上面的例子,我们在调用prefixScan函数时,传入的是一个临时的string对象,如果RocksPrefixIterator中的prefix_是一个string_view,那么在prefixScan函数返回之后,临时string对象已经析构,prefix_就指向了一块已经释放的内存。

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

Prefix Seek API

RocksDB在3.0这个版本退出了一直沿用到现在的prefix seek api,如果不熟悉的,请务必先简单看下官方wiki和之前整理过的文章

我几年间陆陆续续看过好几遍这篇wiki,始终觉得语焉不详。另外由于中英文语境不同,wiki中充斥了各种误导。Let’s see!

首先,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.

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

之后,使用prefix bloom filter的方式就是定义prefix_extractor,这部分可以参考wiki中的示例。当我们给定了prefix_extractorIteratorSeekNext操作的确能够通过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可以在SeekNext这两个操作中使用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条数据。

在原文中有这样一句话也充满误导,请注意,它只是在描述能保证哪些key的顺序关系,而不是说只会读到给定前缀的key

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

重点2:在默认行为下,所谓的prefix seek不等同于prefix scan,不能保证只扫描满足prefix_extractor前缀的数据。

下面我们开始介绍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全局有序。

重点3:如果配置了prefix_extractor,又无法确保迭代器不会越界,那么请设置total_order_seektrue

  // ...
  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,它的行为如下:

  1. IteratorSeek时,会通过prefix_extractor获取到SeekKey的前缀prefix。
  2. 如果有匹配prefix的key,那么会将迭代器指向第一个大于等于它的key,并且对于所有满足前缀的key,能够保证全序关系。
  3. 如果没有,或者是经过若干次Next之后,RocksDB一定会返回Valid()=false

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-3key-4key-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-3key-4key-5

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

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

  1. len(SeekKey) > len(prefix)时,比如SeekKey为key-3,此时prefix为key-。RocksDB会根据prefix bloom filter进行过滤,最终输出所有大于等于SeekKey,且前缀匹配的数据。
  2. len(SeekKey) = len(prefix)时,比如SeekKey和prefix均为key-。此时RocksDB也会根据prefix bloom filter进行过滤,最终输出所有大于等于SeekKey,且前缀匹配的数据。而由于SeekKey和prefix相同,实际上就是我们所要的前缀扫描。
  3. len(SeekKey) < len(prefix),比如SeekKey为key,此时prefix也为key,而RocksDB中没有任何prefix bloom filter匹配的数据,最终什么都不会返回。

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

auto_prefix_mode

由于这个prefix mode实在是太难以正确使用,RocksDB在6.8退出了一个新功能,也就是auto prefix mode。它就像total_order_seekprefix_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中的单测,假设目前有以下这些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_boundb\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();
  }

这里就不贴答案了,如果你真的看明白了这篇文章,我想这个问题是不难的。

正确答案是,当SeekKey的长度和prefix_extractor给定长度相同时,此时可以设置prefix_same_as_start为true,RocksDB能够保证当Seek或者Next若干次之后,前缀不匹配SeekKey的前缀时,Iterator的Valid会返回false,也就不需要再额外手动检查前缀是否匹配。所以一个可能的优化手段时,根据SeekKey的不同长度,在prefixScan函数中调用prefix seek api的不同模式。

Reference

Prefix Seek · facebook/rocksdb Wiki