Rocksdb 获取当前db内部的有效key个数 （估值）

Rocksdb因为是AppendOnly 方式写入，所以没有办法提供db内部唯一key个数的接口（可能存在多版本的key，对用户来说只有一个userkey，但是rocksdb认为是多个internal key）。
不过Rocksdb支持提供获取大概非删除key 的internal-key的个数接口，也能让用户对写入的key有一个大体量级的估计。

本文相关rocksdb代码版本是6.4.6

1. 基本接口

可以通过如下接口来获取：

uint64_t int_num;
dbfull()->GetIntProperty("rocksdb.estimate-num-keys", &int_num)

这一部分key组成从Rocksdb的写入整体架构来看应该由三部分组成:

• memtable-keys
• imm memtable-keys
• sstables-keys

实际获取的时候还需要将删除key(DeleteType)过滤掉。

具体estimate-num-keys的底层实现接口是:

bool InternalStats::HandleEstimateNumKeys(uint64_t* value, DBImpl* /*db*/,
                                          Version* /*version*/) {
  // Estimate number of entries in the column family:
  // Use estimated entries in tables + total entries in memtables.
  const auto* vstorage = cfd_->current()->storage_info(); 
  uint64_t estimate_keys = cfd_->mem()->num_entries() + // memtable 的keys
                           cfd_->imm()->current()->GetTotalNumEntries() + // imm 的keys
                           vstorage->GetEstimatedActiveKeys(); // sstables 的keys
  uint64_t estimate_deletes = // 删除keys
      cfd_->mem()->num_deletes() + cfd_->imm()->current()->GetTotalNumDeletes();
  *value = estimate_keys > estimate_deletes * 2
               ? estimate_keys - (estimate_deletes * 2)
               : 0;
  return true;
}

2. Memtable key个数统计

针对memtable中的keys的统计会获取mem()->num_entries()中的num_entries_的个数，memtable即active-memtable 在写入路径中达到write-buffer-size阈值之前有且仅有一个，所以只需要看一下当前的用户进程中这一个memtable中写入的key的个数即可，而这个数据会在Add memtable的时候同步更新。

3. Immutable Memtable key个数统计

针对immutable memtable的keys统计，因为这种只读的memtable可能存在多个，由一个链表进行管理，后续统一进行flush，所以获取对应的有效key以及删除key的个数的话只需要逐个累加immutable memtable中的key的个数即可。

uint64_t MemTableListVersion::GetTotalNumEntries() const {
  uint64_t total_num = 0;
  for (auto& m : memlist_) {
    total_num += m->num_entries();
  }
  return total_num;
}

这一些有效key的更新是在Status MemTable::Add()函数中进行更新的，包括后续的num_deletes_的个数也一样。

4. Sstables key个数统计

这个数据的统计也是我们想要的数据主体，因为大多数的时候 我们数据还是会持久化到sst文件之中的。
针对sst的有效key的个数统计是通过如下接口实现的：
大体逻辑是

• 如果发现sst文件个数为0，则直接返回0。
  这个场景接口也说明了，在用户写入少量 key没有达到触发flush的条件时，这里获取到的数据就是0
• 如果当前统计的有效key的个数current_num_non_deletions_ 比实际的删除key的个数还少，则也认为这一些有效key后续都会被删除，也返回0
• 有效key的个数 通过current_num_non_deletions_ - current_num_deletions_ 即可获得
• 为了防止返回的key数量过多（compaction完成之后会伴随着文件的删除），这里会重新逐层获取一下sst文件数量，和实际统计的sst文件数量做一个比值再乘上一步有效key的个数 – 双重保险

uint64_t VersionStorageInfo::GetEstimatedActiveKeys() const {
  // Estimation will be inaccurate when:
  // (1) there exist merge keys
  // (2) keys are directly overwritten
  // (3) deletion on non-existing keys
  // (4) low number of samples
  if (current_num_samples_ == 0) {
    return 0;
  }

  if (current_num_non_deletions_ <= current_num_deletions_) {
    return 0;
  }

  uint64_t est = current_num_non_deletions_ - current_num_deletions_;

  uint64_t file_count = 0;
  for (int level = 0; level < num_levels_; ++level) {
    file_count += files_[level].size();
  }

  if (current_num_samples_ < file_count) {
    // casting to avoid overflowing
    return
      static_cast<uint64_t>(
        (est * static_cast<double>(file_count) / current_num_samples_)
      );
  } else {
    return est;
  }
}

而实际的这一些指标的填充都是VersionStorageInfo::UpdateAccumulatedStats 这个函数中，这个函数的调用链实际能贯穿到compaction ，如下调用链

DBImpl::BackgroundCompaction // compaction执行入口
	VersionSet::LogAndApply // compaction 执行结束前需要更新manifest的version信息
		VersionSet::ProcessManifestWrites // manifest更新入口
			Version::PrepareApply // 更新各个指标
				Version::UpdateAccumulatedStats // 更新当前version内所有层的所有文件元信息
					VersionStorageInfo::UpdateAccumulatedStats // 更新每一个sst文件的元信息

最后一个函数的的更新一个sst文件元数据指标方式如下：

void VersionStorageInfo::UpdateAccumulatedStats(FileMetaData* file_meta) {
  assert(file_meta->init_stats_from_file);
  accumulated_file_size_ += file_meta->fd.GetFileSize();
  accumulated_raw_key_size_ += file_meta->raw_key_size;
  accumulated_raw_value_size_ += file_meta->raw_value_size;
  accumulated_num_non_deletions_ +=
      file_meta->num_entries - file_meta->num_deletions;
  accumulated_num_deletions_ += file_meta->num_deletions;

  current_num_non_deletions_ +=
      file_meta->num_entries - file_meta->num_deletions;
  current_num_deletions_ += file_meta->num_deletions;
  current_num_samples_++;
}

5. 疑问

可以看到以上代码中
current_num_non_deletions_的数值是file_meta->num_entries - file_meta->num_deletions; ，已经减去了当前文件被删除的数据条目，细心的同学可能会发现估算sst文件内key的个数的函数GetEstimatedActiveKeys中有一行代码uint64_t est = current_num_non_deletions_ - current_num_deletions_;，这里又减了一次统计的删除key，有点奇怪。

针对删除类型的key减了两次，这里想一想，rocksdb写入一个删除类型的key肯定表示需要删除之前一个已经存在的key，所以我们想要保证返回的是有效key，需要减掉当前删除key个数的两倍才行。