Repository: petermattis/pebble (go)

Issue:
db: support changing {Lower,Upper}Bound on an iterator
CockroachDB likes to specify tight lower and upper bounds when iterating with an iterator, but in order to maintain snapshot semantics, it wants to reuse the same iterator across multiple operations within the same batch of operations. It accomplishes this with RocksDB by abusing the iterator API. `ReadOptions::iterate_{lower,upper}_bound` are configured to point to valid slices, and CockroachDB changes the contents of those slices before calls to `Seek`.

A better API to accomplish this would be to provide an `Iterator.SetBounds` or `Iterator.SetOptions` method. This would likely necessitate an `internalIterator.SetBounds` method. None of this seems particularly complex. I think the semantics of `Iterator.SetBounds` should be that it doesn't change the current state of the iterator. If the iterator was valid before the call it will remain valid after the `SetBounds` call and positioned at the same entry. If the iterator was invalid before the call it will remain invalid after the `SetBounds` call. The new bounds only apply to subsequent positioning operations (`Seek{GE,LT}`, `First`, `Last`, `Next`, `Prev`). If these semantics are burdensome, we could have `SetBounds` always invalidate the iterator.

Cc @ajkr, @Ryanfsdf 

Pull request description:
Add support for dynamic bounds
Add support for dynamically setting `Lower,Upper{Bound}` after an iterator has already been created. This is useful for setting the bounds on the fly before calling `SeekGE`, `SeekPrefixGE`, `SeekLT`, `First`, and `Last`, without having to recreate the iterator.

Fixes #143.

Patch:
diff --git a/batch.go b/batch.go
index c584c94..0988639 100644
--- a/batch.go
+++ b/batch.go
@@ -821,6 +821,10 @@ func (i *batchIter) Close() error {
 	return i.err
 }
 
