Implementing DynamicIterator for TransformRepNoLock

Summary: What @haobo done with TransformRep, now in TransformRepNoLock. Similar implementation, except that I made DynamicIterator a subclass of Iterator which makes me have less iterator initializations.

Test Plan: ./prefix_test. Seeing huge savings vs. TransformRep again!

Reviewers: dhruba, haobo, sdong, kailiu

Reviewed By: haobo

CC: leveldb, haobo

Differential Revision: https://reviews.facebook.net/D13953

db/prefix_test.cc

@@ -93,9 +93,8 @@ class PrefixTest {
       auto prefix_extractor = NewFixedPrefixTransform(8);
       options.prefix_extractor = prefix_extractor;
       if (FLAGS_use_nolock_version) {
-        options.memtable_factory =
-          std::make_shared<rocksdb::PrefixHashRepNoLockFactory>(
-            prefix_extractor, FLAGS_bucket_count);
+        options.memtable_factory.reset(NewHashSkipListRepFactory(
+            prefix_extractor, FLAGS_bucket_count));
       } else {
         options.memtable_factory =
           std::make_shared<rocksdb::PrefixHashRepFactory>(

db/skiplist.h

@@ -63,6 +63,11 @@ class SkipList {
     // The returned iterator is not valid.
     explicit Iterator(const SkipList* list);
 
+    // Change the underlying skiplist used for this iterator
+    // This enables us not changing the iterator without deallocating
+    // an old one and then allocating a new one
+    void SetList(const SkipList* list);
+
     // Returns true iff the iterator is positioned at a valid node.
     bool Valid() const;
 
@@ -194,6 +199,11 @@ SkipList<Key,Comparator>::NewNode(const Key& key, int height) {
 
 template<typename Key, class Comparator>
 inline SkipList<Key,Comparator>::Iterator::Iterator(const SkipList* list) {
+  SetList(list);
+}
+
+template<typename Key, class Comparator>
+inline void SkipList<Key,Comparator>::Iterator::SetList(const SkipList* list) {
   list_ = list;
   node_ = nullptr;
 }

include/rocksdb/memtablerep.h

@@ -260,67 +260,11 @@ public:
   }
 };
 
-
-// NO LOCKS VERSION
-
 // The same as TransformRepFactory except it doesn't use locks.
 // Experimental, will replace TransformRepFactory once we are sure
 // it performs better
-class TransformRepNoLockFactory : public MemTableRepFactory {
- public:
-  explicit TransformRepNoLockFactory(const SliceTransform* transform,
-    size_t bucket_count)
-    : transform_(transform),
-      bucket_count_(bucket_count) { }
-
-  virtual ~TransformRepNoLockFactory() { delete transform_; }
-
-  virtual std::shared_ptr<MemTableRep> CreateMemTableRep(
-    MemTableRep::KeyComparator&, Arena*) override;
-
-  virtual const char* Name() const override {
-    return "TransformRepNoLockFactory";
-  }
-
-  const SliceTransform* GetTransform() { return transform_; }
-
- protected:
-  const SliceTransform* transform_;
-  const size_t bucket_count_;
-};
-
-// UnsortedReps bin user keys based on an identity function transform -- that
-// is, transform(key) = key. This optimizes for point look-ups.
-//
-// Parameters: See TransformRepNoLockFactory.
-class UnsortedRepNoLockFactory : public TransformRepNoLockFactory {
-public:
-  explicit UnsortedRepNoLockFactory(size_t bucket_count = 1000000)
-    : TransformRepNoLockFactory(NewNoopTransform(),
-                                bucket_count) { }
-  virtual const char* Name() const override {
-    return "UnsortedRepNoLockFactory";
-  }
-};
-
-// PrefixHashReps bin user keys based on a fixed-size prefix. This optimizes for
-// short ranged scans over a given prefix.
-//
-// Parameters: See TransformRepNoLockFactory.
-class PrefixHashRepNoLockFactory : public TransformRepNoLockFactory {
-public:
-  explicit PrefixHashRepNoLockFactory(const SliceTransform* prefix_extractor,
-    size_t bucket_count = 1000000)
-    : TransformRepNoLockFactory(prefix_extractor, bucket_count)
-    { }
-
-  virtual std::shared_ptr<MemTableRep> CreateMemTableRep(
-    MemTableRep::KeyComparator&, Arena*) override;
-
-  virtual const char* Name() const override {
-    return "PrefixHashRepNoLockFactory";
-  }
-};
+extern MemTableRepFactory* NewHashSkipListRepFactory(
+    const SliceTransform* transform, size_t bucket_count = 1000000);
 
 }
 

util/transformrepnolock.cc → util/hash_skiplist_rep.cc

@@ -16,9 +16,9 @@
 namespace rocksdb {
 namespace {
 
-class TransformRepNoLock : public MemTableRep {
+class HashSkipListRep : public MemTableRep {
  public:
-  TransformRepNoLock(MemTableRep::KeyComparator& compare, Arena* arena,
+  HashSkipListRep(MemTableRep::KeyComparator& compare, Arena* arena,
     const SliceTransform* transform, size_t bucket_size);
 
   virtual void Insert(const char* key) override;
@@ -27,17 +27,21 @@ class TransformRepNoLock : public MemTableRep {
 
   virtual size_t ApproximateMemoryUsage() override;
 
-  virtual ~TransformRepNoLock();
+  virtual ~HashSkipListRep();
 
   virtual std::shared_ptr<MemTableRep::Iterator> GetIterator() override;
 
   virtual std::shared_ptr<MemTableRep::Iterator> GetIterator(
-    const Slice& slice) override;
+      const Slice& slice) override;
 
-  std::shared_ptr<MemTableRep::Iterator> GetTransformIterator(
-    const Slice& transformed);
+  virtual std::shared_ptr<MemTableRep::Iterator> GetPrefixIterator(
+      const Slice& prefix) override;
+
+  virtual std::shared_ptr<MemTableRep::Iterator> GetDynamicPrefixIterator()
+      override;
 
  private:
+  friend class DynamicIterator;
   typedef SkipList<const char*, MemTableRep::KeyComparator&> Bucket;
 
   size_t bucket_size_;
@@ -76,50 +80,72 @@ class TransformRepNoLock : public MemTableRep {
     virtual ~Iterator() {
       // if we own the list, we should also delete it
       if (own_list_) {
+        assert(list_ != nullptr);
         delete list_;
       }
-    };
+    }
 
     // Returns true iff the iterator is positioned at a valid node.
     virtual bool Valid() const {
-      return iter_.Valid();
+      return list_ != nullptr && iter_.Valid();
     }
 
     // Returns the key at the current position.
     // REQUIRES: Valid()
     virtual const char* key() const {
+      assert(Valid());
       return iter_.key();
     }
 
     // Advances to the next position.
     // REQUIRES: Valid()
     virtual void Next() {
+      assert(Valid());
       iter_.Next();
     }
 
     // Advances to the previous position.
     // REQUIRES: Valid()
     virtual void Prev() {
+      assert(Valid());
       iter_.Prev();
     }
 
     // Advance to the first entry with a key >= target
     virtual void Seek(const char* target) {
-      iter_.Seek(target);
+      if (list_ != nullptr) {
+        iter_.Seek(target);
+      }
     }
 
     // Position at the first entry in collection.
     // Final state of iterator is Valid() iff collection is not empty.
     virtual void SeekToFirst() {
-      iter_.SeekToFirst();
+      if (list_ != nullptr) {
+        iter_.SeekToFirst();
+      }
     }
 
     // Position at the last entry in collection.
     // Final state of iterator is Valid() iff collection is not empty.
     virtual void SeekToLast() {
-      iter_.SeekToLast();
+      if (list_ != nullptr) {
+        iter_.SeekToLast();
+      }
+    }
+   protected:
+    void Reset(Bucket* list) {
+      if (own_list_) {
+        assert(list_ != nullptr);
+        delete list_;
+      }
+      list_ = list;
+      iter_.SetList(list);
+      own_list_ = false;
     }
    private:
+    // if list_ is nullptr, we should NEVER call any methods on iter_
+    // if list_ is nullptr, this Iterator is not Valid()
     Bucket* list_;
     Bucket::Iterator iter_;
     // here we track if we own list_. If we own it, we are also
@@ -127,6 +153,40 @@ class TransformRepNoLock : public MemTableRep {
     bool own_list_;
   };
 
+  class DynamicIterator : public HashSkipListRep::Iterator {
+   public:
+    explicit DynamicIterator(const HashSkipListRep& memtable_rep)
+      : HashSkipListRep::Iterator(nullptr, false),
+        memtable_rep_(memtable_rep) {}
+
+    // Advance to the first entry with a key >= target
+    virtual void Seek(const char* target) {
+      auto transformed = memtable_rep_.transform_->Transform(
+        memtable_rep_.UserKey(target));
+      Reset(memtable_rep_.GetBucket(transformed));
+      HashSkipListRep::Iterator::Seek(target);
+    }
+
+    // Position at the first entry in collection.
+    // Final state of iterator is Valid() iff collection is not empty.
+    virtual void SeekToFirst() {
+      // Prefix iterator does not support total order.
+      // We simply set the iterator to invalid state
+      Reset(nullptr);
+    }
+
+    // Position at the last entry in collection.
+    // Final state of iterator is Valid() iff collection is not empty.
+    virtual void SeekToLast() {
+      // Prefix iterator does not support total order.
+      // We simply set the iterator to invalid state
+      Reset(nullptr);
+    }
+   private:
+    // the underlying memtable
+    const HashSkipListRep& memtable_rep_;
+  };
+
   class EmptyIterator : public MemTableRep::Iterator {
     // This is used when there wasn't a bucket. It is cheaper than
     // instantiating an empty bucket over which to iterate.
@@ -150,17 +210,7 @@ class TransformRepNoLock : public MemTableRep {
   std::shared_ptr<EmptyIterator> empty_iterator_;
 };
 
-class PrefixHashRepNoLock : public TransformRepNoLock {
- public:
-  PrefixHashRepNoLock(MemTableRep::KeyComparator& compare, Arena* arena,
-    const SliceTransform* transform, size_t bucket_size)
-  : TransformRepNoLock(compare, arena, transform, bucket_size) { }
-
-  virtual std::shared_ptr<MemTableRep::Iterator> GetPrefixIterator(
-    const Slice& prefix) override;
-};
-
-TransformRepNoLock::TransformRepNoLock(MemTableRep::KeyComparator& compare,
+HashSkipListRep::HashSkipListRep(MemTableRep::KeyComparator& compare,
     Arena* arena, const SliceTransform* transform, size_t bucket_size)
   : bucket_size_(bucket_size),
     transform_(transform),
@@ -175,11 +225,11 @@ TransformRepNoLock::TransformRepNoLock(MemTableRep::KeyComparator& compare,
   }
 }
 
-TransformRepNoLock::~TransformRepNoLock() {
+HashSkipListRep::~HashSkipListRep() {
   delete[] buckets_;
 }
 
-TransformRepNoLock::Bucket* TransformRepNoLock::GetInitializedBucket(
+HashSkipListRep::Bucket* HashSkipListRep::GetInitializedBucket(
     const Slice& transformed) {
   size_t hash = GetHash(transformed);
   auto bucket = GetBucket(hash);
@@ -191,14 +241,14 @@ TransformRepNoLock::Bucket* TransformRepNoLock::GetInitializedBucket(
   return bucket;
 }
 
-void TransformRepNoLock::Insert(const char* key) {
+void HashSkipListRep::Insert(const char* key) {
   assert(!Contains(key));
   auto transformed = transform_->Transform(UserKey(key));
   auto bucket = GetInitializedBucket(transformed);
   bucket->Insert(key);
 }
 
-bool TransformRepNoLock::Contains(const char* key) const {
+bool HashSkipListRep::Contains(const char* key) const {
   auto transformed = transform_->Transform(UserKey(key));
   auto bucket = GetBucket(transformed);
   if (bucket == nullptr) {
@@ -207,11 +257,11 @@ bool TransformRepNoLock::Contains(const char* key) const {
   return bucket->Contains(key);
 }
 
-size_t TransformRepNoLock::ApproximateMemoryUsage() {
+size_t HashSkipListRep::ApproximateMemoryUsage() {
   return sizeof(buckets_);
 }
 
-std::shared_ptr<MemTableRep::Iterator> TransformRepNoLock::GetIterator() {
+std::shared_ptr<MemTableRep::Iterator> HashSkipListRep::GetIterator() {
   auto list = new Bucket(compare_, arena_);
   for (size_t i = 0; i < bucket_size_; ++i) {
     auto bucket = GetBucket(i);
@@ -225,38 +275,56 @@ std::shared_ptr<MemTableRep::Iterator> TransformRepNoLock::GetIterator() {
   return std::make_shared<Iterator>(list);
 }
 
-std::shared_ptr<MemTableRep::Iterator> TransformRepNoLock::GetTransformIterator(
-  const Slice& transformed) {
-  auto bucket = GetBucket(transformed);
+std::shared_ptr<MemTableRep::Iterator> HashSkipListRep::GetPrefixIterator(
+  const Slice& prefix) {
+  auto bucket = GetBucket(prefix);
   if (bucket == nullptr) {
     return empty_iterator_;
   }
   return std::make_shared<Iterator>(bucket, false);
 }
 
-std::shared_ptr<MemTableRep::Iterator> TransformRepNoLock::GetIterator(
-  const Slice& slice) {
-  auto transformed = transform_->Transform(slice);
-  return GetTransformIterator(transformed);
+std::shared_ptr<MemTableRep::Iterator> HashSkipListRep::GetIterator(
+    const Slice& slice) {
+  return GetPrefixIterator(transform_->Transform(slice));
+}
+
+std::shared_ptr<MemTableRep::Iterator>
+    HashSkipListRep::GetDynamicPrefixIterator() {
+  return std::make_shared<DynamicIterator>(*this);
 }
 
 } // anon namespace
 
-std::shared_ptr<MemTableRep> TransformRepNoLockFactory::CreateMemTableRep(
-  MemTableRep::KeyComparator& compare, Arena* arena) {
-  return std::make_shared<TransformRepNoLock>(compare, arena, transform_,
-    bucket_count_);
-}
+class HashSkipListRepFactory : public MemTableRepFactory {
+ public:
+  explicit HashSkipListRepFactory(const SliceTransform* transform,
+      size_t bucket_count = 1000000)
+    : transform_(transform),
+      bucket_count_(bucket_count) { }
 
-std::shared_ptr<MemTableRep> PrefixHashRepNoLockFactory::CreateMemTableRep(
-  MemTableRep::KeyComparator& compare, Arena* arena) {
-  return std::make_shared<PrefixHashRepNoLock>(compare, arena, transform_,
-    bucket_count_);
-}
+  virtual ~HashSkipListRepFactory() { delete transform_; }
 
-std::shared_ptr<MemTableRep::Iterator> PrefixHashRepNoLock::GetPrefixIterator(
-  const Slice& prefix) {
-  return TransformRepNoLock::GetTransformIterator(prefix);
+  virtual std::shared_ptr<MemTableRep> CreateMemTableRep(
+      MemTableRep::KeyComparator& compare, Arena* arena) override {
+    return std::make_shared<HashSkipListRep>(compare, arena, transform_,
+        bucket_count_);
+  }
+
+  virtual const char* Name() const override {
+    return "HashSkipListRepFactory";
+  }
+
+  const SliceTransform* GetTransform() { return transform_; }
+
+ private:
+  const SliceTransform* transform_;
+  const size_t bucket_count_;
+};
+
+MemTableRepFactory* NewHashSkipListRepFactory(
+    const SliceTransform* transform, size_t bucket_count) {
+  return new HashSkipListRepFactory(transform, bucket_count);
 }
 
 } // namespace rocksdb