Fix flaky db_basic_bench caused by unreleased iterator (#10058)

Summary: Iterator is not freed after test is done (after the main for loop), which could cause db close failed. Pull Request resolved: facebook/rocksdb#10058 Test Plan: Able to reproduce consistently with higher thread number, like 100, make sure it passes after the fix Reviewed By: ajkr Differential Revision: D36685823 Pulled By: jay-zhuang fbshipit-source-id: 4c98b8758d106bfe40cae670e689c3d284765bcf
asu-idi · May 26, 2022 · 5490da2 · 5490da2
1 parent bd170dd
commit 5490da2
Showing 1 changed file with 61 additions and 73 deletions.
diff --git a/microbench/db_basic_bench.cc b/microbench/db_basic_bench.cc
@@ -1109,21 +1109,19 @@ static void IteratorSeek(benchmark::State& state) {
     }
   }
 
-  {
+  for (auto _ : state) {
     std::unique_ptr<Iterator> iter{nullptr};
-    for (auto _ : state) {
-      state.PauseTiming();
-      if (!iter) {
-        iter.reset(db->NewIterator(ReadOptions()));
-      }
-      Slice key = negative_query ? kg.NextNonExist() : kg.Next();
-      if (!iter->status().ok()) {
-        state.SkipWithError(iter->status().ToString().c_str());
-        return;
-      }
-      state.ResumeTiming();
-      iter->Seek(key);
+    state.PauseTiming();
+    if (!iter) {
+      iter.reset(db->NewIterator(ReadOptions()));
+    }
+    Slice key = negative_query ? kg.NextNonExist() : kg.Next();
+    if (!iter->status().ok()) {
+      state.SkipWithError(iter->status().ToString().c_str());
+      return;
     }
+    state.ResumeTiming();
+    iter->Seek(key);
   }
 
   if (state.thread_index() == 0) {
@@ -1202,22 +1200,20 @@ static void IteratorNext(benchmark::State& state) {
     }
   }
 
-  {
+  for (auto _ : state) {
     std::unique_ptr<Iterator> iter{nullptr};
-    for (auto _ : state) {
-      state.PauseTiming();
-      if (!iter) {
-        iter.reset(db->NewIterator(ReadOptions()));
-      }
-      while (!iter->Valid()) {
-        iter->Seek(kg.Next());
-        if (!iter->status().ok()) {
-          state.SkipWithError(iter->status().ToString().c_str());
-        }
+    state.PauseTiming();
+    if (!iter) {
+      iter.reset(db->NewIterator(ReadOptions()));
+    }
+    while (!iter->Valid()) {
+      iter->Seek(kg.Next());
+      if (!iter->status().ok()) {
+        state.SkipWithError(iter->status().ToString().c_str());
       }
-      state.ResumeTiming();
-      iter->Next();
     }
+    state.ResumeTiming();
+    iter->Next();
   }
 
   if (state.thread_index() == 0) {
@@ -1281,31 +1277,27 @@ static void IteratorNextWithPerfContext(benchmark::State& state) {
   SetPerfLevel(kEnableTime);
   get_perf_context()->EnablePerLevelPerfContext();
 
-  {
+  for (auto _ : state) {
     std::unique_ptr<Iterator> iter{nullptr};
-    for (auto _ : state) {
-      state.PauseTiming();
-      if (!iter) {
-        iter.reset(db->NewIterator(ReadOptions()));
-      }
-      while (!iter->Valid()) {
-        iter->Seek(kg.Next());
-        if (!iter->status().ok()) {
-          state.SkipWithError(iter->status().ToString().c_str());
-        }
+    state.PauseTiming();
+    if (!iter) {
+      iter.reset(db->NewIterator(ReadOptions()));
+    }
+    while (!iter->Valid()) {
+      iter->Seek(kg.Next());
+      if (!iter->status().ok()) {
+        state.SkipWithError(iter->status().ToString().c_str());
       }
-      get_perf_context()->Reset();
-      state.ResumeTiming();
-
-      iter->Next();
-      user_key_comparison_count +=
-          get_perf_context()->user_key_comparison_count;
-      internal_key_skipped_count +=
-          get_perf_context()->internal_key_skipped_count;
-      find_next_user_entry_time +=
-          get_perf_context()->find_next_user_entry_time;
-      iter_next_cpu_nanos += get_perf_context()->iter_next_cpu_nanos;
     }
+    get_perf_context()->Reset();
+    state.ResumeTiming();
+
+    iter->Next();
+    user_key_comparison_count += get_perf_context()->user_key_comparison_count;
+    internal_key_skipped_count +=
+        get_perf_context()->internal_key_skipped_count;
+    find_next_user_entry_time += get_perf_context()->find_next_user_entry_time;
+    iter_next_cpu_nanos += get_perf_context()->iter_next_cpu_nanos;
   }
 
   state.counters["user_key_comparison_count"] =
@@ -1370,22 +1362,20 @@ static void IteratorPrev(benchmark::State& state) {
     }
   }
 
-  {
+  for (auto _ : state) {
     std::unique_ptr<Iterator> iter{nullptr};
-    for (auto _ : state) {
-      state.PauseTiming();
-      if (!iter) {
-        iter.reset(db->NewIterator(ReadOptions()));
-      }
-      while (!iter->Valid()) {
-        iter->Seek(kg.Next());
-        if (!iter->status().ok()) {
-          state.SkipWithError(iter->status().ToString().c_str());
-        }
+    state.PauseTiming();
+    if (!iter) {
+      iter.reset(db->NewIterator(ReadOptions()));
+    }
+    while (!iter->Valid()) {
+      iter->Seek(kg.Next());
+      if (!iter->status().ok()) {
+        state.SkipWithError(iter->status().ToString().c_str());
       }
-      state.ResumeTiming();
-      iter->Prev();
     }
+    state.ResumeTiming();
+    iter->Prev();
   }
 
   if (state.thread_index() == 0) {
@@ -1464,19 +1454,17 @@ static void PrefixSeek(benchmark::State& state) {
     }
   }
 
-  {
+  for (auto _ : state) {
     std::unique_ptr<Iterator> iter{nullptr};
-    for (auto _ : state) {
-      state.PauseTiming();
-      if (!iter) {
-        iter.reset(db->NewIterator(ReadOptions()));
-      }
-      state.ResumeTiming();
-      iter->Seek(kg.NextPrefix());
-      if (!iter->status().ok()) {
-        state.SkipWithError(iter->status().ToString().c_str());
-        return;
-      }
+    state.PauseTiming();
+    if (!iter) {
+      iter.reset(db->NewIterator(ReadOptions()));
+    }
+    state.ResumeTiming();
+    iter->Seek(kg.NextPrefix());
+    if (!iter->status().ok()) {
+      state.SkipWithError(iter->status().ToString().c_str());
+      return;
     }
   }