Performance benchmark of leveldb skiplist, compared with unordered_map

skiplist_benchmark.cpp

#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <string>
#include <vector>
#include <set>
#include <unordered_set>
#include <map>
#include "benchmark/benchmark.h"
#include "db/skiplist.h"
#include "util/arena.h"
#include "util/random.h"

// 使用整数 Key 和简单的整数比较器
struct KeyComparator {
  int operator()(const uint64_t& a, const uint64_t& b) const {
    if (a < b) return -1;
    if (a > b) return 1;
    return 0;
  }
};

// 生成具有固定长度且唯一的键
std::vector<uint64_t> GenerateKeys(int num_keys, int key_length_bits, leveldb::Random* rand) {
  std::unordered_set<uint64_t> key_set;
  std::vector<uint64_t> keys;
  keys.reserve(num_keys);
  while (static_cast<int>(key_set.size()) < num_keys) {
    uint64_t key = rand->Next() & ((1ULL << key_length_bits) - 1);
    if (key_set.insert(key).second) { // 插入成功表示键唯一
      keys.push_back(key);
    }
  }
  return keys;
}

// 生成一个不在集合中的随机键
uint64_t GenerateUniqueKey(const std::unordered_set<uint64_t>& existing_keys, int key_length_bits, leveldb::Random* rand) {
  uint64_t key;
  do {
    key = rand->Next() & ((1ULL << key_length_bits) - 1);
  } while (existing_keys.find(key) != existing_keys.end());
  return key;
}

// 测试上下文，包含 map、SkipList 和随机数生成器
struct TestContext {
    std::map<uint64_t, int> test_map;
    leveldb::SkipList<uint64_t, KeyComparator>* skiplist;
    leveldb::Arena* arena;
    uint64_t key_to_insert;
    std::unordered_set<uint64_t> existing_keys;
    std::vector<uint64_t> existing_keys_vector;
    leveldb::Random rand;       
    TestContext() : rand(1234), skiplist(nullptr), arena(nullptr) {}

    ~TestContext() {
        delete skiplist;
        delete arena;
    }
};

static void SetUp(const benchmark::State& state, TestContext* context) {
    const int initial_size = state.range(0);      // 初始链表大小
    const int key_length_bits = 32;   
    context->test_map.clear();
    context->existing_keys.clear();
    context->existing_keys_vector.clear();
    delete context->skiplist;
    delete context->arena;
    context->arena = new leveldb::Arena();
    context->skiplist = new leveldb::SkipList<uint64_t, KeyComparator>(KeyComparator(), context->arena);
    std::vector<uint64_t> initial_keys = GenerateKeys(initial_size, key_length_bits, &context->rand); 
    for (const auto& key : initial_keys) {
      context->skiplist->Insert(key);
      context->test_map[key] = key;
      context->existing_keys.insert(key);
      context->existing_keys_vector.push_back(key);
    }

    context->key_to_insert = GenerateUniqueKey(context->existing_keys, key_length_bits, &context->rand);
}


// 插入性能测试 - SkipList 单个插入
static void BM_SkipListInsertSingle(benchmark::State& state) {
  TestContext context;
  SetUp(state, &context);

  for (auto _ : state) {
    context.skiplist->Insert(context.key_to_insert);
    benchmark::DoNotOptimize(context.skiplist);
    state.PauseTiming();
    SetUp(state, &context);
    state.ResumeTiming();
  }

  state.SetLabel("SkipList Single Insert");
}

BENCHMARK(BM_SkipListInsertSingle)->Arg(1000)->Arg(10000)->Arg(100000);

static void BM_MapInsertSingle(benchmark::State& state) {
  TestContext context;
  SetUp(state, &context);

  for (auto _ : state) {
    context.test_map[context.key_to_insert] = 42;
    benchmark::DoNotOptimize(context.test_map);
    state.PauseTiming();
    context.test_map.erase(context.key_to_insert);
    state.ResumeTiming();
  }

  state.SetLabel("std::map Single Insert");
}

BENCHMARK(BM_MapInsertSingle)->Arg(1000)->Arg(10000)->Arg(100000);

// 查找性能测试 - SkipList
static void BM_SkipListFind(benchmark::State& state) {
  TestContext context;
  SetUp(state, &context);
  uint64_t key_to_find = context.existing_keys_vector[context.rand.Uniform(context.existing_keys_vector.size())];
  for (auto _ : state) {
    bool found = context.skiplist->Contains(key_to_find);
    benchmark::DoNotOptimize(found);
    state.PauseTiming();
    key_to_find = context.existing_keys_vector[context.rand.Uniform(context.existing_keys_vector.size())];
    state.ResumeTiming();
  }

  state.SetLabel("SkipList Find");
}

BENCHMARK(BM_SkipListFind)->Arg(1000)->Arg(10000)->Arg(100000);

// 查找性能测试 - std::map
static void BM_MapFind(benchmark::State& state) {
  TestContext context;
  SetUp(state, &context);
  uint64_t key_to_find = context.existing_keys_vector[context.rand.Uniform(context.existing_keys_vector.size())];
  for (auto _ : state) {
    bool found = context.test_map.find(key_to_find) != context.test_map.end();
    benchmark::DoNotOptimize(found);
    state.PauseTiming();
    key_to_find = context.existing_keys_vector[context.rand.Uniform(context.existing_keys_vector.size())];
    state.ResumeTiming();
  }

  state.SetLabel("std::map Find");
}

BENCHMARK(BM_MapFind)->Arg(1000)->Arg(10000)->Arg(100000);

// 运行所有的 benchmark 测试
BENCHMARK_MAIN();