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