Read from and write to different column families

Summary: This one is big. It adds ability to write to and read from different column families (see the unit test). It also supports recovery of different column families from log, which was the hardest part to reason about. We need to make sure to never delete the log file which has unflushed data from any column family. To support that, I added another concept, which is versions_->MinLogNumber()

Test Plan: Added a unit test in column_family_test

Reviewers: dhruba, haobo, sdong, kailiu

CC: leveldb

Differential Revision: https://reviews.facebook.net/D15537

Author: igorcanadi

db/column_family.cc

@@ -69,7 +69,8 @@ ColumnFamilyData::ColumnFamilyData(uint32_t id, const std::string& name,
       options(options),
       mem(nullptr),
       imm(options.min_write_buffer_number_to_merge),
-      super_version(nullptr) {}
+      super_version(nullptr),
+      log_number(0) {}
 
 ColumnFamilyData::~ColumnFamilyData() {
   if (super_version != nullptr) {
@@ -167,4 +168,18 @@ void ColumnFamilySet::DropColumnFamily(uint32_t id) {
   column_family_data_.erase(cfd);
 }
 
+MemTable* ColumnFamilyMemTablesImpl::GetMemTable(uint32_t column_family_id) {
+  auto cfd = column_family_set_->GetColumnFamily(column_family_id);
+  // TODO(icanadi): this should not be asserting. Rather, it should somehow
+  // return Corruption status back to the Iterator. This will require
+  // API change in WriteBatch::Handler, which is a public API
+  assert(cfd != nullptr);
+
+  if (log_number_ == 0 || log_number_ >= cfd->log_number) {
+    return cfd->mem;
+  } else {
+    return nullptr;
+  }
+}
+
 }  // namespace rocksdb

db/column_family.h

@@ -15,6 +15,7 @@
 
 #include "rocksdb/options.h"
 #include "db/memtablelist.h"
+#include "db/write_batch_internal.h"
 
 namespace rocksdb {
 
@@ -63,6 +64,11 @@ struct ColumnFamilyData {
   MemTableList imm;
   SuperVersion* super_version;
 
+  // This is the earliest log file number that contains data from this
+  // Column Family. All earlier log files must be ignored and not
+  // recovered from
+  uint64_t log_number;
+
   ColumnFamilyData(uint32_t id, const std::string& name,
                    Version* dummy_versions, const ColumnFamilyOptions& options);
   ~ColumnFamilyData();
@@ -122,4 +128,24 @@ class ColumnFamilySet {
   uint32_t max_column_family_;
 };
 
+class ColumnFamilyMemTablesImpl : public ColumnFamilyMemTables {
+ public:
+  explicit ColumnFamilyMemTablesImpl(ColumnFamilySet* column_family_set)
+      : column_family_set_(column_family_set), log_number_(0) {}
+
+  // If column_family_data->log_number is bigger than log_number,
+  // the memtable will not be returned.
+  // If log_number == 0, the memtable will be always returned
+  void SetLogNumber(uint64_t log_number) { log_number_ = log_number; }
+
+  // Returns the column families memtable if log_number == 0 || log_number <=
+  // column_family_data->log_number.
+  // If column family doesn't exist, it asserts
+  virtual MemTable* GetMemTable(uint32_t column_family_id) override;
+
+ private:
+  ColumnFamilySet* column_family_set_;
+  uint64_t log_number_;
+};
+
 }  // namespace rocksdb

db/column_family_test.cc

@@ -8,8 +8,10 @@
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
 
 #include "db/db_impl.h"
+#include "rocksdb/env.h"
 #include "rocksdb/db.h"
 #include "util/testharness.h"
+#include "utilities/merge_operators.h"
 
 #include <algorithm>
 #include <vector>
@@ -22,10 +24,10 @@ using namespace std;
 class ColumnFamilyTest {
  public:
   ColumnFamilyTest() {
+    env_ = Env::Default();
     dbname_ = test::TmpDir() + "/column_family_test";
     db_options_.create_if_missing = true;
-    options_.create_if_missing = true;
-    DestroyDB(dbname_, options_);
+    DestroyDB(dbname_, Options(db_options_, column_family_options_));
   }
 
   void Close() {
@@ -37,18 +39,77 @@ class ColumnFamilyTest {
     vector<ColumnFamilyDescriptor> column_families;
     for (auto x : cf) {
       column_families.push_back(
-          ColumnFamilyDescriptor(x, ColumnFamilyOptions()));
+          ColumnFamilyDescriptor(x, column_family_options_));
     }
-    vector <ColumnFamilyHandle> handles;
     return DB::OpenWithColumnFamilies(db_options_, dbname_, column_families,
-                                      &handles, &db_);
+                                      &handles_, &db_);
   }
 
-  Options options_;
+  void Destroy() {
+    delete db_;
+    db_ = nullptr;
+    ASSERT_OK(DestroyDB(dbname_, Options(db_options_, column_family_options_)));
+  }
+
+  void CreateColumnFamilies(const vector<string>& cfs) {
+    int cfi = handles_.size();
+    handles_.resize(cfi + cfs.size());
+    for (auto cf : cfs) {
+      ASSERT_OK(db_->CreateColumnFamily(column_family_options_, cf,
+                                        &handles_[cfi++]));
+    }
+  }
+
+  Status Put(int cf, const string& key, const string& value) {
+    return db_->Put(WriteOptions(), handles_[cf], Slice(key), Slice(value));
+  }
+  Status Merge(int cf, const string& key, const string& value) {
+    return db_->Merge(WriteOptions(), handles_[cf], Slice(key), Slice(value));
+  }
+
+  string Get(int cf, const string& key) {
+    ReadOptions options;
+    options.verify_checksums = true;
+    string result;
+    Status s = db_->Get(options, handles_[cf], Slice(key), &result);
+    if (s.IsNotFound()) {
+      result = "NOT_FOUND";
+    } else if (!s.ok()) {
+      result = s.ToString();
+    }
+    return result;
+  }
+
+  void CopyFile(const string& source, const string& destination,
+                uint64_t size = 0) {
+    const EnvOptions soptions;
+    unique_ptr<SequentialFile> srcfile;
+    ASSERT_OK(env_->NewSequentialFile(source, &srcfile, soptions));
+    unique_ptr<WritableFile> destfile;
+    ASSERT_OK(env_->NewWritableFile(destination, &destfile, soptions));
+
+    if (size == 0) {
+      // default argument means copy everything
+      ASSERT_OK(env_->GetFileSize(source, &size));
+    }
+
+    char buffer[4096];
+    Slice slice;
+    while (size > 0) {
+      uint64_t one = min(uint64_t(sizeof(buffer)), size);
+      ASSERT_OK(srcfile->Read(one, &slice, buffer));
+      ASSERT_OK(destfile->Append(slice));
+      size -= slice.size();
+    }
+    ASSERT_OK(destfile->Close());
+  }
+
+  vector<ColumnFamilyHandle> handles_;
   ColumnFamilyOptions column_family_options_;
   DBOptions db_options_;
   string dbname_;
   DB* db_;
+  Env* env_;
 };
 
 TEST(ColumnFamilyTest, AddDrop) {
@@ -74,6 +135,108 @@ TEST(ColumnFamilyTest, AddDrop) {
   ASSERT_TRUE(families == vector<string>({"default", "four", "one", "three"}));
 }
 
+TEST(ColumnFamilyTest, ReadWrite) {
+  ASSERT_OK(Open({"default"}));
+  CreateColumnFamilies({"one", "two"});
+  Close();
+  ASSERT_OK(Open({"default", "one", "two"}));
+  ASSERT_OK(Put(0, "foo", "v1"));
+  ASSERT_OK(Put(0, "bar", "v2"));
+  ASSERT_OK(Put(1, "mirko", "v3"));
+  ASSERT_OK(Put(0, "foo", "v2"));
+  ASSERT_OK(Put(2, "fodor", "v5"));
+
+  for (int iter = 0; iter <= 3; ++iter) {
+    ASSERT_EQ("v2", Get(0, "foo"));
+    ASSERT_EQ("v2", Get(0, "bar"));
+    ASSERT_EQ("v3", Get(1, "mirko"));
+    ASSERT_EQ("v5", Get(2, "fodor"));
+    ASSERT_EQ("NOT_FOUND", Get(0, "fodor"));
+    ASSERT_EQ("NOT_FOUND", Get(1, "fodor"));
+    ASSERT_EQ("NOT_FOUND", Get(2, "foo"));
+    if (iter <= 1) {
+      // reopen
+      Close();
+      ASSERT_OK(Open({"default", "one", "two"}));
+    }
+  }
+  Close();
+}
+
+TEST(ColumnFamilyTest, IgnoreRecoveredLog) {
+  string backup_logs = dbname_ + "/backup_logs";
+
+  // delete old files in backup_logs directory
+  env_->CreateDirIfMissing(backup_logs);
+  vector<string> old_files;
+  env_->GetChildren(backup_logs, &old_files);
+  for (auto& file : old_files) {
+    if (file != "." && file != "..") {
+      env_->DeleteFile(backup_logs + "/" + file);
+    }
+  }
+
+  column_family_options_.merge_operator =
+      MergeOperators::CreateUInt64AddOperator();
+  db_options_.wal_dir = dbname_ + "/logs";
+  Destroy();
+  ASSERT_OK(Open({"default"}));
+  CreateColumnFamilies({"cf1", "cf2"});
+
+  // fill up the DB
+  string one, two, three;
+  PutFixed64(&one, 1);
+  PutFixed64(&two, 2);
+  PutFixed64(&three, 3);
+  ASSERT_OK(Merge(0, "foo", one));
+  ASSERT_OK(Merge(1, "mirko", one));
+  ASSERT_OK(Merge(0, "foo", one));
+  ASSERT_OK(Merge(2, "bla", one));
+  ASSERT_OK(Merge(2, "fodor", one));
+  ASSERT_OK(Merge(0, "bar", one));
+  ASSERT_OK(Merge(2, "bla", one));
+  ASSERT_OK(Merge(1, "mirko", two));
+  ASSERT_OK(Merge(1, "franjo", one));
+
+  // copy the logs to backup
+  vector<string> logs;
+  env_->GetChildren(db_options_.wal_dir, &logs);
+  for (auto& log : logs) {
+    if (log != ".." && log != ".") {
+      CopyFile(db_options_.wal_dir + "/" + log, backup_logs + "/" + log);
+    }
+  }
+
+  // recover the DB
+  Close();
+
+  // 1. check consistency
+  // 2. copy the logs from backup back to WAL dir. if the recovery happens
+  // again on the same log files, this should lead to incorrect results
+  // due to applying merge operator twice
+  // 3. check consistency
+  for (int iter = 0; iter < 2; ++iter) {
+    // assert consistency
+    ASSERT_OK(Open({"default", "cf1", "cf2"}));
+    ASSERT_EQ(two, Get(0, "foo"));
+    ASSERT_EQ(one, Get(0, "bar"));
+    ASSERT_EQ(three, Get(1, "mirko"));
+    ASSERT_EQ(one, Get(1, "franjo"));
+    ASSERT_EQ(one, Get(2, "fodor"));
+    ASSERT_EQ(two, Get(2, "bla"));
+    Close();
+
+    if (iter == 0) {
+      // copy the logs from backup back to wal dir
+      for (auto& log : logs) {
+        if (log != ".." && log != ".") {
+          CopyFile(backup_logs + "/" + log, db_options_.wal_dir + "/" + log);
+        }
+      }
+    }
+  }
+}
+
 }  // namespace rocksdb
 
 int main(int argc, char** argv) {