+func (i *batchIter) SetBounds(lower, upper []byte) {
+	i.iter.SetBounds(lower, upper)
+}
+
 type flushableBatchEntry struct {
 	offset   uint32
 	index    uint32
@@ -1100,3 +1104,8 @@ func (i *flushableBatchIter) Error() error {
 func (i *flushableBatchIter) Close() error {
 	return i.err
 }
+
+func (i *flushableBatchIter) SetBounds(lower, upper []byte) {
+	// This should not be called as bounds are not used for this iterator.
+	panic("TODO(ryan): pebble: SetBounds unimplemented")
+}
diff --git a/db.go b/db.go
index ad72892..6f1e522 100644
--- a/db.go
+++ b/db.go
@@ -456,12 +456,14 @@ func (d *DB) newIterInternal(
 	buf := iterAllocPool.Get().(*iterAlloc)
 	dbi := &buf.dbi
 	dbi.alloc = buf
-	dbi.opts = o
 	dbi.cmp = d.cmp
 	dbi.equal = d.equal
 	dbi.merge = d.merge
 	dbi.split = d.split
 	dbi.readState = readState
+	if o != nil {
+		dbi.opts = *o
+	}
 
 	iters := buf.iters[:0]
 	rangeDelIters := buf.rangeDelIters[:0]
@@ -478,8 +480,8 @@ func (d *DB) newIterInternal(
 	memtables := readState.memtables
 	for i := len(memtables) - 1; i >= 0; i-- {
 		mem := memtables[i]
-		iters = append(iters, mem.newIter(o))
-		rangeDelIters = append(rangeDelIters, mem.newRangeDelIter(o))
+		iters = append(iters, mem.newIter(&dbi.opts))
+		rangeDelIters = append(rangeDelIters, mem.newRangeDelIter(&dbi.opts))
 		largestUserKeys = append(largestUserKeys, nil)
 	}
 
@@ -487,7 +489,7 @@ func (d *DB) newIterInternal(
 	current := readState.current
 	for i := len(current.files[0]) - 1; i >= 0; i-- {
 		f := &current.files[0][i]
-		iter, rangeDelIter, err := d.newIters(f, o)
+		iter, rangeDelIter, err := d.newIters(f, &dbi.opts)
 		if err != nil {
 			dbi.err = err
 			return dbi
@@ -525,7 +527,7 @@ func (d *DB) newIterInternal(
 			li = &levelIter{}
 		}
 
-		li.init(o, d.cmp, d.newIters, current.files[level])
+		li.init(&dbi.opts, d.cmp, d.newIters, current.files[level])
 		li.initRangeDel(&rangeDelIters[0])
 		li.initLargestUserKey(&largestUserKeys[0])
 		iters = append(iters, li)
diff --git a/error_iter.go b/error_iter.go
index 23d1d00..8bf5afc 100644
--- a/error_iter.go
+++ b/error_iter.go
@@ -61,3 +61,5 @@ func (c *errorIter) Error() error {
 func (c *errorIter) Close() error {
 	return c.err
 }
+
+func (c *errorIter) SetBounds(lower, upper []byte) {}
diff --git a/get_iter.go b/get_iter.go
index 3197f9c..c109b75 100644
--- a/get_iter.go
+++ b/get_iter.go
@@ -188,3 +188,7 @@ func (g *getIter) Close() error {
 	}
 	return g.err
 }
+
+func (g *getIter) SetBounds(lower, upper []byte) {
+	panic("pebble: SetBounds unimplemented")
+}
diff --git a/internal.go b/internal.go
index 280d6f4..99845b7 100644
--- a/internal.go
+++ b/internal.go
@@ -107,6 +107,11 @@ type internalIterator interface {
 	// It is valid to call Close multiple times. Other methods should not be
 	// called after the iterator has been closed.
 	Close() error
+
+	// SetBounds sets the lower and upper bounds for the iterator. Note that the
+	// result of Next and Prev will be undefined until the iterator has been
+	// repositioned with SeekGE, SeekPrefixGE, SeekLT, First, or Last.
+	SetBounds(lower, upper []byte)
 }
 
 // sstable.Iterator implements the internalIterator interface.
diff --git a/internal/arenaskl/iterator.go b/internal/arenaskl/iterator.go
index ff83d92..589176a 100644
--- a/internal/arenaskl/iterator.go
+++ b/internal/arenaskl/iterator.go
@@ -196,6 +196,14 @@ func (it *Iterator) Valid() bool {
 	return it.nd != it.list.head && it.nd != it.list.tail
 }
 
+// SetBounds sets the lower and upper bounds for the iterator. Note that the
+// result of Next and Prev will be undefined until the iterator has been
+// repositioned with SeekGE, SeekPrefixGE, SeekLT, First, or Last.
+func (it *Iterator) SetBounds(lower, upper []byte) {
+	it.lower = lower
+	it.upper = upper
+}
+
 func (it *Iterator) decodeKey() {
 	b := it.list.arena.getBytes(it.nd.keyOffset, it.nd.keySize)
 	// This is a manual inline of base.DecodeInternalKey, because the Go compiler
diff --git a/internal/batchskl/iterator.go b/internal/batchskl/iterator.go
index 32fb1b4..a01bf17 100644
--- a/internal/batchskl/iterator.go
+++ b/internal/batchskl/iterator.go
@@ -177,6 +177,14 @@ func (it *Iterator) Valid() bool {
 	return it.list != nil && it.nd != it.list.head && it.nd != it.list.tail
 }
 
+// SetBounds sets the lower and upper bounds for the iterator. Note that the
+// result of Next and Prev will be undefined until the iterator has been
+// repositioned with SeekGE, SeekLT, First, or Last.
+func (it *Iterator) SetBounds(lower, upper []byte) {
+	it.lower = lower
+	it.upper = upper
+}
+
 func (it *Iterator) seekForBaseSplice(
 	key []byte, abbreviatedKey uint64,
 ) (prev, next uint32, found bool) {
diff --git a/internal/rangedel/iter.go b/internal/rangedel/iter.go
index d022397..b7daecc 100644
--- a/internal/rangedel/iter.go
+++ b/internal/rangedel/iter.go
@@ -166,3 +166,10 @@ func (i *Iter) Error() error {
 func (i *Iter) Close() error {
 	return nil
 }
+
+// SetBounds implements internalIterator.SetBounds, as documented in the pebble
+// package.
+func (i *Iter) SetBounds(lower, upper []byte) {
+	// This should never be called as bounds are only used for point records.
+	panic("pebble: SetBounds unimplemented")
+}
diff --git a/iterator.go b/iterator.go
index 3033744..ad32df8 100644
--- a/iterator.go
+++ b/iterator.go
@@ -31,7 +31,7 @@ const (
 // key/value pairs are not guaranteed to be a consistent snapshot of that DB
 // at a particular point in time.
 type Iterator struct {
-	opts      *IterOptions
+	opts      IterOptions
 	cmp       Compare
 	equal     Equal
 	merge     Merge
@@ -450,3 +450,18 @@ func (i *Iterator) Close() error {
 	}
 	return err
 }
+
+// SetBounds sets the lower and upper bounds for the iterator. Note that the
+// iterator will always be invalidated and must be repositioned with a call to
+// SeekGE, SeekPrefixGE, SeekLT, First, or Last.
+func (i* Iterator) SetBounds(lower, upper []byte) {
+	i.prefix = nil
+	i.iterKey = nil
+	i.iterValue = nil
+	i.pos = iterPosCur
+	i.valid = false
+
+	i.opts.LowerBound = lower
+	i.opts.UpperBound = upper
+	i.iter.SetBounds(lower, upper)
+}
diff --git a/level_iter.go b/level_iter.go
index dd79f40..cd0ae4e 100644
--- a/level_iter.go
+++ b/level_iter.go
@@ -411,3 +411,11 @@ func (l *levelIter) Close() error {
 	}
 	return l.err
 }
+
+func (l *levelIter) SetBounds(lower, upper []byte) {
+	l.opts.LowerBound = lower
+	l.opts.UpperBound = upper
+	if l.iter != nil {
+		l.iter.SetBounds(lower, upper)
+	}
+}
diff --git a/merging_iter.go b/merging_iter.go
index 2872dc8..17a902c 100644
--- a/merging_iter.go
+++ b/merging_iter.go
@@ -676,6 +676,12 @@ func (m *mergingIter) Close() error {
 	return m.err
 }
 
+func (m *mergingIter) SetBounds(lower, upper []byte) {
+	for _, iter := range m.iters {
+		iter.SetBounds(lower, upper)
+	}
+}
+
 func (m *mergingIter) DebugString() string {
 	var buf bytes.Buffer
 	sep := ""
diff --git a/sstable/block.go b/sstable/block.go
index 4eef8ea..fd3c285 100644
--- a/sstable/block.go
+++ b/sstable/block.go
@@ -644,6 +644,11 @@ func (i *blockIter) Close() error {
 	return i.err
 }
 
+func (i *blockIter) SetBounds(lower, upper []byte) {
+	// This should never be called as bounds are handled by sstable.Iterator.
+	panic("pebble: SetBounds unimplemented")
+}
+
 // invalidate the iterator, positioning it below the first entry.
 func (i *blockIter) invalidateLower() {
 	i.offset = -1
diff --git a/sstable/reader.go b/sstable/reader.go
index 686a126..8bf30ab 100644
--- a/sstable/reader.go
+++ b/sstable/reader.go
@@ -459,6 +459,13 @@ func (i *Iterator) Close() error {
 	return err
 }
 
+// SetBounds implements internalIterator.SetBounds, as documented in the pebble
+// package.
+func (i *Iterator) SetBounds(lower, upper []byte) {
+	i.lower = lower
+	i.upper = upper
+}
+
 type weakCachedBlock struct {
 	bh     blockHandle
 	mu     sync.RWMutex


Test patch:
diff --git a/data_test.go b/data_test.go
index 5e673c5..9a68a26 100644
--- a/data_test.go
+++ b/data_test.go
@@ -57,6 +57,25 @@ func runIterCmd(d *datadriven.TestData, iter *Iterator) string {
 			valid = iter.Next()
 		case "prev":
 			valid = iter.Prev()
+		case "set-bounds":
+			if len(parts) <= 1 || len(parts) > 3 {
+				return fmt.Sprintf("set-bounds lower=<lower> upper=<upper>\n")
+			}
+			var lower []byte
+			var upper []byte
+			for _, part := range parts[1:] {
+				arg := strings.Split(strings.TrimSpace(part), "=")
+				switch arg[0] {
+				case "lower":
+					lower = []byte(arg[1])
+				case "upper":
+					upper = []byte(arg[1])
+				default:
+					return fmt.Sprintf("set-bounds: unknown arg: %s", arg)
+				}
+			}
+			iter.SetBounds(lower, upper)
+			valid = iter.Valid()
 		default:
 			return fmt.Sprintf("unknown op: %s", parts[0])
 		}
@@ -117,6 +136,25 @@ func runInternalIterCmd(d *datadriven.TestData, iter internalIterator, opts ...i
 			iter.Next()
 		case "prev":
 			iter.Prev()
+		case "set-bounds":
+			if len(parts) <= 1 || len(parts) > 3 {
+				return fmt.Sprintf("set-bounds lower=<lower> upper=<upper>\n")
+			}
+			var lower []byte
+			var upper []byte
+			for _, part := range parts[1:] {
+				arg := strings.Split(strings.TrimSpace(part), "=")
+				switch arg[0] {
+				case "lower":
+					lower = []byte(arg[1])
+				case "upper":
+					upper = []byte(arg[1])
+				default:
+					return fmt.Sprintf("set-bounds: unknown arg: %s", arg)
+				}
+			}
+			iter.SetBounds(lower, upper)
+			continue
 		default:
 			return fmt.Sprintf("unknown op: %s", parts[0])
 		}
diff --git a/iterator_test.go b/iterator_test.go
index 9e50bf1..3408a75 100644
--- a/iterator_test.go
+++ b/iterator_test.go
@@ -179,6 +179,11 @@ func (f *fakeIter) Close() error {
 	return f.closeErr
 }
 
+func (f *fakeIter) SetBounds(lower, upper []byte) {
+	f.lower = lower
+	f.upper = upper
+}
+
 // testIterator tests creating a combined iterator from a number of sub-
 // iterators. newFunc is a constructor function. splitFunc returns a random
 // split of the testKeyValuePairs slice such that walking a combined iterator
@@ -296,7 +301,7 @@ func TestIterator(t *testing.T) {
 	var keys []InternalKey
 	var vals [][]byte
 
-	newIter := func(seqNum uint64, opts *IterOptions) *Iterator {
+	newIter := func(seqNum uint64, opts IterOptions) *Iterator {
 		cmp := DefaultComparer.Compare
 		equal := DefaultComparer.Equal
 		split := func(a []byte) int { return len(a) }
@@ -354,7 +359,7 @@ func TestIterator(t *testing.T) {
 				}
 			}
 
-			iter := newIter(uint64(seqNum), &opts)
+			iter := newIter(uint64(seqNum), opts)
 			defer iter.Close()
 			return runIterCmd(d, iter)
 
diff --git a/testdata/iterator b/testdata/iterator
index 367ad05..f7b0ff6 100644
--- a/testdata/iterator
+++ b/testdata/iterator
@@ -827,6 +827,173 @@ next
 b:b
 .
 
+iter seq=2
+set-bounds lower=a
+seek-ge a
+first
+prev
+----
+.
+a:a
+a:a
+.
+
+iter seq=2
+set-bounds lower=b
+seek-ge a
+first
+prev
+----
+.
+b:b
+b:b
+.
+
+iter seq=2
+set-bounds lower=c
+seek-ge a
+first
+prev
+----
+.
+c:c
+c:c
+.
+
+iter seq=2
+set-bounds lower=d
+seek-ge a
+first
+prev
+----
+.
+d:d
+d:d
+.
+
+iter seq=2
+set-bounds lower=e
+seek-ge a
+first
+----
+.
+.
+.
+
+iter seq=2
+set-bounds upper=d
+seek-lt d
+last
+next
+----
+.
+c:c
+c:c
+.
+
+iter seq=2
+set-bounds upper=c
+seek-lt d
+last
+next
+----
+.
+b:b
+b:b
+.
+
+iter seq=2
+set-bounds upper=b
+seek-lt d
+last
+next
+----
+.
+a:a
+a:a
+.
+
+iter seq=2
+set-bounds upper=a
+seek-lt d
+last
+----
+.
+.
+.
+
+iter seq=2
+set-bounds lower=a
+seek-lt d
+next
+next
+----
+.
+c:c
+d:d
+.
+
+iter seq=2
+set-bounds lower=b upper=c
+seek-ge a
+next
+----
+.
+b:b
+.
+
+iter seq=2
+set-bounds lower=b
+seek-ge a
+set-bounds lower=b
+----
+.
+b:b
+.
+
+iter seq=2
+seek-ge a
+set-bounds upper=e
+----
+a:a
+.
+
+iter seq=2
+set-bounds lower=b
+seek-ge a
+set-bounds upper=e
+----
+.
+b:b
+.
+
+iter seq=2
+set-bounds lower=b
+first
+----
+.
+b:b
+
+iter seq=2
+set-bounds upper=b
+first
+----
+.
+a:a
+
+iter seq=2
+set-bounds lower=b
+last
+----
+.
+d:d
+
+iter seq=2
+set-bounds upper=b
+last
+----
+.
+a:a
 
 define
 a.SET.1:a
diff --git a/testdata/level_iter b/testdata/level_iter
index 3c17c13..62b43f7 100644
--- a/testdata/level_iter
+++ b/testdata/level_iter
@@ -166,6 +166,14 @@ first
 a:1
 a:1
 
+iter
+set-bounds lower=a
+seek-ge a
+first
+----
+a:1
+a:1
+
 # levelIter trims lower/upper bound in the options passed to sstables.
 load a
 ----
@@ -202,6 +210,22 @@ first
 c:3
 c:3
 
+iter
+set-bounds lower=b
+seek-ge a
+first
+----
+a:1
+a:1
+
+iter
+set-bounds lower=c
+seek-ge a
+first
+----
+c:3
+c:3
+
 # levelIter only checks lower bound when loading sstables.
 iter lower=d
 seek-ge a
@@ -224,6 +248,30 @@ last
 dd:5
 dd:5
 
+iter
+set-bounds lower=d
+seek-ge a
+first
+----
+c:3
+c:3
+
+iter
+set-bounds lower=e
+seek-ge a
+first
+----
+.
+.
+
+iter
+set-bounds upper=e
+seek-lt e
+last
+----
+dd:5
+dd:5
+
 # levelIter only checks upper bound when loading sstables.
 iter upper=d
 seek-lt e
@@ -239,6 +287,22 @@ last
 b:2
 b:2
 
+iter
+set-bounds upper=d
+seek-lt e
+last
+----
+d:4
+d:4
+
+iter
+set-bounds upper=c
+seek-lt e
+last
+----
+b:2
+b:2
+
 # levelIter only checks lower bound when loading sstables.
 iter upper=b
 seek-lt e
@@ -261,6 +325,30 @@ next
 d:4
 .
 
+iter
+set-bounds upper=b
+seek-lt e
+last
+----
+b:2
+b:2
+
+iter
+set-bounds upper=a
+seek-lt e
+last
+----
+.
+.
+
+iter
+set-bounds upper=dd
+seek-prefix-ge d
+next
+----
+d:4
+.
+
 iter upper=e
 seek-prefix-ge d
 next
@@ -290,3 +378,29 @@ prev
 ----
 dd:5
 .
+
+iter lower=c
+seek-lt c
+----
+.
+
+iter
+seek-lt c
+set-bounds lower=c
+seek-lt c
+----
+b:2
+.
+
+iter upper=c
+seek-ge c
+----
+.
+
+iter
+seek-ge c
+set-bounds upper=c
+seek-ge c
+----
+c:3
+.