@nxtedition/rocksdb 8.2.8 → 9.0.1

package/deps/rocksdb/rocksdb/db/memtable_list.cc

@@ -184,7 +184,7 @@ bool MemTableListVersion::GetFromList(
       assert(*seq != kMaxSequenceNumber || s->IsNotFound());
       return true;
     }
-    if (!done && !s->ok() && !s->IsMergeInProgress() && !s->IsNotFound()) {
+    if (!s->ok() && !s->IsMergeInProgress() && !s->IsNotFound()) {
       return false;
     }
   }
@@ -434,23 +434,57 @@ void MemTableList::PickMemtablesToFlush(uint64_t max_memtable_id,
 }
 
 void MemTableList::RollbackMemtableFlush(const autovector<MemTable*>& mems,
-                                         uint64_t /*file_number*/) {
+                                         bool rollback_succeeding_memtables) {
+  TEST_SYNC_POINT("RollbackMemtableFlush");
   AutoThreadOperationStageUpdater stage_updater(
       ThreadStatus::STAGE_MEMTABLE_ROLLBACK);
-  assert(!mems.empty());
-
-  // If the flush was not successful, then just reset state.
-  // Maybe a succeeding attempt to flush will be successful.
+#ifndef NDEBUG
   for (MemTable* m : mems) {
     assert(m->flush_in_progress_);
     assert(m->file_number_ == 0);
+  }
+#endif
+
+  if (rollback_succeeding_memtables && !mems.empty()) {
+    std::list<MemTable*>& memlist = current_->memlist_;
+    auto it = memlist.rbegin();
+    for (; *it != mems[0] && it != memlist.rend(); ++it) {
+    }
+    // mems should be in memlist
+    assert(*it == mems[0]);
+    if (*it == mems[0]) {
+      ++it;
+    }
+    while (it != memlist.rend()) {
+      MemTable* m = *it;
+      // Only rollback complete, not in-progress,
+      // in_progress can be flushes that are still writing SSTs
+      if (m->flush_completed_) {
+        m->flush_in_progress_ = false;
+        m->flush_completed_ = false;
+        m->edit_.Clear();
+        m->file_number_ = 0;
+        num_flush_not_started_++;
+        ++it;
+      } else {
+        break;
+      }
+    }
+  }
 
-    m->flush_in_progress_ = false;
-    m->flush_completed_ = false;
-    m->edit_.Clear();
-    num_flush_not_started_++;
+  for (MemTable* m : mems) {
+    if (m->flush_in_progress_) {
+      assert(m->file_number_ == 0);
+      m->file_number_ = 0;
+      m->flush_in_progress_ = false;
+      m->flush_completed_ = false;
+      m->edit_.Clear();
+      num_flush_not_started_++;
+    }
+  }
+  if (!mems.empty()) {
+    imm_flush_needed.store(true, std::memory_order_release);
   }
-  imm_flush_needed.store(true, std::memory_order_release);
 }
 
 // Try record a successful flush in the manifest file. It might just return

package/deps/rocksdb/rocksdb/db/memtable_list.h

@@ -271,8 +271,20 @@ class MemTableList {
 
   // Reset status of the given memtable list back to pending state so that
   // they can get picked up again on the next round of flush.
+  //
+  // @param rollback_succeeding_memtables If true, will rollback adjacent
+  // younger memtables whose flush is completed. Specifically, suppose the
+  // current immutable memtables are M_0,M_1...M_N ordered from youngest to
+  // oldest. Suppose that the youngest memtable in `mems` is M_K. We will try to
+  // rollback M_K-1, M_K-2... until the first memtable whose flush is
+  // not completed. These are the memtables that would have been installed
+  // by this flush job if it were to succeed. This flag is currently used
+  // by non atomic_flush rollback.
+  // Note that we also do rollback in `write_manifest_cb` by calling
+  // `RemoveMemTablesOrRestoreFlags()`. There we rollback the entire batch so
+  // it is similar to what we do here with rollback_succeeding_memtables=true.
   void RollbackMemtableFlush(const autovector<MemTable*>& mems,
-                             uint64_t file_number);
+                             bool rollback_succeeding_memtables);
 
   // Try commit a successful flush in the manifest file. It might just return
   // Status::OK letting a concurrent flush to do the actual the recording.
