WritePrepared: optimize read path by avoiding virtual (#5018)

Summary:
The read path includes a callback function, ReadCallback, which would eventually calls IsInSnapshot to figure if a particular seq is in the reading snapshot or not. This callback is virtual, which adds the cost of multiple virtual function call to each read. The first few checks in IsInSnapshot, however, are quite trivial and take care of majority of the cases. The patch moves those to a non-virtual function in the the parent class, ReadCallback, to lower the virtual callback cost.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/5018

Differential Revision: D14226562

Pulled By: maysamyabandeh

fbshipit-source-id: 6feed5b34f3b082e52092c5ef143e29b49c46b44

db/db_iterator_test.cc

@@ -20,7 +20,7 @@ namespace rocksdb {
 
 // A dumb ReadCallback which saying every key is committed.
 class DummyReadCallback : public ReadCallback {
-  bool IsVisible(SequenceNumber /*seq*/) override { return true; }
+  bool IsVisibleFullCheck(SequenceNumber /*seq*/) override { return true; }
 };
 
 // Test param:
@@ -2479,7 +2479,7 @@ TEST_F(DBIteratorWithReadCallbackTest, ReadCallback) {
     explicit TestReadCallback(SequenceNumber last_visible_seq)
         : last_visible_seq_(last_visible_seq) {}
 
-    bool IsVisible(SequenceNumber seq) override {
+    bool IsVisibleFullCheck(SequenceNumber seq) override {
       return seq <= last_visible_seq_;
     }
 

db/db_merge_operator_test.cc

@@ -20,7 +20,7 @@ class TestReadCallback : public ReadCallback {
                    SequenceNumber snapshot_seq)
       : snapshot_checker_(snapshot_checker), snapshot_seq_(snapshot_seq) {}
 
-  bool IsVisible(SequenceNumber seq) override {
+  bool IsVisibleFullCheck(SequenceNumber seq) override {
     return snapshot_checker_->CheckInSnapshot(seq, snapshot_seq_) ==
            SnapshotCheckerResult::kInSnapshot;
   }

db/db_test2.cc

@@ -2645,7 +2645,9 @@ TEST_F(DBTest2, ReadCallbackTest) {
   class TestReadCallback : public ReadCallback {
    public:
     explicit TestReadCallback(SequenceNumber snapshot) : snapshot_(snapshot) {}
-    bool IsVisible(SequenceNumber seq) override { return seq <= snapshot_; }
+    bool IsVisibleFullCheck(SequenceNumber seq) override {
+      return seq <= snapshot_;
+    }
 
    private:
     SequenceNumber snapshot_;

db/read_callback.h

@@ -11,11 +11,27 @@ namespace rocksdb {
 
 class ReadCallback {
  public:
+  ReadCallback() {}
+  ReadCallback(SequenceNumber snapshot) : snapshot_(snapshot) {}
+  ReadCallback(SequenceNumber snapshot, SequenceNumber min_uncommitted)
+      : snapshot_(snapshot), min_uncommitted_(min_uncommitted) {}
+
   virtual ~ReadCallback() {}
 
   // Will be called to see if the seq number visible; if not it moves on to
   // the next seq number.
-  virtual bool IsVisible(SequenceNumber seq) = 0;
+  virtual bool IsVisibleFullCheck(SequenceNumber seq) = 0;
+
+  inline bool IsVisible(SequenceNumber seq) {
+    if (seq == 0 || seq < min_uncommitted_) {
+      assert(seq <= snapshot_);
+      return true;
+    } else if (snapshot_ < seq) {
+      return false;
+    } else {
+      return IsVisibleFullCheck(seq);
+    }
+  }
 
   // This is called to determine the maximum visible sequence number for the
   // current transaction for read-your-own-write semantics. This is so that
@@ -25,6 +41,12 @@ class ReadCallback {
   // For other uses, this returns zero, meaning that the current snapshot
   // sequence number is the maximum visible sequence number.
   inline virtual SequenceNumber MaxUnpreparedSequenceNumber() { return 0; };
+
+ protected:
+  // The snapshot at which the read is performed.
+  const SequenceNumber snapshot_ = kMaxSequenceNumber;
+  // Any seq less than min_uncommitted_ is committed.
+  const SequenceNumber min_uncommitted_ = 0;
 };
 
 }  //  namespace rocksdb

utilities/transactions/write_prepared_txn_db.h

@@ -726,18 +726,16 @@ class WritePreparedTxnReadCallback : public ReadCallback {
  public:
   WritePreparedTxnReadCallback(WritePreparedTxnDB* db, SequenceNumber snapshot,
                                SequenceNumber min_uncommitted)
-      : db_(db), snapshot_(snapshot), min_uncommitted_(min_uncommitted) {}
+      : ReadCallback(snapshot, min_uncommitted), db_(db) {}
 
   // Will be called to see if the seq number visible; if not it moves on to
   // the next seq number.
-  inline virtual bool IsVisible(SequenceNumber seq) override {
+  inline virtual bool IsVisibleFullCheck(SequenceNumber seq) override {
     return db_->IsInSnapshot(seq, snapshot_, min_uncommitted_);
   }
 
  private:
   WritePreparedTxnDB* db_;
-  SequenceNumber snapshot_;
-  SequenceNumber min_uncommitted_;
 };
 
 class AddPreparedCallback : public PreReleaseCallback {

utilities/transactions/write_unprepared_txn.cc

@@ -16,7 +16,7 @@
 
 namespace rocksdb {
 
-bool WriteUnpreparedTxnReadCallback::IsVisible(SequenceNumber seq) {
+bool WriteUnpreparedTxnReadCallback::IsVisibleFullCheck(SequenceNumber seq) {
   auto unprep_seqs = txn_->GetUnpreparedSequenceNumbers();
 
   // Since unprep_seqs maps prep_seq => prepare_batch_cnt, to check if seq is
@@ -31,7 +31,7 @@ bool WriteUnpreparedTxnReadCallback::IsVisible(SequenceNumber seq) {
     }
   }
 
-  return db_->IsInSnapshot(seq, snapshot_, min_uncommitted_);
+  return db_->IsInSnapshot(seq, wup_snapshot_, min_uncommitted_);
 }
 
 SequenceNumber WriteUnpreparedTxnReadCallback::MaxUnpreparedSequenceNumber() {

utilities/transactions/write_unprepared_txn.h

@@ -23,19 +23,20 @@ class WriteUnpreparedTxnReadCallback : public ReadCallback {
                                  SequenceNumber snapshot,
                                  SequenceNumber min_uncommitted,
                                  WriteUnpreparedTxn* txn)
-      : db_(db),
-        snapshot_(snapshot),
-        min_uncommitted_(min_uncommitted),
-        txn_(txn) {}
+      // Disable snapshot check on parent class since it would violate
+      // read-your-own-write guarantee.
+      : ReadCallback(kMaxSequenceNumber, min_uncommitted),
+        db_(db),
+        txn_(txn),
+        wup_snapshot_(snapshot) {}
 
-  virtual bool IsVisible(SequenceNumber seq) override;
+  virtual bool IsVisibleFullCheck(SequenceNumber seq) override;
   virtual SequenceNumber MaxUnpreparedSequenceNumber() override;
 
  private:
   WritePreparedTxnDB* db_;
-  SequenceNumber snapshot_;
-  SequenceNumber min_uncommitted_;
   WriteUnpreparedTxn* txn_;
+  SequenceNumber wup_snapshot_;
 };
 
 class WriteUnpreparedTxn : public WritePreparedTxn {

utilities/transactions/write_unprepared_txn_db.cc

@@ -33,11 +33,11 @@ Status WriteUnpreparedTxnDB::RollbackRecoveredTransaction(
    public:
     InvalidSnapshotReadCallback(WritePreparedTxnDB* db, SequenceNumber snapshot,
                                 SequenceNumber min_uncommitted)
-        : db_(db), snapshot_(snapshot), min_uncommitted_(min_uncommitted) {}
+        : ReadCallback(snapshot, min_uncommitted), db_(db) {}
 
     // Will be called to see if the seq number visible; if not it moves on to
     // the next seq number.
-    inline bool IsVisible(SequenceNumber seq) override {
+    inline bool IsVisibleFullCheck(SequenceNumber seq) override {
       // Becomes true if it cannot tell by comparing seq with snapshot seq since
       // the snapshot_ is not a real snapshot.
       bool released = false;
@@ -48,8 +48,6 @@ Status WriteUnpreparedTxnDB::RollbackRecoveredTransaction(
 
    private:
     WritePreparedTxnDB* db_;
-    SequenceNumber snapshot_;
-    SequenceNumber min_uncommitted_;
   };
 
   // Iterate starting with largest sequence number.