Allow snapshot sequence retrieval from db iterator (#285)

This commit adds utility to DB to extract a sequence number from a passed in
iterator, if possible.

Author: danielin917

db/arena_wrapped_db_iter.h

@@ -99,6 +99,13 @@ class ArenaWrappedDBIter : public Iterator {
     expose_blob_index_ = expose_blob_index;
   }
 
+  SequenceNumber get_sequence() const {
+    if (!db_iter_) {
+      return SequenceNumber(0);
+    }
+    return db_iter_->get_sequence();
+  }
+
  private:
   DBIter* db_iter_ = nullptr;
   Arena arena_;

db/db_impl/db_impl.cc

@@ -4086,6 +4086,14 @@ Status DBImpl::NewIterators(
   return Status::OK();
 }
 
+SequenceNumber DBImpl::GetIteratorSequenceNumber(Iterator* it) {
+  auto arena_wrapped = dynamic_cast<ArenaWrappedDBIter*>(it);
+  if (!arena_wrapped) {
+    return SequenceNumber(0);
+  }
+  return arena_wrapped->get_sequence();
+}
+
 const Snapshot* DBImpl::GetSnapshot() { return GetSnapshotImpl(false); }
 
 #ifndef ROCKSDB_LITE

db/db_impl/db_impl.h

@@ -342,6 +342,8 @@ class DBImpl : public DB {
       const std::vector<ColumnFamilyHandle*>& column_families,
       std::vector<Iterator*>* iterators) override;
 
+  SequenceNumber GetIteratorSequenceNumber(Iterator* it) override;
+
   virtual const Snapshot* GetSnapshot() override;
   virtual void ReleaseSnapshot(const Snapshot* snapshot) override;
   // Create a timestamped snapshot. This snapshot can be shared by multiple

db/db_iter.h

@@ -210,6 +210,8 @@ class DBIter final : public Iterator {
       read_callback_->Refresh(s);
     }
   }
+  SequenceNumber get_sequence() const { return sequence_; }
+
   void set_valid(bool v) { valid_ = v; }
 
  private:

db/db_iterator_test.cc

@@ -815,8 +815,28 @@ TEST_P(DBIteratorTest, IterWithSnapshot) {
         ASSERT_TRUE(!iter->Valid());
       }
     }
+
+    // Check that creating a new iterator snapshot has an increasing and
+    // different sequence number.
+    const Snapshot* snapshot2 = db_->GetSnapshot();
+    ReadOptions options2;
+    options.snapshot = snapshot2;
+    Iterator* iter2 = NewIterator(options2, handles_[1]);
+    ASSERT_GT(db_->GetIteratorSequenceNumber(iter2),
+              db_->GetIteratorSequenceNumber(iter));
+
+    // Now check only the new iterator has the value added after the initial
+    // snapshot.
+    iter->Seek("key0");
+    ASSERT_TRUE(!iter->Valid() || iter->key() != "key0");
+
+    iter2->Seek("key0");
+    ASSERT_EQ(IterStatus(iter2), "key0->val0");
+
     db_->ReleaseSnapshot(snapshot);
+    db_->ReleaseSnapshot(snapshot2);
     delete iter;
+    delete iter2;
   } while (ChangeOptions());
 }
 

db/db_test.cc

@@ -3183,6 +3183,13 @@ class ModelDB : public DB {
                       std::vector<Iterator*>* /*iterators*/) override {
     return Status::NotSupported("Not supported yet");
   }
+
+  SequenceNumber GetIteratorSequenceNumber(Iterator*) override {
+    // No support yet.
+    assert(false);
+    return 0;
+  }
+
   const Snapshot* GetSnapshot() override {
     ModelSnapshot* snapshot = new ModelSnapshot;
     snapshot->map_ = map_;

include/rocksdb/db.h

@@ -836,6 +836,9 @@ class DB {
       const std::vector<ColumnFamilyHandle*>& column_families,
       std::vector<Iterator*>* iterators) = 0;
 
+  // Returns the iterator sequence number
+  virtual SequenceNumber GetIteratorSequenceNumber(Iterator* it) = 0;
+
   // Return a handle to the current DB state.  Iterators created with
   // this handle will all observe a stable snapshot of the current DB
   // state.  The caller must call ReleaseSnapshot(result) when the

include/rocksdb/utilities/stackable_db.h

@@ -237,6 +237,12 @@ class StackableDB : public DB {
     return db_->NewIterators(options, column_families, iterators);
   }
 
+  using DB::GetIteratorSequenceNumber;
+  virtual SequenceNumber GetIteratorSequenceNumber(
+      Iterator* iterator) override {
+    return db_->GetIteratorSequenceNumber(iterator);
+  }
+
   virtual const Snapshot* GetSnapshot() override { return db_->GetSnapshot(); }
 
   virtual void ReleaseSnapshot(const Snapshot* snapshot) override {