RocksDB restarting test for memory profiling

rocksdb_restart_memory_test.cc

#include <cstdio>
#include <cunistd>
#include <iostream>
#include <memory>
#include <string>

#include "rocksdb/convenience.h"
#include "rocksdb/db.h"
#include "rocksdb/slice.h"
#include "rocksdb/options.h"

using namespace ROCKSDB_NAMESPACE;

#if defined(OS_WIN)
std::string kDBPath = "C:\\Windows\\TEMP\\db_close_memory_test";
#else
std::string kDBPath = "/tmp/db_close_memory_test";
#endif

class DBCloseMemoryTestDriver {
public:
  DBCloseMemoryTestDriver() {
    Init();
  }

  ~DBCloseMemoryTestDriver() {
    // Destroy column family handles befor deleting DB object
    m_db->DestroyColumnFamilyHandle(m_default_handle);
    m_db->DestroyColumnFamilyHandle(m_cfh1_handle);
  }

  int ExecuteWriteAndReadOperations(const std::string& key_prefix,
    const int count_of_keys_to_write) {
    // Write
    WriteOptions write_options;
    for (int i = 0; i < count_of_keys_to_write; ++i) {
      std::string value = std::to_string(i);
      std::string key = key_prefix + "," + value;
      Status s = m_db->Put(write_options, m_cfh1_handle, key, value);
      assert(s.ok());
    }

    
    // Read
    ReadOptions read_options;
    int counter = 0;
    std::unique_ptr<Iterator> iterator(m_db->NewIterator(read_options, m_cfh1_handle));
    iterator->SeekToFirst();
    for (iterator->SeekToFirst(); iterator->Valid(); iterator->Next()) {
      ++counter;
    }

    return counter;
  }

  void RestartRocksDB() {
    FlushOptions flush_options;
    Status s;
    // Below flush in the original test contained only default CF handle
    s = m_db->Flush(flush_options, { m_default_handle });
    assert(s.ok());
    s = m_db->FlushWAL(true);
    assert(s.ok());
    m_db->DestroyColumnFamilyHandle(m_default_handle);
    m_db->DestroyColumnFamilyHandle(m_cfh1_handle);
    CancelAllBackgroundWork(m_db.get(), true);
    m_db.reset();
    //std::cerr << "DB objects cleaned" << std::endl;
    //usleep(5000000);
    Init();
  }

private:

  void Init() {
    Options options;

    // create the DB if it's not already present
    options.create_if_missing = true;
    // create column families if they are not present
    options.create_missing_column_families = true;
    options.db_write_buffer_size = 64 * 1024 * 1024;

    // Column descriptors
    ColumnFamilyOptions cf_options;
    ColumnFamilyDescriptor cf_1("cfh1", cf_options);
    ColumnFamilyDescriptor cf_2("default", cf_options);
    std::vector<ColumnFamilyDescriptor> cf_desc_vec{cf_1, cf_2};

    // Vector for column family handles
    std::vector<ColumnFamilyHandle*> cf_handle_vec;

    // open DB
    DB* db;
    Status s = DB::Open(options, kDBPath, cf_desc_vec, &cf_handle_vec, &db);
    assert(s.ok());

    m_db.reset(db);
    m_cfh1_handle = cf_handle_vec[0];
    m_default_handle = cf_handle_vec[1];
  }

  std::unique_ptr<DB> m_db;
  ColumnFamilyHandle* m_cfh1_handle;
  ColumnFamilyHandle* m_default_handle;
};

int main() {
  DBCloseMemoryTestDriver test_driver;

  for (int i = 0; i < 10; ++i) {
    int keys_read_count = test_driver.ExecuteWriteAndReadOperations("Attempt: " + i, 1000000);
    std::cerr << "Read " + std::to_string(keys_read_count) + " keys on attempt " + std::to_string(i) << std::endl;
    test_driver.RestartRocksDB();
  }

  return 0;
}