Fix wrong live data size when encounter rewrite failure (#161) (#175)

* fix wrong live data size when encounter rewrite failure
Signed-off-by: Connor1996 <zbk602423539@gmail.com>

src/blob_gc_job.cc

@@ -347,6 +347,7 @@ Status BlobGCJob::Finish() {
     mutex_->Unlock();
     s = InstallOutputBlobFiles();
     if (s.ok()) {
+      TEST_SYNC_POINT("BlobGCJob::Finish::BeforeRewriteValidKeyToLSM");
       s = RewriteValidKeyToLSM();
       if (!s.ok()) {
         ROCKS_LOG_ERROR(db_options_.info_log,
@@ -368,6 +369,7 @@ Status BlobGCJob::Finish() {
   if (s.ok() && !blob_gc_->GetColumnFamilyData()->IsDropped()) {
     s = DeleteInputBlobFiles();
   }
+  TEST_SYNC_POINT("BlobGCJob::Finish::AfterRewriteValidKeyToLSM");
 
   if (s.ok()) {
     UpdateInternalOpStats();
@@ -451,6 +453,9 @@ Status BlobGCJob::RewriteValidKeyToLSM() {
   WriteOptions wo;
   wo.low_pri = true;
   wo.ignore_missing_column_families = true;
+
+  std::unordered_map<uint64_t, uint64_t>
+      dropped;  // blob_file_number -> dropped_size
   if (!gc_merge_rewrite_) {
     for (auto& write_batch : rewrite_batches_) {
       if (blob_gc_->GetColumnFamilyData()->IsDropped()) {
@@ -473,6 +478,13 @@ Status BlobGCJob::RewriteValidKeyToLSM() {
         metrics_.gc_num_keys_overwritten++;
         metrics_.gc_bytes_overwritten += write_batch.second.blob_record_size();
         // The key is overwritten in the meanwhile. Drop the blob record.
+        // Though record is dropped, the diff won't counted in discardable
+        // ratio,
+        // so we should update the live_data_size here.
+        BlobIndex blob_index;
+        Slice str(write_batch.second.value);
+        blob_index.DecodeFrom(&str);
+        dropped[blob_index.file_number] += blob_index.blob_handle.size;
       } else {
         // We hit an error.
         break;
@@ -508,6 +520,30 @@ Status BlobGCJob::RewriteValidKeyToLSM() {
     s = Status::OK();
   }
 
+  mutex_->Lock();
+  auto cf_id = blob_gc_->column_family_handle()->GetID();
+  for (auto blob_file : dropped) {
+    auto blob_storage = blob_file_set_->GetBlobStorage(cf_id).lock();
+    if (blob_storage) {
+      auto file = blob_storage->FindFile(blob_file.first).lock();
+      if (!file) {
+        ROCKS_LOG_ERROR(db_options_.info_log,
+                        "Blob File %" PRIu64 " not found when GC.",
+                        blob_file.first);
+        continue;
+      }
+      SubStats(stats_, cf_id, file->GetDiscardableRatioLevel(), 1);
+      file->UpdateLiveDataSize(-blob_file.second);
+      AddStats(stats_, cf_id, file->GetDiscardableRatioLevel(), 1);
+
+      blob_storage->ComputeGCScore();
+    } else {
+      ROCKS_LOG_ERROR(db_options_.info_log,
+                      "Column family id:%" PRIu32 " not Found when GC.", cf_id);
+    }
+  }
+  mutex_->Unlock();
+
   if (s.ok()) {
     // Flush and sync WAL.
     s = db_impl->FlushWAL(true /*sync*/);

src/db_impl_gc.cc

@@ -212,6 +212,7 @@ Status TitanDBImpl::BackgroundGC(LogBuffer* log_buffer,
     if (s.ok()) {
       mutex_.Unlock();
       s = blob_gc_job.Run();
+      TEST_SYNC_POINT("TitanDBImpl::BackgroundGC::AfterRunGCJob");
       mutex_.Lock();
     }
     if (s.ok()) {

src/titan_db_test.cc

@@ -155,6 +155,17 @@ class TitanDBTest : public testing::Test {
     return db_impl_->blob_file_set_->GetBlobStorage(cf_handle->GetID());
   }
 
+  void CheckBlobFileCount(int count, ColumnFamilyHandle* cf_handle = nullptr) {
+    db_impl_->TEST_WaitForBackgroundGC();
+    ASSERT_OK(db_impl_->TEST_PurgeObsoleteFiles());
+    std::shared_ptr<BlobStorage> blob_storage =
+        GetBlobStorage(cf_handle).lock();
+    ASSERT_TRUE(blob_storage != nullptr);
+    std::map<uint64_t, std::weak_ptr<BlobFileMeta>> blob_files;
+    blob_storage->ExportBlobFiles(blob_files);
+    ASSERT_EQ(count, blob_files.size());
+  }
+
   ColumnFamilyHandle* GetColumnFamilyHandle(uint32_t cf_id) {
     return db_impl_->db_impl_->GetColumnFamilyHandleUnlocked(cf_id).release();
   }
@@ -220,13 +231,17 @@ class TitanDBTest : public testing::Test {
     }
   }
 
-  void CompactAll() {
+  void CompactAll(ColumnFamilyHandle* cf_handle = nullptr) {
+    if (cf_handle == nullptr) {
+      cf_handle = db_->DefaultColumnFamily();
+    }
     auto opts = db_->GetOptions();
     auto compact_opts = CompactRangeOptions();
     compact_opts.change_level = true;
     compact_opts.target_level = opts.num_levels - 1;
-    compact_opts.bottommost_level_compaction = BottommostLevelCompaction::kSkip;
-    ASSERT_OK(db_->CompactRange(compact_opts, nullptr, nullptr));
+    compact_opts.bottommost_level_compaction =
+        BottommostLevelCompaction::kForce;
+    ASSERT_OK(db_->CompactRange(compact_opts, cf_handle, nullptr, nullptr));
   }
 
   void DeleteFilesInRange(const Slice* begin, const Slice* end) {
@@ -1386,6 +1401,66 @@ TEST_F(TitanDBTest, GCAfterReopen) {
   ASSERT_GT(file2->file_number(), file_number1);
 }
 
+TEST_F(TitanDBTest, CompactionDuringGC) {
+  options_.max_background_gc = 1;
+  options_.disable_background_gc = false;
+  options_.blob_file_discardable_ratio = 0.01;
+  Open();
+
+  ASSERT_OK(db_->Put(WriteOptions(), "k1", std::string(10 * 1024, 'v')));
+  auto snap = db_->GetSnapshot();
+  ASSERT_OK(db_->Put(WriteOptions(), "k1", std::string(100 * 1024, 'v')));
+  Flush();
+
+  db_->ReleaseSnapshot(snap);
+
+  SyncPoint::GetInstance()->LoadDependency(
+      {{"TitanDBImpl::BackgroundGC::AfterRunGCJob",
+        "TitanDBTest::CompactionDuringGC::WaitGCStart"},
+       {"TitanDBTest::CompactionDuringGC::ContinueGC",
+        "BlobGCJob::Finish::BeforeRewriteValidKeyToLSM"},
+       {"BlobGCJob::Finish::AfterRewriteValidKeyToLSM",
+        "TitanDBTest::CompactionDuringGC::WaitGCFinish"}});
+  SyncPoint::GetInstance()->EnableProcessing();
+
+  CheckBlobFileCount(1);
+  CompactAll();
+  std::shared_ptr<BlobStorage> blob_storage = GetBlobStorage().lock();
+  ASSERT_TRUE(blob_storage != nullptr);
+  std::map<uint64_t, std::weak_ptr<BlobFileMeta>> blob_files;
+  blob_storage->ExportBlobFiles(blob_files);
+  ASSERT_EQ(blob_files.size(), 1);
+
+  // trigger GC
+  CompactAll();
+
+  TEST_SYNC_POINT("TitanDBTest::CompactionDuringGC::WaitGCStart");
+
+  ASSERT_OK(db_->Delete(WriteOptions(), "k1"));
+
+  TEST_SYNC_POINT("TitanDBTest::CompactionDuringGC::ContinueGC");
+  TEST_SYNC_POINT("TitanDBTest::CompactionDuringGC::WaitGCFinish");
+
+  blob_storage->ExportBlobFiles(blob_files);
+  // rewriting index to LSM failed, but the output blob file is already
+  // generated
+  ASSERT_EQ(blob_files.size(), 2);
+
+  std::string value;
+  Status status = db_->Get(ReadOptions(), "k1", &value);
+  ASSERT_EQ(status, Status::NotFound());
+
+  SyncPoint::GetInstance()->DisableProcessing();
+  CheckBlobFileCount(1);
+
+  Flush();
+  CompactAll();
+  CompactAll();
+
+  db_impl_->TEST_StartGC(db_impl_->DefaultColumnFamily()->GetID());
+  CheckBlobFileCount(0);
+}
+
 }  // namespace titandb
 }  // namespace rocksdb