From bd4b8d648726b89e4126a79ca0ab63910a08e1c8 Mon Sep 17 00:00:00 2001
From: akankshamahajan <akankshamahajan@fb.com>
Date: Wed, 18 Jan 2023 13:24:37 -0800
Subject: [PATCH] Fix crash in block_cache_trace_analyzer if reference key is
 null in case of MultiGet (#11042)

Summary:
Same as title
Error:
```
block_cache_trace_analyzer: ./db/dbformat.h:421: uint64_t rocksdb::GetInternalKeySeqno(const rocksdb::Slice&): Assertion `n >= kNumInternalBytes' failed.
Aborted (core dumped)
```

Pull Request resolved: https://github.com/facebook/rocksdb/pull/11042

Test Plan:
- Added new unit test which fails without the fix.
                  - Also ran manually on traces to confirm.

Reviewed By: anand1976

Differential Revision: D42481587

Pulled By: akankshamahajan15

fbshipit-source-id: 7c33eb03a4a4d8ffbabcfbe0efa1e4d11bde3ba2
---
 .../block_cache_trace_analyzer_test.cc        | 67 ++++++++++++++++++-
 trace_replay/block_cache_tracer.cc            |  5 ++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/tools/block_cache_analyzer/block_cache_trace_analyzer_test.cc b/tools/block_cache_analyzer/block_cache_trace_analyzer_test.cc
index c5d9b1452..608344805 100644
--- a/tools/block_cache_analyzer/block_cache_trace_analyzer_test.cc
+++ b/tools/block_cache_analyzer/block_cache_trace_analyzer_test.cc
@@ -95,7 +95,8 @@ class BlockCacheTracerTest : public testing::Test {
   }
 
   void WriteBlockAccess(BlockCacheTraceWriter* writer, uint32_t from_key_id,
-                        TraceType block_type, uint32_t nblocks) {
+                        TraceType block_type, uint32_t nblocks,
+                        bool is_referenced_key_null = false) {
     assert(writer);
     for (uint32_t i = 0; i < nblocks; i++) {
       uint32_t key_id = from_key_id + i;
@@ -122,6 +123,11 @@ class BlockCacheTracerTest : public testing::Test {
       record.referenced_key =
           kRefKeyPrefix + std::to_string(key_id) + std::string(8, 0);
       record.referenced_key_exist_in_block = true;
+      if (is_referenced_key_null &&
+          record.caller == TableReaderCaller::kUserMultiGet) {
+        record.referenced_key = "";
+        record.get_from_user_specified_snapshot = true;
+      }
       record.num_keys_in_block = kNumKeysInBlock;
       ASSERT_OK(writer->WriteBlockAccess(
           record, record.block_key, record.cf_name, record.referenced_key));
@@ -717,6 +723,65 @@ TEST_F(BlockCacheTracerTest, MixedBlocks) {
   }
 }
 
+TEST_F(BlockCacheTracerTest, MultiGetWithNullReferenceKey) {
+  {
+    // Generate a trace file containing MultiGet records with reference key
+    // being 0.
+    BlockCacheTraceWriterOptions trace_writer_opt;
+    std::unique_ptr<TraceWriter> trace_writer;
+    const auto& clock = env_->GetSystemClock();
+    ASSERT_OK(NewFileTraceWriter(env_, env_options_, trace_file_path_,
+                                 &trace_writer));
+    std::unique_ptr<BlockCacheTraceWriter> block_cache_trace_writer =
+        NewBlockCacheTraceWriter(clock.get(), trace_writer_opt,
+                                 std::move(trace_writer));
+    ASSERT_NE(block_cache_trace_writer, nullptr);
+    ASSERT_OK(block_cache_trace_writer->WriteHeader());
+    // Write blocks of different types.
+
+    WriteBlockAccess(block_cache_trace_writer.get(), 0,
+                     TraceType::kBlockTraceUncompressionDictBlock, 10, true);
+    WriteBlockAccess(block_cache_trace_writer.get(), 10,
+                     TraceType::kBlockTraceDataBlock, 10, true);
+    WriteBlockAccess(block_cache_trace_writer.get(), 20,
+                     TraceType::kBlockTraceFilterBlock, 10, true);
+    WriteBlockAccess(block_cache_trace_writer.get(), 30,
+                     TraceType::kBlockTraceIndexBlock, 10, true);
+    WriteBlockAccess(block_cache_trace_writer.get(), 40,
+                     TraceType::kBlockTraceRangeDeletionBlock, 10, true);
+    ASSERT_OK(env_->FileExists(trace_file_path_));
+  }
+
+  {
+    // Verify trace file is generated correctly.
+    std::unique_ptr<TraceReader> trace_reader;
+    ASSERT_OK(NewFileTraceReader(env_, env_options_, trace_file_path_,
+                                 &trace_reader));
+    BlockCacheTraceReader reader(std::move(trace_reader));
+    BlockCacheTraceHeader header;
+    ASSERT_OK(reader.ReadHeader(&header));
+    ASSERT_EQ(static_cast<uint32_t>(kMajorVersion),
+              header.rocksdb_major_version);
+    ASSERT_EQ(static_cast<uint32_t>(kMinorVersion),
+              header.rocksdb_minor_version);
+    std::string human_readable_trace_file_path =
+        test_path_ + "/readable_block_cache_trace";
+    // Read blocks.
+    BlockCacheTraceAnalyzer analyzer(
+        trace_file_path_,
+        /*output_dir=*/"",
+        /*human_readable_trace_file_path=*/human_readable_trace_file_path,
+        /*compute_reuse_distance=*/true,
+        /*mrc_only=*/false,
+        /*is_human_readable_trace_file=*/false,
+        /*cache_simulator=*/nullptr);
+    // The analyzer ends when it detects an incomplete access record.
+    ASSERT_EQ(Status::Incomplete(""), analyzer.Analyze());
+
+    ASSERT_OK(env_->DeleteFile(human_readable_trace_file_path));
+  }
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {
diff --git a/trace_replay/block_cache_tracer.cc b/trace_replay/block_cache_tracer.cc
index 508596913..6a3442c5b 100644
--- a/trace_replay/block_cache_tracer.cc
+++ b/trace_replay/block_cache_tracer.cc
@@ -8,6 +8,7 @@
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
+#include <string>
 
 #include "db/db_impl/db_impl.h"
 #include "db/dbformat.h"
@@ -81,6 +82,10 @@ uint64_t BlockCacheTraceHelper::GetSequenceNumber(
   if (!IsGetOrMultiGet(access.caller)) {
     return 0;
   }
+  if (access.caller == TableReaderCaller::kUserMultiGet &&
+      access.referenced_key.size() < 4) {
+    return 0;
+  }
   return access.get_from_user_specified_snapshot
              ? 1 + GetInternalKeySeqno(access.referenced_key)
              : 0;