@@ -374,14 +386,43 @@ class MemTableList {
     return memlist.back()->GetID();
   }
 
-  uint64_t GetLatestMemTableID() const {
+  uint64_t GetLatestMemTableID(bool for_atomic_flush) const {
     auto& memlist = current_->memlist_;
     if (memlist.empty()) {
       return 0;
     }
+    if (for_atomic_flush) {
+      // Scan the memtable list from new to old
+      for (auto it = memlist.begin(); it != memlist.end(); ++it) {
+        MemTable* m = *it;
+        if (m->atomic_flush_seqno_ != kMaxSequenceNumber) {
+          return m->GetID();
+        }
+      }
+      return 0;
+    }
     return memlist.front()->GetID();
   }
 
+  // DB mutex held.
+  // Gets the newest user-defined timestamp for the Memtables in ascending ID
+  // order, up to the `max_memtable_id`. Used by background flush job
+  // to check Memtables' eligibility for flush w.r.t retaining UDTs.
+  std::vector<Slice> GetTablesNewestUDT(uint64_t max_memtable_id) {
+    std::vector<Slice> newest_udts;
+    auto& memlist = current_->memlist_;
+    // Iterating through the memlist starting at the end, the vector<MemTable*>
+    // ret is filled with memtables already sorted in increasing MemTable ID.
+    for (auto it = memlist.rbegin(); it != memlist.rend(); ++it) {
+      MemTable* m = *it;
+      if (m->GetID() > max_memtable_id) {
+        break;
+      }
+      newest_udts.push_back(m->GetNewestUDT());
+    }
+    return newest_udts;
+  }
+
   void AssignAtomicFlushSeq(const SequenceNumber& seq) {
     const auto& memlist = current_->memlist_;
     // Scan the memtable list from new to old

package/deps/rocksdb/rocksdb/db/memtable_list_test.cc

@@ -43,6 +43,9 @@ class MemTableListTest : public testing::Test {
       // Open DB only with default column family
       ColumnFamilyOptions cf_options;
       std::vector<ColumnFamilyDescriptor> cf_descs;
+      if (udt_enabled_) {
+        cf_options.comparator = test::BytewiseComparatorWithU64TsWrapper();
+      }
       cf_descs.emplace_back(kDefaultColumnFamilyName, cf_options);
       Status s = DB::Open(options, dbname, cf_descs, &handles, &db);
       EXPECT_OK(s);
@@ -200,6 +203,9 @@ class MemTableListTest : public testing::Test {
         nullptr /* prep_tracker */, &mutex, file_meta_ptrs,
         committed_flush_jobs_info, to_delete, nullptr, &log_buffer);
   }
+
+ protected:
+  bool udt_enabled_ = false;
 };
 
 TEST_F(MemTableListTest, Empty) {
@@ -676,7 +682,7 @@ TEST_F(MemTableListTest, FlushPendingTest) {
   ASSERT_FALSE(list.imm_flush_needed.load(std::memory_order_acquire));
 
   // Revert flush
-  list.RollbackMemtableFlush(to_flush, 0);
+  list.RollbackMemtableFlush(to_flush, false);
   ASSERT_FALSE(list.IsFlushPending());
   ASSERT_TRUE(list.imm_flush_needed.load(std::memory_order_acquire));
   to_flush.clear();
@@ -726,7 +732,7 @@ TEST_F(MemTableListTest, FlushPendingTest) {
   ASSERT_FALSE(list.imm_flush_needed.load(std::memory_order_acquire));
 
   // Rollback first pick of tables
-  list.RollbackMemtableFlush(to_flush, 0);
+  list.RollbackMemtableFlush(to_flush, false);
   ASSERT_TRUE(list.IsFlushPending());
   ASSERT_TRUE(list.imm_flush_needed.load(std::memory_order_acquire));
   to_flush.clear();
@@ -827,7 +833,7 @@ TEST_F(MemTableListTest, FlushPendingTest) {
   // Add another table
   list.Add(tables[5], &to_delete);
   ASSERT_EQ(1, list.NumNotFlushed());
-  ASSERT_EQ(5, list.GetLatestMemTableID());
+  ASSERT_EQ(5, list.GetLatestMemTableID(false /* for_atomic_flush */));
   memtable_id = 4;
   // Pick tables to flush. The tables to pick must have ID smaller than or
   // equal to 4. Therefore, no table will be selected in this case.
@@ -868,7 +874,7 @@ TEST_F(MemTableListTest, FlushPendingTest) {
   to_delete.clear();
 }
 
-TEST_F(MemTableListTest, EmptyAtomicFlusTest) {
+TEST_F(MemTableListTest, EmptyAtomicFlushTest) {
   autovector<MemTableList*> lists;
   autovector<uint32_t> cf_ids;
   autovector<const MutableCFOptions*> options_list;
@@ -880,7 +886,7 @@ TEST_F(MemTableListTest, EmptyAtomicFlusTest) {
   ASSERT_TRUE(to_delete.empty());
 }
 
-TEST_F(MemTableListTest, AtomicFlusTest) {
+TEST_F(MemTableListTest, AtomicFlushTest) {
   const int num_cfs = 3;
   const int num_tables_per_cf = 2;
   SequenceNumber seq = 1;
@@ -1028,6 +1034,86 @@ TEST_F(MemTableListTest, AtomicFlusTest) {
   }
 }
 
+class MemTableListWithTimestampTest : public MemTableListTest {
+ public:
+  MemTableListWithTimestampTest() : MemTableListTest() {}
+
+  void SetUp() override { udt_enabled_ = true; }
+};
+
+TEST_F(MemTableListWithTimestampTest, GetTableNewestUDT) {
+  const int num_tables = 3;
+  const int num_entries = 5;
+  SequenceNumber seq = 1;
+
+  auto factory = std::make_shared<SkipListFactory>();
+  options.memtable_factory = factory;
+  options.persist_user_defined_timestamps = false;
+  ImmutableOptions ioptions(options);
+  const Comparator* ucmp = test::BytewiseComparatorWithU64TsWrapper();
+  InternalKeyComparator cmp(ucmp);
+  WriteBufferManager wb(options.db_write_buffer_size);
+
+  // Create MemTableList
+  int min_write_buffer_number_to_merge = 1;
+  int max_write_buffer_number_to_maintain = 4;
+  int64_t max_write_buffer_size_to_maintain =
+      4 * static_cast<int>(options.write_buffer_size);
+  MemTableList list(min_write_buffer_number_to_merge,
+                    max_write_buffer_number_to_maintain,
+                    max_write_buffer_size_to_maintain);
+
+  // Create some MemTables
+  uint64_t memtable_id = 0;
+  std::vector<MemTable*> tables;
+  MutableCFOptions mutable_cf_options(options);
+  uint64_t current_ts = 0;
+  autovector<MemTable*> to_delete;
+  std::vector<std::string> newest_udts;
+
+  std::string key;
+  std::string write_ts;
+  for (int i = 0; i < num_tables; i++) {
+    MemTable* mem = new MemTable(cmp, ioptions, mutable_cf_options, &wb,
+                                 kMaxSequenceNumber, 0 /* column_family_id */);
+    mem->SetID(memtable_id++);
+    mem->Ref();
+
+    std::string value;
+    MergeContext merge_context;
+
+    for (int j = 0; j < num_entries; j++) {
+      key = "key1";
+      write_ts.clear();
+      PutFixed64(&write_ts, current_ts);
+      key.append(write_ts);
+      ASSERT_OK(mem->Add(++seq, kTypeValue, key, std::to_string(i),
+                         nullptr /* kv_prot_info */));
+      current_ts++;
+    }
+
+    tables.push_back(mem);
+    list.Add(tables.back(), &to_delete);
+    newest_udts.push_back(write_ts);
+  }
+
+  ASSERT_EQ(num_tables, list.NumNotFlushed());
+  ASSERT_TRUE(list.IsFlushPending());
+  std::vector<Slice> tables_newest_udts = list.GetTablesNewestUDT(num_tables);
+  ASSERT_EQ(newest_udts.size(), tables_newest_udts.size());
+  for (size_t i = 0; i < tables_newest_udts.size(); i++) {
+    const Slice& table_newest_udt = tables_newest_udts[i];
+    const Slice expected_newest_udt = newest_udts[i];
+    ASSERT_EQ(expected_newest_udt, table_newest_udt);
+  }
+
+  list.current()->Unref(&to_delete);
+  for (MemTable* m : to_delete) {
+    delete m;
+  }
+  to_delete.clear();
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {