Improve RocksDB iterator

lib/util.py

@@ -128,20 +128,13 @@ def int_to_bytes(value):
 
 
 def increment_byte_string(bs):
-    bs = bytearray(bs)
+    '''Return the lexicographically next byte string of the same length.
-    incremented = False
+
-    for i in reversed(range(len(bs))):
+    Return None if there is none (when the input is all 0xff bytes).'''
-        if bs[i] < 0xff:
+    for n in range(1, len(bs) + 1):
-            # This is easy
+        if bs[-n] != 0xff:
-            bs[i] += 1
+            return bs[:-n] + bytes([bs[-n] + 1]) + bytes(n - 1)
-            incremented = True
+    return None
-            break
-        # Otherwise we need to look at the previous character
-        bs[i] = 0
-    if not incremented:
-        # This can only happen if all characters are 0xff
-        bs = bytes([1]) + bs
-    return bytes(bs)
 
 
 class LogicalFile(object):

server/storage.py

@@ -115,7 +115,16 @@ class RocksDB(Storage):
         import gc
         gc.collect()
 
-    class WriteBatch(object):
+    def write_batch(self):
+        return RocksDBWriteBatch(self.db)
+
+    def iterator(self, prefix=b'', reverse=False):
+        return RocksDBIterator(self.db, prefix, reverse)
+
+
+class RocksDBWriteBatch(object):
+    '''A write batch for RocksDB.'''
+
     def __init__(self, db):
         self.batch = RocksDB.module.WriteBatch()
         self.db = db
@@ -127,34 +136,32 @@ class RocksDB(Storage):
         if not exc_val:
             self.db.write(self.batch)
 
-    def write_batch(self):
-        return RocksDB.WriteBatch(self.db)
 
-    class Iterator(object):
+class RocksDBIterator(object):
+    '''An iterator for RocksDB.'''
+
     def __init__(self, db, prefix, reverse):
-            self.it = db.iteritems()
-            self.reverse = reverse
         self.prefix = prefix
-            # Whether we are at the first item
+        if reverse:
-            self.first = True
+            self.iterator = reversed(db.iteritems())
+            nxt_prefix = util.increment_byte_string(prefix)
+            if nxt_prefix:
+                self.iterator.seek(nxt_prefix)
+                try:
+                    next(self.iterator)
+                except StopIteration:
+                    self.iterator.seek(nxt_prefix)
+            else:
+                self.iterator.seek_to_last()
+        else:
+            self.iterator = db.iteritems()
+            self.iterator.seek(prefix)
 
     def __iter__(self):
-            prefix = self.prefix
-            if self.reverse:
-                prefix = increment_byte_string(prefix)
-                self.it = reversed(self.it)
-            self.it.seek(prefix)
         return self
 
     def __next__(self):
-            k, v = self.it.__next__()
+        k, v = next(self.iterator)
-            if self.first and self.reverse and not k.startswith(self.prefix):
-                k, v = self.it.__next__()
-            self.first = False
         if not k.startswith(self.prefix):
-                # We're already ahead of the prefix
             raise StopIteration
         return k, v
-
-    def iterator(self, prefix=b'', reverse=False):
-        return RocksDB.Iterator(self.db, prefix, reverse)