From 7045b74b47c2049560eb806d69cce6339b22d599 Mon Sep 17 00:00:00 2001
From: Yanqin Jin <yanqin@fb.com>
Date: Tue, 27 Sep 2022 09:04:57 -0700
Subject: [PATCH] Remove timestamp before inserting to WBWI's index (#10742)

Summary:
Currently, this original behavior should not lead to incorrect result, but will violate the contract of CompareWithTimestamp() that when a_has_ts or b_has_ts is false, the slice does not include timestamp.

Resolves https://github.com/facebook/rocksdb/issues/10709

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

Test Plan: make check

Reviewed By: ltamasi

Differential Revision: D39834096

Pulled By: riversand963

fbshipit-source-id: c597600f5a7820734f07d0926cdc224cea5eabe1
---
 .../write_batch_with_index.cc                 |  7 +++++++
 .../write_batch_with_index_test.cc            | 20 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/utilities/write_batch_with_index/write_batch_with_index.cc b/utilities/write_batch_with_index/write_batch_with_index.cc
index 9f65216f7..f43630a48 100644
--- a/utilities/write_batch_with_index/write_batch_with_index.cc
+++ b/utilities/write_batch_with_index/write_batch_with_index.cc
@@ -153,6 +153,13 @@ void WriteBatchWithIndex::Rep::AddNewEntry(uint32_t column_family_id) {
 #endif
   assert(success);
 
+  const Comparator* const ucmp = comparator.GetComparator(column_family_id);
+  size_t ts_sz = ucmp ? ucmp->timestamp_size() : 0;
+
+  if (ts_sz > 0) {
+    key.remove_suffix(ts_sz);
+  }
+
   auto* mem = arena.Allocate(sizeof(WriteBatchIndexEntry));
   auto* index_entry =
       new (mem) WriteBatchIndexEntry(last_entry_offset, column_family_id,
diff --git a/utilities/write_batch_with_index/write_batch_with_index_test.cc b/utilities/write_batch_with_index/write_batch_with_index_test.cc
index a57a95108..50d64963a 100644
--- a/utilities/write_batch_with_index/write_batch_with_index_test.cc
+++ b/utilities/write_batch_with_index/write_batch_with_index_test.cc
@@ -2387,6 +2387,26 @@ TEST_P(WriteBatchWithIndexTest, ColumnFamilyWithTimestamp) {
   }
 }
 
+TEST_P(WriteBatchWithIndexTest, IndexNoTs) {
+  const Comparator* const ucmp = test::BytewiseComparatorWithU64TsWrapper();
+  ColumnFamilyHandleImplDummy cf(1, ucmp);
+  WriteBatchWithIndex wbwi;
+  ASSERT_OK(wbwi.Put(&cf, "a", "a0"));
+  ASSERT_OK(wbwi.Put(&cf, "a", "a1"));
+  {
+    std::string ts;
+    PutFixed64(&ts, 10000);
+    ASSERT_OK(wbwi.GetWriteBatch()->UpdateTimestamps(
+        ts, [](uint32_t cf_id) { return cf_id == 1 ? 8 : 0; }));
+  }
+  {
+    std::string value;
+    Status s = wbwi.GetFromBatch(&cf, options_, "a", &value);
+    ASSERT_OK(s);
+    ASSERT_EQ("a1", value);
+  }
+}
+
 INSTANTIATE_TEST_CASE_P(WBWI, WriteBatchWithIndexTest, testing::Bool());
 }  // namespace ROCKSDB_NAMESPACE