Add `IteratorTraceExecutionResult` for iterator related trace records. (#8687)

Summary:
- Allow to get `Valid()`, `status()`, `key()` and `value()` of an iterator from `IteratorTraceExecutionResult`.
- Move lower bound and upper bound from `IteratorSeekQueryTraceRecord` to `IteratorQueryTraceRecord`.

Added test in `DBTest2.TraceAndReplay`.

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

Reviewed By: zhichao-cao

Differential Revision: D30457630

Pulled By: autopear

fbshipit-source-id: be433099a25895b3aa6f0c00f95ad7b1d7489c1d

db/db_test2.cc

@@ -4256,13 +4256,6 @@ class TraceExecutionResultHandler : public TraceRecordResult::Handler {
         writes_++;
         break;
       }
-      case kTraceIteratorSeek:
-      case kTraceIteratorSeekForPrev: {
-        total_latency_ += result.GetLatency();
-        cnt_++;
-        seeks_++;
-        break;
-      }
       default:
         return Status::Corruption("Type mismatch.");
     }
@@ -4309,6 +4302,25 @@ class TraceExecutionResultHandler : public TraceRecordResult::Handler {
     return Status::OK();
   }
 
+  virtual Status Handle(const IteratorTraceExecutionResult& result) override {
+    if (result.GetStartTimestamp() > result.GetEndTimestamp()) {
+      return Status::InvalidArgument("Invalid timestamps.");
+    }
+    result.GetStatus().PermitUncheckedError();
+    switch (result.GetTraceType()) {
+      case kTraceIteratorSeek:
+      case kTraceIteratorSeekForPrev: {
+        total_latency_ += result.GetLatency();
+        cnt_++;
+        seeks_++;
+        break;
+      }
+      default:
+        return Status::Corruption("Type mismatch.");
+    }
+    return Status::OK();
+  }
+
   void Reset() {
     total_latency_ = 0;
     cnt_ = 0;
@@ -4644,23 +4656,39 @@ TEST_F(DBTest2, TraceAndManualReplay) {
         continue;
       }
       if (s.ok()) {
+        ASSERT_OK(replayer->Execute(record, &result));
+        if (result != nullptr) {
+          ASSERT_OK(result->Accept(&res_handler));
           if (record->GetTraceType() == kTraceIteratorSeek ||
               record->GetTraceType() == kTraceIteratorSeekForPrev) {
-          IteratorSeekQueryTraceRecord* iter_r =
+            IteratorSeekQueryTraceRecord* iter_rec =
                 dynamic_cast<IteratorSeekQueryTraceRecord*>(record.get());
+            IteratorTraceExecutionResult* iter_res =
+                dynamic_cast<IteratorTraceExecutionResult*>(result.get());
             // Check if lower/upper bounds are correctly saved and decoded.
-          lower_bound = iter_r->GetLowerBound();
+            std::string lower_str = iter_rec->GetLowerBound().ToString();
-          if (!lower_bound.empty()) {
+            std::string upper_str = iter_rec->GetUpperBound().ToString();
-            ASSERT_EQ(lower_bound.ToString(), "iter-1");
+            std::string iter_key = iter_res->GetKey().ToString();
+            std::string iter_value = iter_res->GetValue().ToString();
+            if (!lower_str.empty() && !upper_str.empty()) {
+              ASSERT_EQ(lower_str, "iter-1");
+              ASSERT_EQ(upper_str, "iter-3");
+              if (iter_res->GetValid()) {
+                // If iterator is valid, then lower_bound <= key < upper_bound.
+                ASSERT_GE(iter_key, lower_str);
+                ASSERT_LT(iter_key, upper_str);
+              } else {
+                // If iterator is invalid, then
+                //   key < lower_bound or key >= upper_bound.
+                ASSERT_TRUE(iter_key < lower_str || iter_key >= upper_str);
               }
-          upper_bound = iter_r->GetUpperBound();
+            }
-          if (!upper_bound.empty()) {
+            // If iterator is invalid, the key and value should be empty.
-            ASSERT_EQ(upper_bound.ToString(), "iter-3");
+            if (!iter_res->GetValid()) {
+              ASSERT_TRUE(iter_key.empty());
+              ASSERT_TRUE(iter_value.empty());
             }
           }
-        ASSERT_OK(replayer->Execute(record, &result));
-        if (result != nullptr) {
-          ASSERT_OK(result->Accept(&res_handler));
           result.reset();
         }
       }

include/rocksdb/trace_record.h

@@ -65,19 +65,15 @@ class TraceRecord {
    public:
     virtual ~Handler() = default;
 
-    // Handle WriteQueryTraceRecord
     virtual Status Handle(const WriteQueryTraceRecord& record,
                           std::unique_ptr<TraceRecordResult>* result) = 0;
 
-    // Handle GetQueryTraceRecord
     virtual Status Handle(const GetQueryTraceRecord& record,
                           std::unique_ptr<TraceRecordResult>* result) = 0;
 
-    // Handle IteratorSeekQueryTraceRecord
     virtual Status Handle(const IteratorSeekQueryTraceRecord& record,
                           std::unique_ptr<TraceRecordResult>* result) = 0;
 
-    // Handle MultiGetQueryTraceRecord
     virtual Status Handle(const MultiGetQueryTraceRecord& record,
                           std::unique_ptr<TraceRecordResult>* result) = 0;
   };
@@ -152,6 +148,23 @@ class GetQueryTraceRecord : public QueryTraceRecord {
 class IteratorQueryTraceRecord : public QueryTraceRecord {
  public:
   explicit IteratorQueryTraceRecord(uint64_t timestamp);
+
+  IteratorQueryTraceRecord(PinnableSlice&& lower_bound,
+                           PinnableSlice&& upper_bound, uint64_t timestamp);
+
+  IteratorQueryTraceRecord(const std::string& lower_bound,
+                           const std::string& upper_bound, uint64_t timestamp);
+
+  virtual ~IteratorQueryTraceRecord() override;
+
+  // Get the iterator's lower/upper bound. They may be used in ReadOptions to
+  // create an Iterator instance.
+  virtual Slice GetLowerBound() const;
+  virtual Slice GetUpperBound() const;
+
+ private:
+  PinnableSlice lower_;
+  PinnableSlice upper_;
 };
 
 // Trace record for Iterator::Seek() and Iterator::SeekForPrev() operation.
@@ -193,12 +206,6 @@ 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;
 
@@ -206,8 +213,6 @@ class IteratorSeekQueryTraceRecord : public IteratorQueryTraceRecord {
   SeekType type_;
   uint32_t cf_id_;
   PinnableSlice key_;
-  PinnableSlice lower_;
-  PinnableSlice upper_;
 };
 
 // Trace record for DB::MultiGet() operation.

include/rocksdb/trace_record_result.h

@@ -9,11 +9,13 @@
 #include <vector>
 
 #include "rocksdb/rocksdb_namespace.h"
+#include "rocksdb/slice.h"
 #include "rocksdb/status.h"
 #include "rocksdb/trace_record.h"
 
 namespace ROCKSDB_NAMESPACE {
 
+class IteratorTraceExecutionResult;
 class MultiValuesTraceExecutionResult;
 class SingleValueTraceExecutionResult;
 class StatusOnlyTraceExecutionResult;
@@ -34,42 +36,14 @@ class TraceRecordResult {
    public:
     virtual ~Handler() = default;
 
-    // Handle StatusOnlyTraceExecutionResult
     virtual Status Handle(const StatusOnlyTraceExecutionResult& result) = 0;
 
-    // Handle SingleValueTraceExecutionResult
     virtual Status Handle(const SingleValueTraceExecutionResult& result) = 0;
 
-    // Handle MultiValuesTraceExecutionResult
     virtual Status Handle(const MultiValuesTraceExecutionResult& result) = 0;
-  };
 
-  /*
+    virtual Status Handle(const IteratorTraceExecutionResult& result) = 0;
-   * Example handler to just print the trace record execution results.
+  };
-   *
-   * class ResultPrintHandler : public TraceRecordResult::Handler {
-   *  public:
-   *   ResultPrintHandler();
-   *   ~ResultPrintHandler() override {}
-   *
-   *   Status Handle(const StatusOnlyTraceExecutionResult& result) override {
-   *     std::cout << "Status: " << result.GetStatus().ToString() << std::endl;
-   *   }
-   *
-   *   Status Handle(const SingleValueTraceExecutionResult& result) override {
-   *     std::cout << "Status: " << result.GetStatus().ToString()
-   *               << ", value: " << result.GetValue() << std::endl;
-   *   }
-   *
-   *   Status Handle(const MultiValuesTraceExecutionResult& result) override {
-   *     size_t size = result.GetMultiStatus().size();
-   *     for (size_t i = 0; i < size; i++) {
-   *       std::cout << "Status: " << result.GetMultiStatus()[i].ToString()
-   *                 << ", value: " << result.GetValues()[i] << std::endl;
-   *     }
-   *   }
-   * };
-   * */
 
   // Accept the handler.
   virtual Status Accept(Handler* handler) = 0;
@@ -106,8 +80,7 @@ class TraceExecutionResult : public TraceRecordResult {
 };
 
 // Result for operations that only return a single Status.
-// Example operations: DB::Write(), Iterator::Seek() and
+// Example operation: DB::Write()
-// Iterator::SeekForPrev().
 class StatusOnlyTraceExecutionResult : public TraceExecutionResult {
  public:
   StatusOnlyTraceExecutionResult(Status status, uint64_t start_timestamp,
@@ -138,7 +111,7 @@ class SingleValueTraceExecutionResult : public TraceExecutionResult {
 
   virtual ~SingleValueTraceExecutionResult() override;
 
-  // Return status of DB::Get(), etc.
+  // Return status of DB::Get().
   virtual const Status& GetStatus() const;
 
   // Value for the searched key.
@@ -151,7 +124,7 @@ class SingleValueTraceExecutionResult : public TraceExecutionResult {
   std::string value_;
 };
 
-// Result for operations that return multiple Status(es) and values.
+// Result for operations that return multiple Status(es) and values as vectors.
 // Example operation: DB::MultiGet()
 class MultiValuesTraceExecutionResult : public TraceExecutionResult {
  public:
@@ -162,7 +135,7 @@ class MultiValuesTraceExecutionResult : public TraceExecutionResult {
 
   virtual ~MultiValuesTraceExecutionResult() override;
 
-  // Returned Status(es) of DB::MultiGet(), etc.
+  // Returned Status(es) of DB::MultiGet().
   virtual const std::vector<Status>& GetMultiStatus() const;
 
   // Returned values for the searched keys.
@@ -175,4 +148,40 @@ class MultiValuesTraceExecutionResult : public TraceExecutionResult {
   std::vector<std::string> values_;
 };
 
+// Result for Iterator operations.
+// Example operations: Iterator::Seek(), Iterator::SeekForPrev()
+class IteratorTraceExecutionResult : public TraceExecutionResult {
+ public:
+  IteratorTraceExecutionResult(bool valid, Status status, PinnableSlice&& key,
+                               PinnableSlice&& value, uint64_t start_timestamp,
+                               uint64_t end_timestamp, TraceType trace_type);
+
+  IteratorTraceExecutionResult(bool valid, Status status,
+                               const std::string& key, const std::string& value,
+                               uint64_t start_timestamp, uint64_t end_timestamp,
+                               TraceType trace_type);
+
+  virtual ~IteratorTraceExecutionResult() override;
+
+  // Return if the Iterator is valid.
+  virtual bool GetValid() const;
+
+  // Return the status of the Iterator.
+  virtual const Status& GetStatus() const;
+
+  // Key of the current iterating entry, empty if GetValid() is false.
+  virtual Slice GetKey() const;
+
+  // Value of the current iterating entry, empty if GetValid() is false.
+  virtual Slice GetValue() const;
+
+  virtual Status Accept(Handler* handler) override;
+
+ private:
+  bool valid_;
+  Status status_;
+  PinnableSlice key_;
+  PinnableSlice value_;
+};
+
 }  // namespace ROCKSDB_NAMESPACE

trace_replay/trace_record.cc

@@ -82,6 +82,27 @@ Status GetQueryTraceRecord::Accept(Handler* handler,
 IteratorQueryTraceRecord::IteratorQueryTraceRecord(uint64_t timestamp)
     : QueryTraceRecord(timestamp) {}
 
+IteratorQueryTraceRecord::IteratorQueryTraceRecord(PinnableSlice&& lower_bound,
+                                                   PinnableSlice&& upper_bound,
+                                                   uint64_t timestamp)
+    : QueryTraceRecord(timestamp),
+      lower_(std::move(lower_bound)),
+      upper_(std::move(upper_bound)) {}
+
+IteratorQueryTraceRecord::IteratorQueryTraceRecord(
+    const std::string& lower_bound, const std::string& upper_bound,
+    uint64_t timestamp)
+    : QueryTraceRecord(timestamp) {
+  lower_.PinSelf(lower_bound);
+  upper_.PinSelf(upper_bound);
+}
+
+IteratorQueryTraceRecord::~IteratorQueryTraceRecord() {}
+
+Slice IteratorQueryTraceRecord::GetLowerBound() const { return Slice(lower_); }
+
+Slice IteratorQueryTraceRecord::GetUpperBound() const { return Slice(upper_); }
+
 // IteratorSeekQueryTraceRecord
 IteratorSeekQueryTraceRecord::IteratorSeekQueryTraceRecord(
     SeekType seek_type, uint32_t column_family_id, PinnableSlice&& key,
@@ -104,23 +125,20 @@ IteratorSeekQueryTraceRecord::IteratorSeekQueryTraceRecord(
     SeekType seek_type, uint32_t column_family_id, PinnableSlice&& key,
     PinnableSlice&& lower_bound, PinnableSlice&& upper_bound,
     uint64_t timestamp)
-    : IteratorQueryTraceRecord(timestamp),
+    : IteratorQueryTraceRecord(std::move(lower_bound), std::move(upper_bound),
+                               timestamp),
       type_(seek_type),
       cf_id_(column_family_id),
-      key_(std::move(key)),
+      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),
+    : IteratorQueryTraceRecord(lower_bound, upper_bound, timestamp),
       type_(seek_type),
       cf_id_(column_family_id) {
   key_.PinSelf(key);
-  lower_.PinSelf(lower_bound);
-  upper_.PinSelf(upper_bound);
 }
 
 IteratorSeekQueryTraceRecord::~IteratorSeekQueryTraceRecord() { key_.clear(); }
@@ -140,14 +158,6 @@ 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

@@ -120,12 +120,21 @@ Status TraceExecutionHandler::Handle(
   uint64_t end = clock_->NowMicros();
 
   Status s = single_iter->status();
-  delete single_iter;
-
   if (s.ok() && result != nullptr) {
-    result->reset(new StatusOnlyTraceExecutionResult(s, start, end,
+    if (single_iter->Valid()) {
+      PinnableSlice ps_key;
+      ps_key.PinSelf(single_iter->key());
+      PinnableSlice ps_value;
+      ps_value.PinSelf(single_iter->value());
+      result->reset(new IteratorTraceExecutionResult(
+          true, s, std::move(ps_key), std::move(ps_value), start, end,
           record.GetTraceType()));
+    } else {
+      result->reset(new IteratorTraceExecutionResult(
+          false, s, "", "", start, end, record.GetTraceType()));
     }
+  }
+  delete single_iter;
 
   return s;
 }

trace_replay/trace_record_result.cc

@@ -103,4 +103,44 @@ Status MultiValuesTraceExecutionResult::Accept(Handler* handler) {
   return handler->Handle(*this);
 }
 
+// IteratorTraceExecutionResult
+IteratorTraceExecutionResult::IteratorTraceExecutionResult(
+    bool valid, Status status, PinnableSlice&& key, PinnableSlice&& value,
+    uint64_t start_timestamp, uint64_t end_timestamp, TraceType trace_type)
+    : TraceExecutionResult(start_timestamp, end_timestamp, trace_type),
+      valid_(valid),
+      status_(std::move(status)),
+      key_(std::move(key)),
+      value_(std::move(value)) {}
+
+IteratorTraceExecutionResult::IteratorTraceExecutionResult(
+    bool valid, Status status, const std::string& key, const std::string& value,
+    uint64_t start_timestamp, uint64_t end_timestamp, TraceType trace_type)
+    : TraceExecutionResult(start_timestamp, end_timestamp, trace_type),
+      valid_(valid),
+      status_(std::move(status)) {
+  key_.PinSelf(key);
+  value_.PinSelf(value);
+}
+
+IteratorTraceExecutionResult::~IteratorTraceExecutionResult() {
+  key_.clear();
+  value_.clear();
+}
+
+bool IteratorTraceExecutionResult::GetValid() const { return valid_; }
+
+const Status& IteratorTraceExecutionResult::GetStatus() const {
+  return status_;
+}
+
+Slice IteratorTraceExecutionResult::GetKey() const { return Slice(key_); }
+
+Slice IteratorTraceExecutionResult::GetValue() const { return Slice(value_); }
+
+Status IteratorTraceExecutionResult::Accept(Handler* handler) {
+  assert(handler != nullptr);
+  return handler->Handle(*this);
+}
+
 }  // namespace ROCKSDB_NAMESPACE