github.com/cockroachdb/pebble@v1.1.2/internal/manifest/btree_test.go (about)

     1  // Copyright 2020 The LevelDB-Go and Pebble Authors. All rights reserved. Use
     2  // of this source code is governed by a BSD-style license that can be found in
     3  // the LICENSE file.
     4  
     5  package manifest
     6  
     7  import (
     8  	"fmt"
     9  	"math/rand"
    10  	"reflect"
    11  	"sort"
    12  	"sync"
    13  	"testing"
    14  	"time"
    15  
    16  	"github.com/cockroachdb/errors"
    17  	"github.com/cockroachdb/pebble/internal/base"
    18  	"github.com/cockroachdb/pebble/internal/invariants"
    19  	"github.com/stretchr/testify/require"
    20  )
    21  
    22  func newItem(k InternalKey) *FileMetadata {
    23  	m := (&FileMetadata{}).ExtendPointKeyBounds(
    24  		base.DefaultComparer.Compare, k, k,
    25  	)
    26  	m.InitPhysicalBacking()
    27  	return m
    28  }
    29  
    30  func cmp(a, b *FileMetadata) int {
    31  	return cmpKey(a.Smallest, b.Smallest)
    32  }
    33  
    34  func cmpKey(a, b InternalKey) int {
    35  	return base.InternalCompare(base.DefaultComparer.Compare, a, b)
    36  }
    37  
    38  //////////////////////////////////////////
    39  //        Invariant verification        //
    40  //////////////////////////////////////////
    41  
    42  // Verify asserts that the tree's structural invariants all hold.
    43  func (t *btree) Verify(tt *testing.T) {
    44  	if t.Count() == 0 {
    45  		require.Nil(tt, t.root)
    46  		return
    47  	}
    48  	t.verifyLeafSameDepth(tt)
    49  	t.verifyCountAllowed(tt)
    50  	t.isSorted(tt)
    51  	t.root.verifyInvariants()
    52  }
    53  
    54  func (t *btree) verifyLeafSameDepth(tt *testing.T) {
    55  	h := t.height()
    56  	t.root.verifyDepthEqualToHeight(tt, 1, h)
    57  }
    58  
    59  func (n *node) verifyDepthEqualToHeight(t *testing.T, depth, height int) {
    60  	if n.leaf {
    61  		require.Equal(t, height, depth, "all leaves should have the same depth as the tree height")
    62  	}
    63  	n.recurse(func(child *node, _ int16) {
    64  		child.verifyDepthEqualToHeight(t, depth+1, height)
    65  	})
    66  }
    67  
    68  func (t *btree) verifyCountAllowed(tt *testing.T) {
    69  	t.root.verifyCountAllowed(tt, true)
    70  }
    71  
    72  // height returns the height of the tree.
    73  func (t *btree) height() int {
    74  	if t.root == nil {
    75  		return 0
    76  	}
    77  	h := 1
    78  	n := t.root
    79  	for !n.leaf {
    80  		n = n.children[0]
    81  		h++
    82  	}
    83  	return h
    84  }
    85  
    86  func (n *node) verifyCountAllowed(t *testing.T, root bool) {
    87  	if !root {
    88  		require.GreaterOrEqual(t, n.count, int16(minItems), "item count %d must be in range [%d,%d]", n.count, minItems, maxItems)
    89  		require.LessOrEqual(t, n.count, int16(maxItems), "item count %d must be in range [%d,%d]", n.count, minItems, maxItems)
    90  	}
    91  	for i, item := range n.items {
    92  		if i < int(n.count) {
    93  			require.NotNil(t, item, "item below count")
    94  		} else {
    95  			require.Nil(t, item, "item above count")
    96  		}
    97  	}
    98  	if !n.leaf {
    99  		for i, child := range n.children {
   100  			if i <= int(n.count) {
   101  				require.NotNil(t, child, "node below count")
   102  			} else {
   103  				require.Nil(t, child, "node above count")
   104  			}
   105  		}
   106  	}
   107  	n.recurse(func(child *node, _ int16) {
   108  		child.verifyCountAllowed(t, false)
   109  	})
   110  }
   111  
   112  func (t *btree) isSorted(tt *testing.T) {
   113  	t.root.isSorted(tt, t.cmp)
   114  }
   115  
   116  func (n *node) isSorted(t *testing.T, cmp func(*FileMetadata, *FileMetadata) int) {
   117  	for i := int16(1); i < n.count; i++ {
   118  		require.LessOrEqual(t, cmp(n.items[i-1], n.items[i]), 0)
   119  	}
   120  	if !n.leaf {
   121  		for i := int16(0); i < n.count; i++ {
   122  			prev := n.children[i]
   123  			next := n.children[i+1]
   124  
   125  			require.LessOrEqual(t, cmp(prev.items[prev.count-1], n.items[i]), 0)
   126  			require.LessOrEqual(t, cmp(n.items[i], next.items[0]), 0)
   127  		}
   128  	}
   129  	n.recurse(func(child *node, _ int16) {
   130  		child.isSorted(t, cmp)
   131  	})
   132  }
   133  
   134  func (n *node) recurse(f func(child *node, pos int16)) {
   135  	if !n.leaf {
   136  		for i := int16(0); i <= n.count; i++ {
   137  			f(n.children[i], i)
   138  		}
   139  	}
   140  }
   141  
   142  //////////////////////////////////////////
   143  //              Unit Tests              //
   144  //////////////////////////////////////////
   145  
   146  func key(i int) InternalKey {
   147  	if i < 0 || i > 99999 {
   148  		panic("key out of bounds")
   149  	}
   150  	return base.MakeInternalKey([]byte(fmt.Sprintf("%05d", i)), 0, base.InternalKeyKindSet)
   151  }
   152  
   153  func keyWithMemo(i int, memo map[int]InternalKey) InternalKey {
   154  	if s, ok := memo[i]; ok {
   155  		return s
   156  	}
   157  	s := key(i)
   158  	memo[i] = s
   159  	return s
   160  }
   161  
   162  func checkIterRelative(t *testing.T, it *iterator, start, end int, keyMemo map[int]InternalKey) {
   163  	t.Helper()
   164  	i := start
   165  	for ; it.valid(); it.next() {
   166  		item := it.cur()
   167  		expected := keyWithMemo(i, keyMemo)
   168  		if cmpKey(expected, item.Smallest) != 0 {
   169  			t.Fatalf("expected %s, but found %s", expected, item.Smallest)
   170  		}
   171  		i++
   172  	}
   173  	if i != end {
   174  		t.Fatalf("expected %d, but at %d", end, i)
   175  	}
   176  }
   177  
   178  func checkIter(t *testing.T, it iterator, start, end int, keyMemo map[int]InternalKey) {
   179  	t.Helper()
   180  	i := start
   181  	for it.first(); it.valid(); it.next() {
   182  		item := it.cur()
   183  		expected := keyWithMemo(i, keyMemo)
   184  		if cmpKey(expected, item.Smallest) != 0 {
   185  			t.Fatalf("expected %s, but found %s", expected, item.Smallest)
   186  		}
   187  		require.Equal(t, i-start, it.countLeft())
   188  		i++
   189  	}
   190  	if i != end {
   191  		t.Fatalf("expected %d, but at %d", end, i)
   192  	}
   193  
   194  	for it.last(); it.valid(); it.prev() {
   195  		i--
   196  		item := it.cur()
   197  		expected := keyWithMemo(i, keyMemo)
   198  		if cmpKey(expected, item.Smallest) != 0 {
   199  			t.Fatalf("expected %s, but found %s", expected, item.Smallest)
   200  		}
   201  		require.Equal(t, i-start, it.countLeft())
   202  	}
   203  	if i != start {
   204  		t.Fatalf("expected %d, but at %d: %+v", start, i, it)
   205  	}
   206  }
   207  
   208  // TestBTree tests basic btree operations.
   209  func TestBTree(t *testing.T) {
   210  	var tr btree
   211  	tr.cmp = cmp
   212  	keyMemo := make(map[int]InternalKey)
   213  
   214  	// With degree == 16 (max-items/node == 31) we need 513 items in order for
   215  	// there to be 3 levels in the tree. The count here is comfortably above
   216  	// that.
   217  	const count = 768
   218  	items := rang(0, count-1)
   219  
   220  	// Add keys in sorted order.
   221  	for i := 0; i < count; i++ {
   222  		require.NoError(t, tr.Insert(items[i]))
   223  		tr.Verify(t)
   224  		if e := i + 1; e != tr.Count() {
   225  			t.Fatalf("expected length %d, but found %d", e, tr.Count())
   226  		}
   227  		checkIter(t, tr.Iter(), 0, i+1, keyMemo)
   228  	}
   229  
   230  	// delete keys in sorted order.
   231  	for i := 0; i < count; i++ {
   232  		obsolete := tr.Delete(items[i])
   233  		if !obsolete {
   234  			t.Fatalf("expected item %d to be obsolete", i)
   235  		}
   236  		tr.Verify(t)
   237  		if e := count - (i + 1); e != tr.Count() {
   238  			t.Fatalf("expected length %d, but found %d", e, tr.Count())
   239  		}
   240  		checkIter(t, tr.Iter(), i+1, count, keyMemo)
   241  	}
   242  
   243  	// Add keys in reverse sorted order.
   244  	for i := 1; i <= count; i++ {
   245  		require.NoError(t, tr.Insert(items[count-i]))
   246  		tr.Verify(t)
   247  		if i != tr.Count() {
   248  			t.Fatalf("expected length %d, but found %d", i, tr.Count())
   249  		}
   250  		checkIter(t, tr.Iter(), count-i, count, keyMemo)
   251  	}
   252  
   253  	// delete keys in reverse sorted order.
   254  	for i := 1; i <= count; i++ {
   255  		obsolete := tr.Delete(items[count-i])
   256  		if !obsolete {
   257  			t.Fatalf("expected item %d to be obsolete", i)
   258  		}
   259  		tr.Verify(t)
   260  		if e := count - i; e != tr.Count() {
   261  			t.Fatalf("expected length %d, but found %d", e, tr.Count())
   262  		}
   263  		checkIter(t, tr.Iter(), 0, count-i, keyMemo)
   264  	}
   265  }
   266  
   267  func TestIterClone(t *testing.T) {
   268  	const count = 65536
   269  
   270  	var tr btree
   271  	tr.cmp = cmp
   272  	keyMemo := make(map[int]InternalKey)
   273  
   274  	for i := 0; i < count; i++ {
   275  		require.NoError(t, tr.Insert(newItem(key(i))))
   276  	}
   277  
   278  	it := tr.Iter()
   279  	i := 0
   280  	for it.first(); it.valid(); it.next() {
   281  		if i%500 == 0 {
   282  			c := it.clone()
   283  
   284  			require.Equal(t, 0, cmpIter(it, c))
   285  			checkIterRelative(t, &c, i, count, keyMemo)
   286  			if i < count {
   287  				require.Equal(t, -1, cmpIter(it, c))
   288  				require.Equal(t, +1, cmpIter(c, it))
   289  			}
   290  		}
   291  		i++
   292  	}
   293  }
   294  
   295  func TestIterCmpEdgeCases(t *testing.T) {
   296  	var tr btree
   297  	tr.cmp = cmp
   298  	t.Run("empty", func(t *testing.T) {
   299  		a := tr.Iter()
   300  		b := tr.Iter()
   301  		require.Equal(t, 0, cmpIter(a, b))
   302  	})
   303  	require.NoError(t, tr.Insert(newItem(key(5))))
   304  	t.Run("exhausted_next", func(t *testing.T) {
   305  		a := tr.Iter()
   306  		b := tr.Iter()
   307  		a.first()
   308  		b.first()
   309  		require.Equal(t, 0, cmpIter(a, b))
   310  		b.next()
   311  		require.False(t, b.valid())
   312  		require.Equal(t, -1, cmpIter(a, b))
   313  	})
   314  	t.Run("exhausted_prev", func(t *testing.T) {
   315  		a := tr.Iter()
   316  		b := tr.Iter()
   317  		a.first()
   318  		b.first()
   319  		b.prev()
   320  		require.False(t, b.valid())
   321  		require.Equal(t, 1, cmpIter(a, b))
   322  		b.next()
   323  		require.Equal(t, 0, cmpIter(a, b))
   324  	})
   325  }
   326  
   327  func TestIterCmpRand(t *testing.T) {
   328  	const itemCount = 65536
   329  	const iterCount = 1000
   330  
   331  	var tr btree
   332  	tr.cmp = cmp
   333  	for i := 0; i < itemCount; i++ {
   334  		require.NoError(t, tr.Insert(newItem(key(i))))
   335  	}
   336  
   337  	seed := time.Now().UnixNano()
   338  	rng := rand.New(rand.NewSource(seed))
   339  	iters1 := make([]*LevelIterator, iterCount)
   340  	iters2 := make([]*LevelIterator, iterCount)
   341  	for i := 0; i < iterCount; i++ {
   342  		k := rng.Intn(itemCount)
   343  		iter := LevelIterator{iter: tr.Iter()}
   344  		iter.SeekGE(base.DefaultComparer.Compare, key(k).UserKey)
   345  		iters1[i] = &iter
   346  		iters2[i] = &iter
   347  	}
   348  
   349  	// All the iterators should be positioned, so sorting them by items and by
   350  	// iterator comparisons should equal identical orderings.
   351  	sort.SliceStable(iters1, func(i, j int) bool { return cmpIter(iters1[i].iter, iters1[j].iter) < 0 })
   352  	sort.SliceStable(iters2, func(i, j int) bool { return cmp(iters2[i].iter.cur(), iters2[j].iter.cur()) < 0 })
   353  	for i := 0; i < iterCount; i++ {
   354  		if iters1[i] != iters2[i] {
   355  			t.Fatalf("seed %d: iters out of order at index %d:\n%s\n\n%s",
   356  				seed, i, iters1[i], iters2[i])
   357  		}
   358  	}
   359  }
   360  
   361  // TestBTreeSeek tests basic btree iterator operations on an iterator wrapped
   362  // by a LevelIterator.
   363  func TestBTreeSeek(t *testing.T) {
   364  	const count = 513
   365  
   366  	var tr btree
   367  	tr.cmp = cmp
   368  	for i := 0; i < count; i++ {
   369  		require.NoError(t, tr.Insert(newItem(key(i*2))))
   370  	}
   371  
   372  	it := LevelIterator{iter: tr.Iter()}
   373  	for i := 0; i < 2*count-1; i++ {
   374  		item := it.SeekGE(base.DefaultComparer.Compare, key(i).UserKey)
   375  		if item == nil {
   376  			t.Fatalf("%d: expected valid iterator", i)
   377  		}
   378  		expected := key(2 * ((i + 1) / 2))
   379  		if cmpKey(expected, item.Smallest) != 0 {
   380  			t.Fatalf("%d: expected %s, but found %s", i, expected, item.Smallest)
   381  		}
   382  	}
   383  	it.SeekGE(base.DefaultComparer.Compare, key(2*count-1).UserKey)
   384  	if it.iter.valid() {
   385  		t.Fatalf("expected invalid iterator")
   386  	}
   387  
   388  	for i := 1; i < 2*count; i++ {
   389  		item := it.SeekLT(base.DefaultComparer.Compare, key(i).UserKey)
   390  		if item == nil {
   391  			t.Fatalf("%d: expected valid iterator", i)
   392  		}
   393  		expected := key(2 * ((i - 1) / 2))
   394  		if cmpKey(expected, item.Smallest) != 0 {
   395  			t.Fatalf("%d: expected %s, but found %s", i, expected, item.Smallest)
   396  		}
   397  	}
   398  	it.SeekLT(base.DefaultComparer.Compare, key(0).UserKey)
   399  	if it.iter.valid() {
   400  		t.Fatalf("expected invalid iterator")
   401  	}
   402  }
   403  
   404  func TestBTreeInsertDuplicateError(t *testing.T) {
   405  	var tr btree
   406  	tr.cmp = cmp
   407  	require.NoError(t, tr.Insert(newItem(key(1))))
   408  	require.NoError(t, tr.Insert(newItem(key(2))))
   409  	require.NoError(t, tr.Insert(newItem(key(3))))
   410  	wantErr := errors.Errorf("files %s and %s collided on sort keys",
   411  		errors.Safe(base.FileNum(000000)), errors.Safe(base.FileNum(000000)))
   412  	require.Error(t, wantErr, tr.Insert(newItem(key(2))))
   413  }
   414  
   415  // TestBTreeCloneConcurrentOperations tests that cloning a btree returns a new
   416  // btree instance which is an exact logical copy of the original but that can be
   417  // modified independently going forward.
   418  func TestBTreeCloneConcurrentOperations(t *testing.T) {
   419  	const cloneTestSize = 1000
   420  	p := perm(cloneTestSize)
   421  
   422  	var trees []*btree
   423  	treeC, treeDone := make(chan *btree), make(chan struct{})
   424  	go func() {
   425  		for b := range treeC {
   426  			trees = append(trees, b)
   427  		}
   428  		close(treeDone)
   429  	}()
   430  
   431  	var wg sync.WaitGroup
   432  	var populate func(tr *btree, start int)
   433  	populate = func(tr *btree, start int) {
   434  		t.Logf("Starting new clone at %v", start)
   435  		treeC <- tr
   436  		for i := start; i < cloneTestSize; i++ {
   437  			require.NoError(t, tr.Insert(p[i]))
   438  			if i%(cloneTestSize/5) == 0 {
   439  				wg.Add(1)
   440  				c := tr.Clone()
   441  				go populate(&c, i+1)
   442  			}
   443  		}
   444  		wg.Done()
   445  	}
   446  
   447  	wg.Add(1)
   448  	var tr btree
   449  	tr.cmp = cmp
   450  	go populate(&tr, 0)
   451  	wg.Wait()
   452  	close(treeC)
   453  	<-treeDone
   454  
   455  	t.Logf("Starting equality checks on %d trees", len(trees))
   456  	want := rang(0, cloneTestSize-1)
   457  	for i, tree := range trees {
   458  		if got := all(tree); !reflect.DeepEqual(strReprs(got), strReprs(want)) {
   459  			t.Errorf("tree %v mismatch", i)
   460  		}
   461  	}
   462  
   463  	t.Log("Removing half of items from first half")
   464  	toRemove := want[cloneTestSize/2:]
   465  	for i := 0; i < len(trees)/2; i++ {
   466  		tree := trees[i]
   467  		wg.Add(1)
   468  		go func() {
   469  			for _, item := range toRemove {
   470  				tree.Delete(item)
   471  			}
   472  			wg.Done()
   473  		}()
   474  	}
   475  	wg.Wait()
   476  
   477  	t.Log("Checking all values again")
   478  	for i, tree := range trees {
   479  		var wantpart []*FileMetadata
   480  		if i < len(trees)/2 {
   481  			wantpart = want[:cloneTestSize/2]
   482  		} else {
   483  			wantpart = want
   484  		}
   485  		if got := all(tree); !reflect.DeepEqual(strReprs(got), strReprs(wantpart)) {
   486  			t.Errorf("tree %v mismatch, want %#v got %#v", i, strReprs(wantpart), strReprs(got))
   487  		}
   488  	}
   489  
   490  	var obsolete []*FileBacking
   491  	for i := range trees {
   492  		obsolete = append(obsolete, trees[i].Release()...)
   493  	}
   494  	if len(obsolete) != len(p) {
   495  		t.Errorf("got %d obsolete trees, expected %d", len(obsolete), len(p))
   496  	}
   497  }
   498  
   499  // TestIterStack tests the interface of the iterStack type.
   500  func TestIterStack(t *testing.T) {
   501  	f := func(i int) iterFrame { return iterFrame{pos: int16(i)} }
   502  	var is iterStack
   503  	for i := 1; i <= 2*len(iterStackArr{}); i++ {
   504  		var j int
   505  		for j = 0; j < i; j++ {
   506  			is.push(f(j))
   507  		}
   508  		require.Equal(t, j, is.len())
   509  		for j--; j >= 0; j-- {
   510  			require.Equal(t, f(j), is.pop())
   511  		}
   512  		is.reset()
   513  	}
   514  }
   515  
   516  func TestIterEndSentinel(t *testing.T) {
   517  	var tr btree
   518  	tr.cmp = cmp
   519  	require.NoError(t, tr.Insert(newItem(key(1))))
   520  	require.NoError(t, tr.Insert(newItem(key(2))))
   521  	require.NoError(t, tr.Insert(newItem(key(3))))
   522  	iter := LevelIterator{iter: tr.Iter()}
   523  	iter.SeekGE(base.DefaultComparer.Compare, key(3).UserKey)
   524  	require.True(t, iter.iter.valid())
   525  	iter.Next()
   526  	require.False(t, iter.iter.valid())
   527  
   528  	// If we seek into the end sentinel, prev should return us to a valid
   529  	// position.
   530  	iter.SeekGE(base.DefaultComparer.Compare, key(4).UserKey)
   531  	require.False(t, iter.iter.valid())
   532  	iter.Prev()
   533  	require.True(t, iter.iter.valid())
   534  }
   535  
   536  type orderStatistic struct{}
   537  
   538  func (o orderStatistic) Zero(dst interface{}) interface{} {
   539  	if dst == nil {
   540  		return new(int)
   541  	}
   542  	v := dst.(*int)
   543  	*v = 0
   544  	return v
   545  }
   546  
   547  func (o orderStatistic) Accumulate(meta *FileMetadata, dst interface{}) (interface{}, bool) {
   548  	v := dst.(*int)
   549  	*v++
   550  	return v, true
   551  }
   552  
   553  func (o orderStatistic) Merge(src interface{}, dst interface{}) interface{} {
   554  	srcv := src.(*int)
   555  	dstv := dst.(*int)
   556  	*dstv = *dstv + *srcv
   557  	return dstv
   558  }
   559  
   560  func TestAnnotationOrderStatistic(t *testing.T) {
   561  	const count = 1000
   562  	ann := orderStatistic{}
   563  
   564  	var tr btree
   565  	tr.cmp = cmp
   566  	for i := 1; i <= count; i++ {
   567  		require.NoError(t, tr.Insert(newItem(key(i))))
   568  
   569  		v, ok := tr.root.Annotation(ann)
   570  		require.True(t, ok)
   571  		vtyped := v.(*int)
   572  		require.Equal(t, i, *vtyped)
   573  	}
   574  
   575  	v, ok := tr.root.Annotation(ann)
   576  	require.True(t, ok)
   577  	vtyped := v.(*int)
   578  	require.Equal(t, count, *vtyped)
   579  
   580  	v, ok = tr.root.Annotation(ann)
   581  	vtyped = v.(*int)
   582  	require.True(t, ok)
   583  	require.Equal(t, count, *vtyped)
   584  }
   585  
   586  // TestRandomizedBTree tests a random set of Insert, Delete and iteration
   587  // operations, checking for equivalence with a map of filenums.
   588  func TestRandomizedBTree(t *testing.T) {
   589  	const maxFileNum = 50_000
   590  
   591  	seed := time.Now().UnixNano()
   592  	t.Log("seed", seed)
   593  	rng := rand.New(rand.NewSource(seed))
   594  
   595  	var numOps int
   596  	if invariants.RaceEnabled {
   597  		// Reduce the number of ops in race mode so the test doesn't take very long.
   598  		numOps = 1_000 + rng.Intn(4_000)
   599  	} else {
   600  		numOps = 10_000 + rng.Intn(40_000)
   601  	}
   602  
   603  	var metadataAlloc [maxFileNum]FileMetadata
   604  	for i := 0; i < len(metadataAlloc); i++ {
   605  		metadataAlloc[i].FileNum = base.FileNum(i)
   606  		metadataAlloc[i].InitPhysicalBacking()
   607  	}
   608  
   609  	// Use a btree comparator that sorts by file number to make it easier to
   610  	// prevent duplicates or overlaps.
   611  	tree := btree{
   612  		cmp: func(a *FileMetadata, b *FileMetadata) int {
   613  			switch {
   614  			case a.FileNum < b.FileNum:
   615  				return -1
   616  			case a.FileNum > b.FileNum:
   617  				return +1
   618  			default:
   619  				return 0
   620  			}
   621  		},
   622  	}
   623  
   624  	type opDecl struct {
   625  		fn     func()
   626  		weight int
   627  	}
   628  	ref := map[base.FileNum]bool{}
   629  	ops := []opDecl{
   630  		{
   631  			// Insert
   632  			fn: func() {
   633  				f := &metadataAlloc[rng.Intn(maxFileNum)]
   634  				err := tree.Insert(f)
   635  				if ref[f.FileNum] {
   636  					require.Error(t, err, "btree.Insert should error if file already exists")
   637  				} else {
   638  					ref[f.FileNum] = true
   639  					require.NoError(t, err)
   640  				}
   641  			},
   642  			weight: 20,
   643  		},
   644  		{
   645  			// Delete
   646  			fn: func() {
   647  				f := &metadataAlloc[rng.Intn(maxFileNum)]
   648  				tree.Delete(f)
   649  				delete(ref, f.FileNum)
   650  			},
   651  			weight: 10,
   652  		},
   653  		{
   654  			// Iterate
   655  			fn: func() {
   656  				iter := tree.Iter()
   657  				count := 0
   658  				var prev base.FileNum
   659  				for iter.first(); iter.valid(); iter.next() {
   660  					fn := iter.cur().FileNum
   661  					require.True(t, ref[fn])
   662  					if count > 0 {
   663  						require.Less(t, prev, fn)
   664  					}
   665  					count++
   666  				}
   667  				require.Equal(t, count, len(ref))
   668  			},
   669  			weight: 1,
   670  		},
   671  	}
   672  	weightSum := 0
   673  	for i := range ops {
   674  		weightSum += ops[i].weight
   675  	}
   676  
   677  	for i := 0; i < numOps; i++ {
   678  		w := rng.Intn(weightSum)
   679  		for j := range ops {
   680  			w -= ops[j].weight
   681  			if w < 0 {
   682  				ops[j].fn()
   683  				break
   684  			}
   685  		}
   686  	}
   687  }
   688  
   689  //////////////////////////////////////////
   690  //              Benchmarks              //
   691  //////////////////////////////////////////
   692  
   693  // perm returns a random permutation of items with keys in the range [0, n).
   694  func perm(n int) (out []*FileMetadata) {
   695  	for _, i := range rand.Perm(n) {
   696  		out = append(out, newItem(key(i)))
   697  	}
   698  	return out
   699  }
   700  
   701  // rang returns an ordered list of items with keys in the range [m, n].
   702  func rang(m, n int) (out []*FileMetadata) {
   703  	for i := m; i <= n; i++ {
   704  		out = append(out, newItem(key(i)))
   705  	}
   706  	return out
   707  }
   708  
   709  func strReprs(items []*FileMetadata) []string {
   710  	s := make([]string, len(items))
   711  	for i := range items {
   712  		s[i] = items[i].String()
   713  	}
   714  	return s
   715  }
   716  
   717  // all extracts all items from a tree in order as a slice.
   718  func all(tr *btree) (out []*FileMetadata) {
   719  	it := tr.Iter()
   720  	it.first()
   721  	for it.valid() {
   722  		out = append(out, it.cur())
   723  		it.next()
   724  	}
   725  	return out
   726  }
   727  
   728  func forBenchmarkSizes(b *testing.B, f func(b *testing.B, count int)) {
   729  	for _, count := range []int{16, 128, 1024, 8192, 65536} {
   730  		b.Run(fmt.Sprintf("count=%d", count), func(b *testing.B) {
   731  			f(b, count)
   732  		})
   733  	}
   734  }
   735  
   736  // BenchmarkBTreeInsert measures btree insertion performance.
   737  func BenchmarkBTreeInsert(b *testing.B) {
   738  	forBenchmarkSizes(b, func(b *testing.B, count int) {
   739  		insertP := perm(count)
   740  		b.ResetTimer()
   741  		for i := 0; i < b.N; {
   742  			var tr btree
   743  			tr.cmp = cmp
   744  			for _, item := range insertP {
   745  				if err := tr.Insert(item); err != nil {
   746  					b.Fatal(err)
   747  				}
   748  				i++
   749  				if i >= b.N {
   750  					return
   751  				}
   752  			}
   753  		}
   754  	})
   755  }
   756  
   757  // BenchmarkBTreeDelete measures btree deletion performance.
   758  func BenchmarkBTreeDelete(b *testing.B) {
   759  	forBenchmarkSizes(b, func(b *testing.B, count int) {
   760  		insertP, removeP := perm(count), perm(count)
   761  		b.ResetTimer()
   762  		for i := 0; i < b.N; {
   763  			b.StopTimer()
   764  			var tr btree
   765  			tr.cmp = cmp
   766  			for _, item := range insertP {
   767  				if err := tr.Insert(item); err != nil {
   768  					b.Fatal(err)
   769  				}
   770  			}
   771  			b.StartTimer()
   772  			for _, item := range removeP {
   773  				tr.Delete(item)
   774  				i++
   775  				if i >= b.N {
   776  					return
   777  				}
   778  			}
   779  			if tr.Count() > 0 {
   780  				b.Fatalf("tree not empty: %s", &tr)
   781  			}
   782  		}
   783  	})
   784  }
   785  
   786  // BenchmarkBTreeDeleteInsert measures btree deletion and insertion performance.
   787  func BenchmarkBTreeDeleteInsert(b *testing.B) {
   788  	forBenchmarkSizes(b, func(b *testing.B, count int) {
   789  		insertP := perm(count)
   790  		var tr btree
   791  		tr.cmp = cmp
   792  		for _, item := range insertP {
   793  			if err := tr.Insert(item); err != nil {
   794  				b.Fatal(err)
   795  			}
   796  		}
   797  		b.ResetTimer()
   798  		for i := 0; i < b.N; i++ {
   799  			item := insertP[i%count]
   800  			tr.Delete(item)
   801  			if err := tr.Insert(item); err != nil {
   802  				b.Fatal(err)
   803  			}
   804  		}
   805  	})
   806  }
   807  
   808  // BenchmarkBTreeDeleteInsertCloneOnce measures btree deletion and insertion
   809  // performance after the tree has been copy-on-write cloned once.
   810  func BenchmarkBTreeDeleteInsertCloneOnce(b *testing.B) {
   811  	forBenchmarkSizes(b, func(b *testing.B, count int) {
   812  		insertP := perm(count)
   813  		var tr btree
   814  		tr.cmp = cmp
   815  		for _, item := range insertP {
   816  			if err := tr.Insert(item); err != nil {
   817  				b.Fatal(err)
   818  			}
   819  		}
   820  		tr = tr.Clone()
   821  		b.ResetTimer()
   822  		for i := 0; i < b.N; i++ {
   823  			item := insertP[i%count]
   824  			tr.Delete(item)
   825  			if err := tr.Insert(item); err != nil {
   826  				b.Fatal(err)
   827  			}
   828  		}
   829  	})
   830  }
   831  
   832  // BenchmarkBTreeDeleteInsertCloneEachTime measures btree deletion and insertion
   833  // performance while the tree is repeatedly copy-on-write cloned.
   834  func BenchmarkBTreeDeleteInsertCloneEachTime(b *testing.B) {
   835  	for _, release := range []bool{false, true} {
   836  		b.Run(fmt.Sprintf("release=%t", release), func(b *testing.B) {
   837  			forBenchmarkSizes(b, func(b *testing.B, count int) {
   838  				insertP := perm(count)
   839  				var tr, trRelease btree
   840  				tr.cmp = cmp
   841  				trRelease.cmp = cmp
   842  				for _, item := range insertP {
   843  					if err := tr.Insert(item); err != nil {
   844  						b.Fatal(err)
   845  					}
   846  				}
   847  				b.ResetTimer()
   848  				for i := 0; i < b.N; i++ {
   849  					item := insertP[i%count]
   850  					if release {
   851  						trRelease.Release()
   852  						trRelease = tr
   853  					}
   854  					tr = tr.Clone()
   855  					tr.Delete(item)
   856  					if err := tr.Insert(item); err != nil {
   857  						b.Fatal(err)
   858  					}
   859  				}
   860  			})
   861  		})
   862  	}
   863  }
   864  
   865  // BenchmarkBTreeIter measures the cost of creating a btree iterator.
   866  func BenchmarkBTreeIter(b *testing.B) {
   867  	var tr btree
   868  	tr.cmp = cmp
   869  	for i := 0; i < b.N; i++ {
   870  		it := tr.Iter()
   871  		it.first()
   872  	}
   873  }
   874  
   875  // BenchmarkBTreeIterSeekGE measures the cost of seeking a btree iterator
   876  // forward.
   877  func BenchmarkBTreeIterSeekGE(b *testing.B) {
   878  	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
   879  	forBenchmarkSizes(b, func(b *testing.B, count int) {
   880  		var keys []InternalKey
   881  		var tr btree
   882  		tr.cmp = cmp
   883  
   884  		for i := 0; i < count; i++ {
   885  			s := key(i)
   886  			keys = append(keys, s)
   887  			if err := tr.Insert(newItem(s)); err != nil {
   888  				b.Fatal(err)
   889  			}
   890  		}
   891  
   892  		b.ResetTimer()
   893  		for i := 0; i < b.N; i++ {
   894  			k := keys[rng.Intn(len(keys))]
   895  			it := LevelIterator{iter: tr.Iter()}
   896  			f := it.SeekGE(base.DefaultComparer.Compare, k.UserKey)
   897  			if testing.Verbose() {
   898  				if f == nil {
   899  					b.Fatal("expected to find key")
   900  				}
   901  				if cmpKey(k, f.Smallest) != 0 {
   902  					b.Fatalf("expected %s, but found %s", k, f.Smallest)
   903  				}
   904  			}
   905  		}
   906  	})
   907  }
   908  
   909  // BenchmarkBTreeIterSeekLT measures the cost of seeking a btree iterator
   910  // backward.
   911  func BenchmarkBTreeIterSeekLT(b *testing.B) {
   912  	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
   913  	forBenchmarkSizes(b, func(b *testing.B, count int) {
   914  		var keys []InternalKey
   915  		var tr btree
   916  		tr.cmp = cmp
   917  
   918  		for i := 0; i < count; i++ {
   919  			k := key(i)
   920  			keys = append(keys, k)
   921  			if err := tr.Insert(newItem(k)); err != nil {
   922  				b.Fatal(err)
   923  			}
   924  		}
   925  
   926  		b.ResetTimer()
   927  		for i := 0; i < b.N; i++ {
   928  			j := rng.Intn(len(keys))
   929  			k := keys[j]
   930  			it := LevelIterator{iter: tr.Iter()}
   931  			f := it.SeekLT(base.DefaultComparer.Compare, k.UserKey)
   932  			if testing.Verbose() {
   933  				if j == 0 {
   934  					if f != nil {
   935  						b.Fatal("unexpected key")
   936  					}
   937  				} else {
   938  					if f == nil {
   939  						b.Fatal("expected to find key")
   940  					}
   941  					k := keys[j-1]
   942  					if cmpKey(k, f.Smallest) != 0 {
   943  						b.Fatalf("expected %s, but found %s", k, f.Smallest)
   944  					}
   945  				}
   946  			}
   947  		}
   948  	})
   949  }
   950  
   951  // BenchmarkBTreeIterNext measures the cost of seeking a btree iterator to the
   952  // next item in the tree.
   953  func BenchmarkBTreeIterNext(b *testing.B) {
   954  	var tr btree
   955  	tr.cmp = cmp
   956  
   957  	const count = 8 << 10
   958  	for i := 0; i < count; i++ {
   959  		item := newItem(key(i))
   960  		if err := tr.Insert(item); err != nil {
   961  			b.Fatal(err)
   962  		}
   963  	}
   964  
   965  	it := tr.Iter()
   966  	b.ResetTimer()
   967  	for i := 0; i < b.N; i++ {
   968  		if !it.valid() {
   969  			it.first()
   970  		}
   971  		it.next()
   972  	}
   973  }
   974  
   975  // BenchmarkBTreeIterPrev measures the cost of seeking a btree iterator to the
   976  // previous item in the tree.
   977  func BenchmarkBTreeIterPrev(b *testing.B) {
   978  	var tr btree
   979  	tr.cmp = cmp
   980  
   981  	const count = 8 << 10
   982  	for i := 0; i < count; i++ {
   983  		item := newItem(key(i))
   984  		if err := tr.Insert(item); err != nil {
   985  			b.Fatal(err)
   986  		}
   987  	}
   988  
   989  	it := tr.Iter()
   990  	b.ResetTimer()
   991  	for i := 0; i < b.N; i++ {
   992  		if !it.valid() {
   993  			it.first()
   994  		}
   995  		it.prev()
   996  	}
   997  }