From 685044dff278ff059d47b393d97aedc8f316026d Mon Sep 17 00:00:00 2001
From: Yanqin Jin <yanqin@fb.com>
Date: Wed, 9 Feb 2022 09:49:35 -0800
Subject: [PATCH] Remove timestamp from key in expected state (#9525)

Summary:
The keys as part of write batch read from trace file can contain trailing timestamps.
This PR removes them before calling `ExpectedState`.

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

Test Plan:
make check
make crash_test_with_ts

Reviewed By: ajkr

Differential Revision: D34082358

Pulled By: riversand963

fbshipit-source-id: 78c925659e2a19e4a8278fb4a8ddf5070e265c04
---
 db/dbformat.h                    | 11 ++++++-----
 db_stress_tool/expected_state.cc | 27 ++++++++++++++++++++-------
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/db/dbformat.h b/db/dbformat.h
index 339ce9ff4..137f23aa6 100644
--- a/db/dbformat.h
+++ b/db/dbformat.h
@@ -191,14 +191,15 @@ inline Slice ExtractUserKey(const Slice& internal_key) {
 
 inline Slice ExtractUserKeyAndStripTimestamp(const Slice& internal_key,
                                              size_t ts_sz) {
-  assert(internal_key.size() >= kNumInternalBytes + ts_sz);
-  return Slice(internal_key.data(),
-               internal_key.size() - kNumInternalBytes - ts_sz);
+  Slice ret = internal_key;
+  ret.remove_suffix(kNumInternalBytes + ts_sz);
+  return ret;
 }
 
 inline Slice StripTimestampFromUserKey(const Slice& user_key, size_t ts_sz) {
-  assert(user_key.size() >= ts_sz);
-  return Slice(user_key.data(), user_key.size() - ts_sz);
+  Slice ret = user_key;
+  ret.remove_suffix(ts_sz);
+  return ret;
 }
 
 inline Slice ExtractTimestampFromUserKey(const Slice& user_key, size_t ts_sz) {
diff --git a/db_stress_tool/expected_state.cc b/db_stress_tool/expected_state.cc
index 487ad193f..88e8b7daf 100644
--- a/db_stress_tool/expected_state.cc
+++ b/db_stress_tool/expected_state.cc
@@ -381,8 +381,10 @@ class ExpectedStateTraceRecordHandler : public TraceRecord::Handler,
   // object, but it's convenient and works to share state with the
   // `TraceRecord::Handler`.
 
-  Status PutCF(uint32_t column_family_id, const Slice& key,
+  Status PutCF(uint32_t column_family_id, const Slice& key_with_ts,
                const Slice& value) override {
+    Slice key =
+        StripTimestampFromUserKey(key_with_ts, FLAGS_user_timestamp_size);
     uint64_t key_id;
     if (!GetIntVal(key.ToString(), &key_id)) {
       return Status::Corruption("unable to parse key", key.ToString());
@@ -395,7 +397,10 @@ class ExpectedStateTraceRecordHandler : public TraceRecord::Handler,
     return Status::OK();
   }
 
-  Status DeleteCF(uint32_t column_family_id, const Slice& key) override {
+  Status DeleteCF(uint32_t column_family_id,
+                  const Slice& key_with_ts) override {
+    Slice key =
+        StripTimestampFromUserKey(key_with_ts, FLAGS_user_timestamp_size);
     uint64_t key_id;
     if (!GetIntVal(key.ToString(), &key_id)) {
       return Status::Corruption("unable to parse key", key.ToString());
@@ -407,12 +412,18 @@ class ExpectedStateTraceRecordHandler : public TraceRecord::Handler,
     return Status::OK();
   }
 
-  Status SingleDeleteCF(uint32_t column_family_id, const Slice& key) override {
-    return DeleteCF(column_family_id, key);
+  Status SingleDeleteCF(uint32_t column_family_id,
+                        const Slice& key_with_ts) override {
+    return DeleteCF(column_family_id, key_with_ts);
   }
 
-  Status DeleteRangeCF(uint32_t column_family_id, const Slice& begin_key,
-                       const Slice& end_key) override {
+  Status DeleteRangeCF(uint32_t column_family_id,
+                       const Slice& begin_key_with_ts,
+                       const Slice& end_key_with_ts) override {
+    Slice begin_key =
+        StripTimestampFromUserKey(begin_key_with_ts, FLAGS_user_timestamp_size);
+    Slice end_key =
+        StripTimestampFromUserKey(end_key_with_ts, FLAGS_user_timestamp_size);
     uint64_t begin_key_id, end_key_id;
     if (!GetIntVal(begin_key.ToString(), &begin_key_id)) {
       return Status::Corruption("unable to parse begin key",
@@ -428,8 +439,10 @@ class ExpectedStateTraceRecordHandler : public TraceRecord::Handler,
     return Status::OK();
   }
 
-  Status MergeCF(uint32_t column_family_id, const Slice& key,
+  Status MergeCF(uint32_t column_family_id, const Slice& key_with_ts,
                  const Slice& value) override {
+    Slice key =
+        StripTimestampFromUserKey(key_with_ts, FLAGS_user_timestamp_size);
     return PutCF(column_family_id, key, value);
   }