Add iterator's lower and upper bounds to TraceRecord (#8677)

Summary:
Trace file V2 added lower/upper bounds to `Iterator::Seek()` and `Iterator::SeekForPrev()`. They were not used anywhere during the execution of a `TraceRecord`. Now they are added to be used by `ReadOptions` during `Iterator::Seek()` and `Iterator::SeekForPrev()` if they are set.

Added test cases in `DBTest2.TraceAndManualReplay`.

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

Reviewed By: zhichao-cao

Differential Revision: D30438255

Pulled By: autopear

fbshipit-source-id: 82563006be0b69155990e506a74951c18af8d288

db/db_test2.cc

@@ -4538,6 +4538,37 @@ TEST_F(DBTest2, TraceAndManualReplay) {
   single_iter = db_->NewIterator(ro);
   single_iter->Seek("f");
   single_iter->SeekForPrev("g");
+  ASSERT_OK(single_iter->status());
+  delete single_iter;
+
+  // Write some sequenced keys for testing lower/upper bounds of iterator.
+  batch.Clear();
+  ASSERT_OK(batch.Put("iter-0", "iter-0"));
+  ASSERT_OK(batch.Put("iter-1", "iter-1"));
+  ASSERT_OK(batch.Put("iter-2", "iter-2"));
+  ASSERT_OK(batch.Put("iter-3", "iter-3"));
+  ASSERT_OK(batch.Put("iter-4", "iter-4"));
+  ASSERT_OK(db_->Write(wo, &batch));
+
+  ReadOptions bounded_ro = ro;
+  Slice lower_bound("iter-1");
+  Slice upper_bound("iter-3");
+  bounded_ro.iterate_lower_bound = &lower_bound;
+  bounded_ro.iterate_upper_bound = &upper_bound;
+  single_iter = db_->NewIterator(bounded_ro);
+  single_iter->Seek("iter-0");
+  ASSERT_EQ(single_iter->key().ToString(), "iter-1");
+  single_iter->Seek("iter-2");
+  ASSERT_EQ(single_iter->key().ToString(), "iter-2");
+  single_iter->Seek("iter-4");
+  ASSERT_FALSE(single_iter->Valid());
+  single_iter->SeekForPrev("iter-0");
+  ASSERT_FALSE(single_iter->Valid());
+  single_iter->SeekForPrev("iter-2");
+  ASSERT_EQ(single_iter->key().ToString(), "iter-2");
+  single_iter->SeekForPrev("iter-4");
+  ASSERT_EQ(single_iter->key().ToString(), "iter-2");
+  ASSERT_OK(single_iter->status());
   delete single_iter;
 
   ASSERT_EQ("1", Get(0, "a"));
@@ -4548,6 +4579,9 @@ TEST_F(DBTest2, TraceAndManualReplay) {
   ASSERT_EQ("NOT_FOUND", Get(1, "leveldb"));
 
   // Same as TraceAndReplay, Write x 8, Get x 3, Seek x 2.
+  // Plus 1 WriteBatch for iterator with lower/upper bounds, and 6
+  // Seek(ForPrev)s.
+  // Total Write x 9, Get x 3, Seek x 8
   ASSERT_OK(db_->EndTrace());
   // These should not get into the trace file as it is after EndTrace.
   ASSERT_OK(Put("hello", "world"));
@@ -4610,6 +4644,20 @@ TEST_F(DBTest2, TraceAndManualReplay) {
         continue;
       }
       if (s.ok()) {
+        if (record->GetTraceType() == kTraceIteratorSeek ||
+            record->GetTraceType() == kTraceIteratorSeekForPrev) {
+          IteratorSeekQueryTraceRecord* iter_r =
+              dynamic_cast<IteratorSeekQueryTraceRecord*>(record.get());
+          // Check if lower/upper bounds are correctly saved and decoded.
+          lower_bound = iter_r->GetLowerBound();
+          if (!lower_bound.empty()) {
+            ASSERT_EQ(lower_bound.ToString(), "iter-1");
+          }
+          upper_bound = iter_r->GetUpperBound();
+          if (!upper_bound.empty()) {
+            ASSERT_EQ(upper_bound.ToString(), "iter-3");
+          }
+        }
         ASSERT_OK(replayer->Execute(record, &result));
         if (result != nullptr) {
           ASSERT_OK(result->Accept(&res_handler));
@@ -4622,9 +4670,9 @@ TEST_F(DBTest2, TraceAndManualReplay) {
     ASSERT_TRUE(s.IsIncomplete());
     ASSERT_TRUE(replayer->Next(nullptr).IsIncomplete());
     ASSERT_GT(res_handler.GetAvgLatency(), 0.0);
-    ASSERT_EQ(res_handler.GetNumWrites(), 8);
+    ASSERT_EQ(res_handler.GetNumWrites(), 9);
     ASSERT_EQ(res_handler.GetNumGets(), 3);
-    ASSERT_EQ(res_handler.GetNumIterSeeks(), 2);
+    ASSERT_EQ(res_handler.GetNumIterSeeks(), 8);
     ASSERT_EQ(res_handler.GetNumMultiGets(), 0);
     res_handler.Reset();
   }

include/rocksdb/trace_record.h

@@ -169,6 +169,16 @@ class IteratorSeekQueryTraceRecord : public IteratorQueryTraceRecord {
   IteratorSeekQueryTraceRecord(SeekType seekType, uint32_t column_family_id,
                                const std::string& key, uint64_t timestamp);
 
+  IteratorSeekQueryTraceRecord(SeekType seekType, uint32_t column_family_id,
+                               PinnableSlice&& key, PinnableSlice&& lower_bound,
+                               PinnableSlice&& upper_bound, uint64_t timestamp);
+
+  IteratorSeekQueryTraceRecord(SeekType seekType, uint32_t column_family_id,
+                               const std::string& key,
+                               const std::string& lower_bound,
+                               const std::string& upper_bound,
+                               uint64_t timestamp);
+
   virtual ~IteratorSeekQueryTraceRecord() override;
 
   // Trace type matches the seek type.
@@ -183,6 +193,12 @@ class IteratorSeekQueryTraceRecord : public IteratorQueryTraceRecord {
   // Key to seek to.
   virtual Slice GetKey() const;
 
+  // Iterate lower bound.
+  virtual Slice GetLowerBound() const;
+
+  // Iterate upper bound.
+  virtual Slice GetUpperBound() const;
+
   Status Accept(Handler* handler,
                 std::unique_ptr<TraceRecordResult>* result) override;
 
@@ -190,6 +206,8 @@ class IteratorSeekQueryTraceRecord : public IteratorQueryTraceRecord {
   SeekType type_;
   uint32_t cf_id_;
   PinnableSlice key_;
+  PinnableSlice lower_;
+  PinnableSlice upper_;
 };
 
 // Trace record for DB::MultiGet() operation.

trace_replay/trace_record.cc

@@ -100,6 +100,29 @@ IteratorSeekQueryTraceRecord::IteratorSeekQueryTraceRecord(
   key_.PinSelf(key);
 }
 
+IteratorSeekQueryTraceRecord::IteratorSeekQueryTraceRecord(
+    SeekType seek_type, uint32_t column_family_id, PinnableSlice&& key,
+    PinnableSlice&& lower_bound, PinnableSlice&& upper_bound,
+    uint64_t timestamp)
+    : IteratorQueryTraceRecord(timestamp),
+      type_(seek_type),
+      cf_id_(column_family_id),
+      key_(std::move(key)),
+      lower_(std::move(lower_bound)),
+      upper_(std::move(upper_bound)) {}
+
+IteratorSeekQueryTraceRecord::IteratorSeekQueryTraceRecord(
+    SeekType seek_type, uint32_t column_family_id, const std::string& key,
+    const std::string& lower_bound, const std::string& upper_bound,
+    uint64_t timestamp)
+    : IteratorQueryTraceRecord(timestamp),
+      type_(seek_type),
+      cf_id_(column_family_id) {
+  key_.PinSelf(key);
+  lower_.PinSelf(lower_bound);
+  upper_.PinSelf(upper_bound);
+}
+
 IteratorSeekQueryTraceRecord::~IteratorSeekQueryTraceRecord() { key_.clear(); }
 
 TraceType IteratorSeekQueryTraceRecord::GetTraceType() const {
@@ -117,6 +140,14 @@ uint32_t IteratorSeekQueryTraceRecord::GetColumnFamilyID() const {
 
 Slice IteratorSeekQueryTraceRecord::GetKey() const { return Slice(key_); }
 
+Slice IteratorSeekQueryTraceRecord::GetLowerBound() const {
+  return Slice(lower_);
+}
+
+Slice IteratorSeekQueryTraceRecord::GetUpperBound() const {
+  return Slice(upper_);
+}
+
 Status IteratorSeekQueryTraceRecord::Accept(
     Handler* handler, std::unique_ptr<TraceRecordResult>* result) {
   assert(handler != nullptr);

trace_replay/trace_record_handler.cc

@@ -93,7 +93,16 @@ Status TraceExecutionHandler::Handle(
     return Status::Corruption("Invalid Column Family ID.");
   }
 
-  Iterator* single_iter = db_->NewIterator(read_opts_, it->second);
+  ReadOptions r_opts = read_opts_;
+  Slice lower = record.GetLowerBound();
+  if (!lower.empty()) {
+    r_opts.iterate_lower_bound = &lower;
+  }
+  Slice upper = record.GetUpperBound();
+  if (!upper.empty()) {
+    r_opts.iterate_upper_bound = &upper;
+  }
+  Iterator* single_iter = db_->NewIterator(r_opts, it->second);
 
   uint64_t start = clock_->NowMicros();
 

trace_replay/trace_replay.cc

@@ -225,14 +225,12 @@ Status TracerHelper::DecodeIterRecord(Trace* trace, int trace_file_version,
 
   uint32_t cf_id = 0;
   Slice iter_key;
+  Slice lower_bound;
+  Slice upper_bound;
 
   if (trace_file_version < 2) {
     DecodeCFAndKey(trace->payload, &cf_id, &iter_key);
   } else {
-    // Are these two used anywhere?
-    Slice lower_bound;
-    Slice upper_bound;
-
     Slice buf(trace->payload);
     GetFixed64(&buf, &trace->payload_map);
     int64_t payload_map = static_cast<int64_t>(trace->payload_map);
@@ -264,9 +262,14 @@ Status TracerHelper::DecodeIterRecord(Trace* trace, int trace_file_version,
   if (record != nullptr) {
     PinnableSlice ps_key;
     ps_key.PinSelf(iter_key);
+    PinnableSlice ps_lower;
+    ps_lower.PinSelf(lower_bound);
+    PinnableSlice ps_upper;
+    ps_upper.PinSelf(upper_bound);
     record->reset(new IteratorSeekQueryTraceRecord(
         static_cast<IteratorSeekQueryTraceRecord::SeekType>(trace->type), cf_id,
-        std::move(ps_key), trace->ts));
+        std::move(ps_key), std::move(ps_lower), std::move(ps_upper),
+        trace->ts));
   }
 
   return Status::OK();