storage/cmdq: prevent iterator from escaping to heap

The iterator returned by `tree.MakeIter` was escaping to the heap
for two reasons:
1. it was capturing a reference to the tree itself.
2. it was pointing a slice into its own array.

This change addresses both of these problems and prevents the iterator
from escaping when used. The fixes were:
1. copy the tree's root pointer reference instead of a reference
   to the tree itself.
2. avoid creating the self-referential slice reference. This mistakenly
   escapes because of golang/go#7921, which
   also caused issues with `bytes.Buffer` (https://golang.org/cl/133715).

This change also adds a new benchmark which demonstrates whether `MakeIter`
escapes or not:

```
name             old time/op    new time/op    delta
BTreeMakeIter-4     131ns ±14%      25ns ± 1%   -81.23%  (p=0.000 n=9+9)

name             old alloc/op   new alloc/op   delta
BTreeMakeIter-4      144B ± 0%        0B       -100.00%  (p=0.000 n=10+10)

name             old allocs/op  new allocs/op  delta
BTreeMakeIter-4      1.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
```

Release note: None

Author: nvanbenschoten

pkg/storage/cmdq/interval_btree.go

@@ -589,10 +589,12 @@ func (t *btree) Set(c *cmd) {
 	}
 }
 
-// MakeIter returns a new iterator object. Note that it is safe for an
-// iterator to be copied by value.
+// MakeIter returns a new iterator object. Note that it is safe for an iterator
+// to be copied by value. However, it is not safe to make modification to the
+// tree after an iterator is created. If modifications are made, create a new
+// iterator.
 func (t *btree) MakeIter() iterator {
-	return iterator{t: t, pos: -1}
+	return iterator{r: t.root, pos: -1}
 }
 
 // Height returns the height of the tree.
@@ -645,45 +647,87 @@ func (n *node) writeString(b *strings.Builder) {
 	}
 }
 
-// iterator is responsible for search and traversal within a btree.
-type iterator struct {
-	t    *btree
-	n    *node
-	pos  int16
+// iterStack represents a stack of (node, pos) tuples, which captures
+// iteration state as an iterator descends a btree.
+type iterStack struct {
+	a    iterStackArr
+	aLen int16 // -1 when using s
 	s    []iterFrame
-	sBuf [3]iterFrame // avoids allocation of s up to height 4
-	o    overlapScan
 }
 
+// Used to avoid allocations for stacks below a certain size.
+type iterStackArr [3]iterFrame
+
 type iterFrame struct {
 	n   *node
 	pos int16
 }
 
+func (is *iterStack) push(f iterFrame) {
+	if is.aLen == -1 {
+		is.s = append(is.s, f)
+	} else if int(is.aLen) == len(is.a) {
+		is.s = make([]iterFrame, int(is.aLen)+1, 2*int(is.aLen))
+		copy(is.s, is.a[:])
+		is.s[int(is.aLen)] = f
+		is.aLen = -1
+	} else {
+		is.a[is.aLen] = f
+		is.aLen++
+	}
+}
+
+func (is *iterStack) pop() iterFrame {
+	if is.aLen == -1 {
+		f := is.s[len(is.s)-1]
+		is.s = is.s[:len(is.s)-1]
+		return f
+	}
+	is.aLen--
+	return is.a[is.aLen]
+}
+
+func (is *iterStack) len() int {
+	if is.aLen == -1 {
+		return len(is.s)
+	}
+	return int(is.aLen)
+}
+
+func (is *iterStack) reset() {
+	if is.aLen == -1 {
+		is.s = is.s[:0]
+	} else {
+		is.aLen = 0
+	}
+}
+
+// iterator is responsible for search and traversal within a btree.
+type iterator struct {
+	r   *node
+	n   *node
+	pos int16
+	s   iterStack
+	o   overlapScan
+}
+
 func (i *iterator) reset() {
-	i.n = i.t.root
-	i.s = i.s[:0]
+	i.n = i.r
 	i.pos = -1
+	i.s.reset()
 	i.o = overlapScan{}
 }
 
-// descend descends into the node's child at position pos. It maintains a
-// stack of (parent, position) pairs so that the tree can be ascended again
-// later.
 func (i *iterator) descend(n *node, pos int16) {
-	if i.s == nil {
-		i.s = i.sBuf[:0]
-	}
-	i.s = append(i.s, iterFrame{n: n, pos: pos})
+	i.s.push(iterFrame{n: n, pos: pos})
 	i.n = n.children[pos]
 	i.pos = 0
 }
 
 // ascend ascends up to the current node's parent and resets the position
 // to the one previously set for this parent node.
 func (i *iterator) ascend() {
-	f := i.s[len(i.s)-1]
-	i.s = i.s[:len(i.s)-1]
+	f := i.s.pop()
 	i.n = f.n
 	i.pos = f.pos
 }
@@ -763,7 +807,7 @@ func (i *iterator) Next() {
 		if i.pos < i.n.count {
 			return
 		}
-		for len(i.s) > 0 && i.pos >= i.n.count {
+		for i.s.len() > 0 && i.pos >= i.n.count {
 			i.ascend()
 		}
 		return
@@ -788,7 +832,7 @@ func (i *iterator) Prev() {
 		if i.pos >= 0 {
 			return
 		}
-		for len(i.s) > 0 && i.pos < 0 {
+		for i.s.len() > 0 && i.pos < 0 {
 			i.ascend()
 			i.pos--
 		}
@@ -921,10 +965,6 @@ func (i *iterator) findNextOverlap() {
 	for {
 		if i.pos > i.n.count {
 			// Iterate up tree.
-			if len(i.s) == 0 {
-				// Should have already hit upper-bound constraint.
-				panic("unreachable")
-			}
 			i.ascend()
 		} else if !i.n.leaf {
 			// Iterate down tree.

pkg/storage/cmdq/interval_btree_test.go

@@ -516,6 +516,22 @@ func TestBTreeCmp(t *testing.T) {
 	}
 }
 
+func TestIterStack(t *testing.T) {
+	f := func(i int) iterFrame { return iterFrame{pos: int16(i)} }
+	var is iterStack
+	for i := 1; i <= 2*len(iterStackArr{}); i++ {
+		var j int
+		for j = 0; j < i; j++ {
+			is.push(f(j))
+		}
+		require.Equal(t, j, is.len())
+		for j--; j >= 0; j-- {
+			require.Equal(t, f(j), is.pop())
+		}
+		is.reset()
+	}
+}
+
 //////////////////////////////////////////
 //              Benchmarks              //
 //////////////////////////////////////////
@@ -594,6 +610,14 @@ func BenchmarkBTreeDeleteInsert(b *testing.B) {
 	})
 }
 
+func BenchmarkBTreeMakeIter(b *testing.B) {
+	var tr btree
+	for i := 0; i < b.N; i++ {
+		it := tr.MakeIter()
+		it.First()
+	}
+}
+
 func BenchmarkBTreeIterSeekGE(b *testing.B) {
 	forBenchmarkSizes(b, func(b *testing.B, count int) {
 		var spans []roachpb.Span