Skip to content

Commit

Permalink
storage/cmdq: prevent iterator from escaping to heap
Browse files Browse the repository at this point in the history
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
  • Loading branch information
nvanbenschoten committed Nov 15, 2018
1 parent 4a02b8c commit a0955d2
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 27 deletions.
94 changes: 67 additions & 27 deletions pkg/storage/cmdq/interval_btree.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
Expand All @@ -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--
}
Expand Down Expand Up @@ -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.
Expand Down
24 changes: 24 additions & 0 deletions pkg/storage/cmdq/interval_btree_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 //
//////////////////////////////////////////
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit a0955d2

Please sign in to comment.