Integrate block cache tracer into db_impl (#5433)

Summary:
This PR integrates the block cache tracer class into db_impl.cc.
db_impl.cc contains a member variable of AtomicBlockCacheTraceWriter class and passes its reference to the block_based_table_reader.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/5433

Differential Revision: D15728016

Pulled By: HaoyuHuang

fbshipit-source-id: 23d5659e8c82d556833dcc1a5558aac8c1f7db71

1 files changed, 50 insertions, 16 deletions

diff --git a/trace_replay/block_cache_tracer.cc b/trace_replay/block_cache_tracer.cc
index 58c7df70b..565511e5a 100644
--- a/trace_replay/block_cache_tracer.cc
+++ b/trace_replay/block_cache_tracer.cc
@@ -23,30 +23,29 @@ bool ShouldTraceReferencedKey(const BlockCacheTraceRecord& record) {
           record.caller == BlockCacheLookupCaller::kUserMGet);
 }
 
-BlockCacheTraceWriter::BlockCacheTraceWriter(
-    Env* env, const TraceOptions& trace_options,
-    std::unique_ptr<TraceWriter>&& trace_writer)
-    : env_(env),
-      trace_options_(trace_options),
-      trace_writer_(std::move(trace_writer)) {}
-
-bool BlockCacheTraceWriter::ShouldTrace(
-    const BlockCacheTraceRecord& record) const {
-  if (trace_options_.sampling_frequency == 0 ||
-      trace_options_.sampling_frequency == 1) {
+bool ShouldTrace(const BlockCacheTraceRecord& record,
+                 const TraceOptions& trace_options) {
+  if (trace_options.sampling_frequency == 0 ||
+      trace_options.sampling_frequency == 1) {
     return true;
   }
   // We use spatial downsampling so that we have a complete access history for a
   // block.
   const uint64_t hash = GetSliceNPHash64(Slice(record.block_key));
-  return hash % trace_options_.sampling_frequency == 0;
+  return hash % trace_options.sampling_frequency == 0;
 }
 
+BlockCacheTraceWriter::BlockCacheTraceWriter(
+    Env* env, const TraceOptions& trace_options,
+    std::unique_ptr<TraceWriter>&& trace_writer)
+    : env_(env),
+      trace_options_(trace_options),
+      trace_writer_(std::move(trace_writer)) {}
+
 Status BlockCacheTraceWriter::WriteBlockAccess(
     const BlockCacheTraceRecord& record) {
   uint64_t trace_file_size = trace_writer_->GetFileSize();
-  if (trace_file_size > trace_options_.max_trace_file_size ||
-      !ShouldTrace(record)) {
+  if (trace_file_size > trace_options_.max_trace_file_size) {
     return Status::OK();
   }
   Trace trace;
@@ -68,7 +67,6 @@ Status BlockCacheTraceWriter::WriteBlockAccess(
   }
   std::string encoded_trace;
   TracerHelper::EncodeTrace(trace, &encoded_trace);
-  InstrumentedMutexLock lock_guard(&trace_writer_mutex_);
   return trace_writer_->Write(encoded_trace);
 }
 
@@ -81,7 +79,6 @@ Status BlockCacheTraceWriter::WriteHeader() {
   PutFixed32(&trace.payload, kMinorVersion);
   std::string encoded_trace;
   TracerHelper::EncodeTrace(trace, &encoded_trace);
-  InstrumentedMutexLock lock_guard(&trace_writer_mutex_);
   return trace_writer_->Write(encoded_trace);
 }
 
@@ -216,4 +213,41 @@ Status BlockCacheTraceReader::ReadAccess(BlockCacheTraceRecord* record) {
   return Status::OK();
 }
 
+BlockCacheTracer::BlockCacheTracer() { writer_.store(nullptr); }
+
+BlockCacheTracer::~BlockCacheTracer() { EndTrace(); }
+
+Status BlockCacheTracer::StartTrace(
+    Env* env, const TraceOptions& trace_options,
+    std::unique_ptr<TraceWriter>&& trace_writer) {
+  InstrumentedMutexLock lock_guard(&trace_writer_mutex_);
+  if (writer_.load()) {
+    return Status::OK();
+  }
+  trace_options_ = trace_options;
+  writer_.store(
+      new BlockCacheTraceWriter(env, trace_options, std::move(trace_writer)));
+  return writer_.load()->WriteHeader();
+}
+
+void BlockCacheTracer::EndTrace() {
+  InstrumentedMutexLock lock_guard(&trace_writer_mutex_);
+  if (!writer_.load()) {
+    return;
+  }
+  delete writer_.load();
+  writer_.store(nullptr);
+}
+
+Status BlockCacheTracer::WriteBlockAccess(const BlockCacheTraceRecord& record) {
+  if (!writer_.load() || !ShouldTrace(record, trace_options_)) {
+    return Status::OK();
+  }
+  InstrumentedMutexLock lock_guard(&trace_writer_mutex_);
+  if (!writer_.load()) {
+    return Status::OK();
+  }
+  return writer_.load()->WriteBlockAccess(record);
+}
+
 }  // namespace rocksdb