Address the review comments

facebook · Jun 18, 2021 · b39a15d · b39a15d
1 parent 42c657a
commit b39a15d
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 4 deletions.
diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc
@@ -2201,6 +2201,16 @@ void DBImpl::MultiGet(const ReadOptions& read_options, const size_t num_keys,
   }
 #endif  // NDEBUG
 
+  if (tracer_) {
+    // TODO: This mutex should be removed later, to improve performance when
+    // tracing is enabled.
+    InstrumentedMutexLock lock(&trace_mutex_);
+    if (tracer_) {
+      // TODO: maybe handle the tracing status?
+      tracer_->MultiGet(num_keys, column_families, keys).PermitUncheckedError();
+    }
+  }
+
   autovector<KeyContext, MultiGetContext::MAX_BATCH_SIZE> key_context;
   autovector<KeyContext*, MultiGetContext::MAX_BATCH_SIZE> sorted_keys;
   sorted_keys.resize(num_keys);
@@ -2363,6 +2373,15 @@ void DBImpl::MultiGet(const ReadOptions& read_options,
                       const Slice* keys, PinnableSlice* values,
                       std::string* timestamps, Status* statuses,
                       const bool sorted_input) {
+  if (tracer_) {
+    // TODO: This mutex should be removed later, to improve performance when
+    // tracing is enabled.
+    InstrumentedMutexLock lock(&trace_mutex_);
+    if (tracer_) {
+      // TODO: maybe handle the tracing status?
+      tracer_->MultiGet(num_keys, column_family, keys).PermitUncheckedError();
+    }
+  }
   autovector<KeyContext, MultiGetContext::MAX_BATCH_SIZE> key_context;
   autovector<KeyContext*, MultiGetContext::MAX_BATCH_SIZE> sorted_keys;
   sorted_keys.resize(num_keys);

diff --git a/tools/trace_analyzer_test.cc b/tools/trace_analyzer_test.cc
@@ -92,6 +92,13 @@ class TraceAnalyzerTest : public testing::Test {
     ASSERT_GE(ss.size(), 0);
     ASSERT_OK(ss[0]);
     ASSERT_NOK(ss[2]);
+    std::vector<ColumnFamilyHandle*> cfs(2, db_->DefaultColumnFamily());
+    std::vector<PinnableSlice> values2(keys.size());
+    db_->MultiGet(ro, 2, cfs.data(), keys.data(), values2.data(), ss.data(),
+                  false);
+    ASSERT_OK(ss[0]);
+    db_->MultiGet(ro, db_->DefaultColumnFamily(), 2, keys.data() + 3,
+                  values2.data(), ss.data(), false);
     ASSERT_OK(db_->Get(ro, "a", &value));
 
     single_iter = db_->NewIterator(ro);

diff --git a/trace_replay/trace_replay.cc b/trace_replay/trace_replay.cc
@@ -343,9 +343,42 @@ Status Tracer::IteratorSeekForPrev(const uint32_t& cf_id, const Slice& key,
   return WriteTrace(trace);
 }
 
-Status Tracer::MultiGet(const std::vector<ColumnFamilyHandle*>& column_family,
+Status Tracer::MultiGet(const size_t num_keys,
+                        ColumnFamilyHandle** column_families,
+                        const Slice* keys) {
+  if (num_keys == 0) {
+    return Status::OK();
+  }
+  std::vector<ColumnFamilyHandle*> v_column_families;
+  std::vector<Slice> v_keys;
+  v_column_families.resize(num_keys);
+  v_keys.resize(num_keys);
+  for (size_t i = 0; i < num_keys; i++) {
+    v_column_families[i] = column_families[i];
+    v_keys[i] = keys[i];
+  }
+  return MultiGet(v_column_families, v_keys);
+}
+
+Status Tracer::MultiGet(const size_t num_keys,
+                        ColumnFamilyHandle* column_family, const Slice* keys) {
+  if (num_keys == 0) {
+    return Status::OK();
+  }
+  std::vector<ColumnFamilyHandle*> column_families;
+  std::vector<Slice> v_keys;
+  column_families.resize(num_keys);
+  v_keys.resize(num_keys);
+  for (size_t i = 0; i < num_keys; i++) {
+    column_families[i] = column_family;
+    v_keys[i] = keys[i];
+  }
+  return MultiGet(column_families, v_keys);
+}
+
+Status Tracer::MultiGet(const std::vector<ColumnFamilyHandle*>& column_families,
                         const std::vector<Slice>& keys) {
-  if (column_family.size() != keys.size()) {
+  if (column_families.size() != keys.size()) {
     return Status::Corruption("the CFs size and keys size does not match!");
   }
   TraceType trace_type = kTraceMultiGet;
@@ -368,9 +401,9 @@ Status Tracer::MultiGet(const std::vector<ColumnFamilyHandle*>& column_family,
   std::string cfids_payload;
   std::string keys_payload;
   for (uint32_t i = 0; i < multiget_size; i++) {
-    assert(i < column_family.size());
+    assert(i < column_families.size());
     assert(i < keys.size());
-    PutFixed32(&cfids_payload, column_family[i]->GetID());
+    PutFixed32(&cfids_payload, column_families[i]->GetID());
     PutLengthPrefixedSlice(&keys_payload, keys[i]);
   }
   // Encode the Get struct members into payload. Make sure add them in order.

diff --git a/trace_replay/trace_replay.h b/trace_replay/trace_replay.h
@@ -182,6 +182,13 @@ class Tracer {
                              const Slice& lower_bound, const Slice upper_bound);
 
   // Trace MultiGet
+
+  Status MultiGet(const size_t num_keys, ColumnFamilyHandle** column_families,
+                  const Slice* keys);
+
+  Status MultiGet(const size_t num_keys, ColumnFamilyHandle* column_family,
+                  const Slice* keys);
+
   Status MultiGet(const std::vector<ColumnFamilyHandle*>& column_family,
                   const std::vector<Slice>& keys);