github.com/cockroachdb/pebble@v0.0.0-20231214172447-ab4952c5f87b/level_iter_test.go (about)

     1  // Copyright 2018 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 pebble
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"fmt"
    11  	"strings"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/cockroachdb/datadriven"
    16  	"github.com/cockroachdb/pebble/bloom"
    17  	"github.com/cockroachdb/pebble/internal/base"
    18  	"github.com/cockroachdb/pebble/internal/itertest"
    19  	"github.com/cockroachdb/pebble/internal/keyspan"
    20  	"github.com/cockroachdb/pebble/internal/manifest"
    21  	"github.com/cockroachdb/pebble/internal/rangedel"
    22  	"github.com/cockroachdb/pebble/internal/testkeys"
    23  	"github.com/cockroachdb/pebble/objstorage/objstorageprovider"
    24  	"github.com/cockroachdb/pebble/sstable"
    25  	"github.com/cockroachdb/pebble/vfs"
    26  	"github.com/stretchr/testify/require"
    27  	"golang.org/x/exp/rand"
    28  )
    29  
    30  const (
    31  	level = 1
    32  )
    33  
    34  func TestLevelIter(t *testing.T) {
    35  	var iters []*fakeIter
    36  	var files manifest.LevelSlice
    37  
    38  	newIters := func(
    39  		_ context.Context, file *manifest.FileMetadata, opts *IterOptions, _ internalIterOpts,
    40  	) (internalIterator, keyspan.FragmentIterator, error) {
    41  		f := *iters[file.FileNum]
    42  		f.lower = opts.GetLowerBound()
    43  		f.upper = opts.GetUpperBound()
    44  		return &f, nil, nil
    45  	}
    46  
    47  	datadriven.RunTest(t, "testdata/level_iter", func(t *testing.T, d *datadriven.TestData) string {
    48  		switch d.Cmd {
    49  		case "define":
    50  			iters = nil
    51  			var metas []*fileMetadata
    52  			for _, line := range strings.Split(d.Input, "\n") {
    53  				f := &fakeIter{}
    54  				for _, key := range strings.Fields(line) {
    55  					j := strings.Index(key, ":")
    56  					f.keys = append(f.keys, base.ParseInternalKey(key[:j]))
    57  					f.vals = append(f.vals, []byte(key[j+1:]))
    58  				}
    59  				iters = append(iters, f)
    60  
    61  				meta := (&fileMetadata{
    62  					FileNum: FileNum(len(metas)),
    63  				}).ExtendPointKeyBounds(
    64  					DefaultComparer.Compare,
    65  					f.keys[0],
    66  					f.keys[len(f.keys)-1],
    67  				)
    68  				meta.InitPhysicalBacking()
    69  				metas = append(metas, meta)
    70  			}
    71  			files = manifest.NewLevelSliceKeySorted(base.DefaultComparer.Compare, metas)
    72  
    73  			return ""
    74  
    75  		case "iter":
    76  			var opts IterOptions
    77  			for _, arg := range d.CmdArgs {
    78  				if len(arg.Vals) != 1 {
    79  					return fmt.Sprintf("%s: %s=<value>", d.Cmd, arg.Key)
    80  				}
    81  				switch arg.Key {
    82  				case "lower":
    83  					opts.LowerBound = []byte(arg.Vals[0])
    84  				case "upper":
    85  					opts.UpperBound = []byte(arg.Vals[0])
    86  				default:
    87  					return fmt.Sprintf("%s: unknown arg: %s", d.Cmd, arg.Key)
    88  				}
    89  			}
    90  
    91  			iter := newLevelIter(context.Background(), opts, testkeys.Comparer, newIters, files.Iter(), manifest.Level(level), internalIterOpts{})
    92  			defer iter.Close()
    93  			// Fake up the range deletion initialization.
    94  			iter.initRangeDel(new(keyspan.FragmentIterator))
    95  			iter.disableInvariants = true
    96  			return itertest.RunInternalIterCmd(t, d, iter, itertest.Verbose)
    97  
    98  		case "load":
    99  			// The "load" command allows testing the iterator options passed to load
   100  			// sstables.
   101  			//
   102  			// load <key> [lower=<key>] [upper=<key>]
   103  			var opts IterOptions
   104  			var key string
   105  			for _, arg := range d.CmdArgs {
   106  				if len(arg.Vals) == 0 {
   107  					key = arg.Key
   108  					continue
   109  				}
   110  				if len(arg.Vals) != 1 {
   111  					return fmt.Sprintf("%s: %s=<value>", d.Cmd, arg.Key)
   112  				}
   113  				switch arg.Key {
   114  				case "lower":
   115  					opts.LowerBound = []byte(arg.Vals[0])
   116  				case "upper":
   117  					opts.UpperBound = []byte(arg.Vals[0])
   118  				default:
   119  					return fmt.Sprintf("%s: unknown arg: %s", d.Cmd, arg.Key)
   120  				}
   121  			}
   122  
   123  			var tableOpts *IterOptions
   124  			newIters2 := func(
   125  				ctx context.Context, file *manifest.FileMetadata, opts *IterOptions,
   126  				internalOpts internalIterOpts,
   127  			) (internalIterator, keyspan.FragmentIterator, error) {
   128  				tableOpts = opts
   129  				return newIters(ctx, file, opts, internalOpts)
   130  			}
   131  
   132  			iter := newLevelIter(context.Background(), opts, testkeys.Comparer, newIters2, files.Iter(), manifest.Level(level), internalIterOpts{})
   133  			iter.SeekGE([]byte(key), base.SeekGEFlagsNone)
   134  			lower, upper := tableOpts.GetLowerBound(), tableOpts.GetUpperBound()
   135  			return fmt.Sprintf("[%s,%s]\n", lower, upper)
   136  
   137  		default:
   138  			return fmt.Sprintf("unknown command: %s", d.Cmd)
   139  		}
   140  	})
   141  }
   142  
   143  type levelIterTest struct {
   144  	cmp          base.Comparer
   145  	mem          vfs.FS
   146  	readers      []*sstable.Reader
   147  	metas        []*fileMetadata
   148  	itersCreated int
   149  }
   150  
   151  func newLevelIterTest() *levelIterTest {
   152  	lt := &levelIterTest{
   153  		cmp: *DefaultComparer,
   154  		mem: vfs.NewMem(),
   155  	}
   156  	lt.cmp.Split = func(a []byte) int { return len(a) }
   157  	return lt
   158  }
   159  
   160  func (lt *levelIterTest) newIters(
   161  	ctx context.Context, file *manifest.FileMetadata, opts *IterOptions, iio internalIterOpts,
   162  ) (internalIterator, keyspan.FragmentIterator, error) {
   163  	lt.itersCreated++
   164  	iter, err := lt.readers[file.FileNum].NewIterWithBlockPropertyFiltersAndContextEtc(
   165  		ctx, opts.LowerBound, opts.UpperBound, nil, false, true, iio.stats, sstable.CategoryAndQoS{},
   166  		nil, sstable.TrivialReaderProvider{Reader: lt.readers[file.FileNum]})
   167  	if err != nil {
   168  		return nil, nil, err
   169  	}
   170  	rangeDelIter, err := lt.readers[file.FileNum].NewRawRangeDelIter()
   171  	if err != nil {
   172  		return nil, nil, err
   173  	}
   174  	return iter, rangeDelIter, nil
   175  }
   176  
   177  func (lt *levelIterTest) runClear(d *datadriven.TestData) string {
   178  	lt.mem = vfs.NewMem()
   179  	for _, r := range lt.readers {
   180  		r.Close()
   181  	}
   182  	lt.readers = nil
   183  	lt.metas = nil
   184  	lt.itersCreated = 0
   185  	return ""
   186  }
   187  
   188  func (lt *levelIterTest) runBuild(d *datadriven.TestData) string {
   189  	fileNum := FileNum(len(lt.readers))
   190  	name := fmt.Sprint(fileNum)
   191  	f0, err := lt.mem.Create(name)
   192  	if err != nil {
   193  		return err.Error()
   194  	}
   195  
   196  	tableFormat := sstable.TableFormatRocksDBv2
   197  	for _, arg := range d.CmdArgs {
   198  		if arg.Key == "format" {
   199  			switch arg.Vals[0] {
   200  			case "rocksdbv2":
   201  				tableFormat = sstable.TableFormatRocksDBv2
   202  			case "pebblev2":
   203  				tableFormat = sstable.TableFormatPebblev2
   204  			}
   205  		}
   206  	}
   207  	fp := bloom.FilterPolicy(10)
   208  	w := sstable.NewWriter(objstorageprovider.NewFileWritable(f0), sstable.WriterOptions{
   209  		Comparer:     &lt.cmp,
   210  		FilterPolicy: fp,
   211  		TableFormat:  tableFormat,
   212  	})
   213  	var tombstones []keyspan.Span
   214  	f := keyspan.Fragmenter{
   215  		Cmp:    lt.cmp.Compare,
   216  		Format: lt.cmp.FormatKey,
   217  		Emit: func(fragmented keyspan.Span) {
   218  			tombstones = append(tombstones, fragmented)
   219  		},
   220  	}
   221  	for _, key := range strings.Split(d.Input, "\n") {
   222  		j := strings.Index(key, ":")
   223  		ikey := base.ParseInternalKey(key[:j])
   224  		value := []byte(key[j+1:])
   225  		switch ikey.Kind() {
   226  		case InternalKeyKindRangeDelete:
   227  			f.Add(rangedel.Decode(ikey, value, nil))
   228  		case InternalKeyKindRangeKeySet, InternalKeyKindRangeKeyUnset, InternalKeyKindRangeKeyDelete:
   229  			if err := w.AddRangeKey(ikey, value); err != nil {
   230  				return err.Error()
   231  			}
   232  		default:
   233  			if err := w.Add(ikey, value); err != nil {
   234  				return err.Error()
   235  			}
   236  		}
   237  	}
   238  	f.Finish()
   239  	for _, v := range tombstones {
   240  		if err := rangedel.Encode(&v, w.Add); err != nil {
   241  			return err.Error()
   242  		}
   243  	}
   244  	if err := w.Close(); err != nil {
   245  		return err.Error()
   246  	}
   247  	meta, err := w.Metadata()
   248  	if err != nil {
   249  		return err.Error()
   250  	}
   251  
   252  	f1, err := lt.mem.Open(name)
   253  	if err != nil {
   254  		return err.Error()
   255  	}
   256  	readable, err := sstable.NewSimpleReadable(f1)
   257  	if err != nil {
   258  		return err.Error()
   259  	}
   260  	r, err := sstable.NewReader(readable, sstable.ReaderOptions{
   261  		Filters: map[string]FilterPolicy{
   262  			fp.Name(): fp,
   263  		},
   264  	})
   265  	if err != nil {
   266  		return err.Error()
   267  	}
   268  	lt.readers = append(lt.readers, r)
   269  	m := &fileMetadata{FileNum: fileNum}
   270  	if meta.HasPointKeys {
   271  		m.ExtendPointKeyBounds(lt.cmp.Compare, meta.SmallestPoint, meta.LargestPoint)
   272  	}
   273  	if meta.HasRangeDelKeys {
   274  		m.ExtendPointKeyBounds(lt.cmp.Compare, meta.SmallestRangeDel, meta.LargestRangeDel)
   275  	}
   276  	if meta.HasRangeKeys {
   277  		m.ExtendRangeKeyBounds(lt.cmp.Compare, meta.SmallestRangeKey, meta.LargestRangeKey)
   278  	}
   279  	m.InitPhysicalBacking()
   280  	lt.metas = append(lt.metas, m)
   281  
   282  	var buf bytes.Buffer
   283  	for _, f := range lt.metas {
   284  		fmt.Fprintf(&buf, "%d: %s-%s\n", f.FileNum, f.Smallest, f.Largest)
   285  	}
   286  	return buf.String()
   287  }
   288  
   289  func TestLevelIterBoundaries(t *testing.T) {
   290  	lt := newLevelIterTest()
   291  	defer lt.runClear(nil)
   292  
   293  	var iter *levelIter
   294  	datadriven.RunTest(t, "testdata/level_iter_boundaries", func(t *testing.T, d *datadriven.TestData) string {
   295  		switch d.Cmd {
   296  		case "clear":
   297  			return lt.runClear(d)
   298  
   299  		case "build":
   300  			return lt.runBuild(d)
   301  
   302  		case "iter":
   303  			// The save and continue parameters allow us to save the iterator
   304  			// for later continued use.
   305  			save := false
   306  			cont := false
   307  			for _, arg := range d.CmdArgs {
   308  				switch arg.Key {
   309  				case "save":
   310  					save = true
   311  				case "continue":
   312  					cont = true
   313  				default:
   314  					return fmt.Sprintf("%s: unknown arg: %s", d.Cmd, arg.Key)
   315  				}
   316  			}
   317  			if !cont && iter != nil {
   318  				return "preceding iter was not closed"
   319  			}
   320  			if cont && iter == nil {
   321  				return "no existing iter"
   322  			}
   323  			if iter == nil {
   324  				slice := manifest.NewLevelSliceKeySorted(lt.cmp.Compare, lt.metas)
   325  				iter = newLevelIter(context.Background(), IterOptions{}, testkeys.Comparer, lt.newIters, slice.Iter(), manifest.Level(level), internalIterOpts{})
   326  				// Fake up the range deletion initialization.
   327  				iter.initRangeDel(new(keyspan.FragmentIterator))
   328  			}
   329  			if !save {
   330  				defer func() {
   331  					iter.Close()
   332  					iter = nil
   333  				}()
   334  			}
   335  			return itertest.RunInternalIterCmd(t, d, iter, itertest.Verbose)
   336  
   337  		case "file-pos":
   338  			// Returns the FileNum at which the iterator is positioned.
   339  			if iter == nil {
   340  				return "nil levelIter"
   341  			}
   342  			if iter.iterFile == nil {
   343  				return "nil iterFile"
   344  			}
   345  			return fmt.Sprintf("file %d", iter.iterFile.FileNum)
   346  
   347  		default:
   348  			return fmt.Sprintf("unknown command: %s", d.Cmd)
   349  		}
   350  	})
   351  }
   352  
   353  // levelIterTestIter allows a datadriven test to use runInternalIterCmd and
   354  // perform parallel operations on both both a levelIter and rangeDelIter.
   355  type levelIterTestIter struct {
   356  	*levelIter
   357  	rangeDelIter keyspan.FragmentIterator
   358  }
   359  
   360  func (i *levelIterTestIter) rangeDelSeek(
   361  	key []byte, ikey *InternalKey, val base.LazyValue, dir int,
   362  ) (*InternalKey, base.LazyValue) {
   363  	var tombstone keyspan.Span
   364  	if i.rangeDelIter != nil {
   365  		var t *keyspan.Span
   366  		if dir < 0 {
   367  			t = keyspan.SeekLE(i.levelIter.cmp, i.rangeDelIter, key)
   368  		} else {
   369  			t = i.rangeDelIter.SeekGE(key)
   370  		}
   371  		if t != nil {
   372  			tombstone = t.Visible(1000)
   373  		}
   374  	}
   375  	if ikey == nil {
   376  		return &InternalKey{
   377  			UserKey: []byte(fmt.Sprintf("./%s", tombstone)),
   378  		}, base.LazyValue{}
   379  	}
   380  	return &InternalKey{
   381  		UserKey: []byte(fmt.Sprintf("%s/%s", ikey.UserKey, tombstone)),
   382  		Trailer: ikey.Trailer,
   383  	}, val
   384  }
   385  
   386  func (i *levelIterTestIter) String() string {
   387  	return "level-iter-test"
   388  }
   389  
   390  func (i *levelIterTestIter) SeekGE(
   391  	key []byte, flags base.SeekGEFlags,
   392  ) (*InternalKey, base.LazyValue) {
   393  	ikey, val := i.levelIter.SeekGE(key, flags)
   394  	return i.rangeDelSeek(key, ikey, val, 1)
   395  }
   396  
   397  func (i *levelIterTestIter) SeekPrefixGE(
   398  	prefix, key []byte, flags base.SeekGEFlags,
   399  ) (*base.InternalKey, base.LazyValue) {
   400  	ikey, val := i.levelIter.SeekPrefixGE(prefix, key, flags)
   401  	return i.rangeDelSeek(key, ikey, val, 1)
   402  }
   403  
   404  func (i *levelIterTestIter) SeekLT(
   405  	key []byte, flags base.SeekLTFlags,
   406  ) (*InternalKey, base.LazyValue) {
   407  	ikey, val := i.levelIter.SeekLT(key, flags)
   408  	return i.rangeDelSeek(key, ikey, val, -1)
   409  }
   410  
   411  func TestLevelIterSeek(t *testing.T) {
   412  	lt := newLevelIterTest()
   413  	defer lt.runClear(nil)
   414  
   415  	datadriven.RunTest(t, "testdata/level_iter_seek", func(t *testing.T, d *datadriven.TestData) string {
   416  		switch d.Cmd {
   417  		case "clear":
   418  			return lt.runClear(d)
   419  
   420  		case "build":
   421  			return lt.runBuild(d)
   422  
   423  		case "iter":
   424  			var stats base.InternalIteratorStats
   425  			slice := manifest.NewLevelSliceKeySorted(lt.cmp.Compare, lt.metas)
   426  			iter := &levelIterTestIter{levelIter: &levelIter{}}
   427  			iter.init(context.Background(), IterOptions{}, testkeys.Comparer, lt.newIters, slice.Iter(),
   428  				manifest.Level(level), internalIterOpts{stats: &stats})
   429  			defer iter.Close()
   430  			iter.initRangeDel(&iter.rangeDelIter)
   431  			return itertest.RunInternalIterCmd(t, d, iter, itertest.Verbose, itertest.WithStats(&stats))
   432  
   433  		case "iters-created":
   434  			return fmt.Sprintf("%d", lt.itersCreated)
   435  		default:
   436  			return fmt.Sprintf("unknown command: %s", d.Cmd)
   437  		}
   438  	})
   439  }
   440  
   441  func buildLevelIterTables(
   442  	b *testing.B, blockSize, restartInterval, count int,
   443  ) ([]*sstable.Reader, manifest.LevelSlice, [][]byte, func()) {
   444  	mem := vfs.NewMem()
   445  	files := make([]vfs.File, count)
   446  	for i := range files {
   447  		f, err := mem.Create(fmt.Sprintf("bench%d", i))
   448  		if err != nil {
   449  			b.Fatal(err)
   450  		}
   451  		files[i] = f
   452  	}
   453  
   454  	writers := make([]*sstable.Writer, len(files))
   455  	for i := range files {
   456  		writers[i] = sstable.NewWriter(objstorageprovider.NewFileWritable(files[i]), sstable.WriterOptions{
   457  			BlockRestartInterval: restartInterval,
   458  			BlockSize:            blockSize,
   459  			Compression:          NoCompression,
   460  		})
   461  	}
   462  
   463  	var keys [][]byte
   464  	var i int
   465  	const targetSize = 2 << 20
   466  	for _, w := range writers {
   467  		for ; w.EstimatedSize() < targetSize; i++ {
   468  			key := []byte(fmt.Sprintf("%08d", i))
   469  			keys = append(keys, key)
   470  			ikey := base.MakeInternalKey(key, 0, InternalKeyKindSet)
   471  			w.Add(ikey, nil)
   472  		}
   473  		if err := w.Close(); err != nil {
   474  			b.Fatal(err)
   475  		}
   476  	}
   477  
   478  	opts := sstable.ReaderOptions{Cache: NewCache(128 << 20), Comparer: DefaultComparer}
   479  	defer opts.Cache.Unref()
   480  	readers := make([]*sstable.Reader, len(files))
   481  	for i := range files {
   482  		f, err := mem.Open(fmt.Sprintf("bench%d", i))
   483  		if err != nil {
   484  			b.Fatal(err)
   485  		}
   486  		readable, err := sstable.NewSimpleReadable(f)
   487  		if err != nil {
   488  			b.Fatal(err)
   489  		}
   490  		readers[i], err = sstable.NewReader(readable, opts)
   491  		if err != nil {
   492  			b.Fatal(err)
   493  		}
   494  	}
   495  
   496  	cleanup := func() {
   497  		for _, r := range readers {
   498  			require.NoError(b, r.Close())
   499  		}
   500  	}
   501  
   502  	meta := make([]*fileMetadata, len(readers))
   503  	for i := range readers {
   504  		iter, err := readers[i].NewIter(nil /* lower */, nil /* upper */)
   505  		require.NoError(b, err)
   506  		smallest, _ := iter.First()
   507  		meta[i] = &fileMetadata{}
   508  		meta[i].FileNum = FileNum(i)
   509  		largest, _ := iter.Last()
   510  		meta[i].ExtendPointKeyBounds(opts.Comparer.Compare, (*smallest).Clone(), (*largest).Clone())
   511  		meta[i].InitPhysicalBacking()
   512  	}
   513  	slice := manifest.NewLevelSliceKeySorted(base.DefaultComparer.Compare, meta)
   514  	return readers, slice, keys, cleanup
   515  }
   516  
   517  func BenchmarkLevelIterSeekGE(b *testing.B) {
   518  	const blockSize = 32 << 10
   519  
   520  	for _, restartInterval := range []int{16} {
   521  		b.Run(fmt.Sprintf("restart=%d", restartInterval),
   522  			func(b *testing.B) {
   523  				for _, count := range []int{5} {
   524  					b.Run(fmt.Sprintf("count=%d", count),
   525  						func(b *testing.B) {
   526  							readers, metas, keys, cleanup := buildLevelIterTables(b, blockSize, restartInterval, count)
   527  							defer cleanup()
   528  							newIters := func(
   529  								_ context.Context, file *manifest.FileMetadata, _ *IterOptions, _ internalIterOpts,
   530  							) (internalIterator, keyspan.FragmentIterator, error) {
   531  								iter, err := readers[file.FileNum].NewIter(nil /* lower */, nil /* upper */)
   532  								return iter, nil, err
   533  							}
   534  							l := newLevelIter(context.Background(), IterOptions{}, DefaultComparer, newIters, metas.Iter(), manifest.Level(level), internalIterOpts{})
   535  							rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
   536  
   537  							b.ResetTimer()
   538  							for i := 0; i < b.N; i++ {
   539  								l.SeekGE(keys[rng.Intn(len(keys))], base.SeekGEFlagsNone)
   540  							}
   541  							l.Close()
   542  						})
   543  				}
   544  			})
   545  	}
   546  }
   547  
   548  // A benchmark that simulates the behavior of a levelIter being used as part
   549  // of a mergingIter where narrow bounds are repeatedly set and used to Seek
   550  // and then iterate over the keys within the bounds. This resembles MVCC
   551  // scanning by CockroachDB when doing a lookup/index join with a large number
   552  // of left rows, that are batched and reuse the same iterator, and which can
   553  // have good locality of access. This results in the successive bounds being
   554  // in the same file.
   555  func BenchmarkLevelIterSeqSeekGEWithBounds(b *testing.B) {
   556  	const blockSize = 32 << 10
   557  
   558  	for _, restartInterval := range []int{16} {
   559  		b.Run(fmt.Sprintf("restart=%d", restartInterval),
   560  			func(b *testing.B) {
   561  				for _, count := range []int{5} {
   562  					b.Run(fmt.Sprintf("count=%d", count),
   563  						func(b *testing.B) {
   564  							readers, metas, keys, cleanup :=
   565  								buildLevelIterTables(b, blockSize, restartInterval, count)
   566  							defer cleanup()
   567  							// This newIters is cheaper than in practice since it does not do
   568  							// tableCacheShard.findNode.
   569  							newIters := func(
   570  								_ context.Context, file *manifest.FileMetadata, opts *IterOptions, _ internalIterOpts,
   571  							) (internalIterator, keyspan.FragmentIterator, error) {
   572  								iter, err := readers[file.FileNum].NewIter(
   573  									opts.LowerBound, opts.UpperBound)
   574  								return iter, nil, err
   575  							}
   576  							l := newLevelIter(context.Background(), IterOptions{}, DefaultComparer, newIters, metas.Iter(), manifest.Level(level), internalIterOpts{})
   577  							// Fake up the range deletion initialization, to resemble the usage
   578  							// in a mergingIter.
   579  							l.initRangeDel(new(keyspan.FragmentIterator))
   580  							keyCount := len(keys)
   581  							b.ResetTimer()
   582  							for i := 0; i < b.N; i++ {
   583  								pos := i % (keyCount - 1)
   584  								l.SetBounds(keys[pos], keys[pos+1])
   585  								// SeekGE will return keys[pos].
   586  								k, _ := l.SeekGE(keys[pos], base.SeekGEFlagsNone)
   587  								// Next() will get called once and return nil.
   588  								for k != nil {
   589  									k, _ = l.Next()
   590  								}
   591  							}
   592  							l.Close()
   593  						})
   594  				}
   595  			})
   596  	}
   597  }
   598  
   599  // BenchmarkLevelIterSeqSeekPrefixGE simulates the behavior of a levelIter
   600  // being used as part of a mergingIter where SeekPrefixGE is used to seek in a
   601  // monotonically increasing manner. This resembles key-value lookups done by
   602  // CockroachDB when evaluating Put operations.
   603  func BenchmarkLevelIterSeqSeekPrefixGE(b *testing.B) {
   604  	const blockSize = 32 << 10
   605  	const restartInterval = 16
   606  	readers, metas, keys, cleanup :=
   607  		buildLevelIterTables(b, blockSize, restartInterval, 5)
   608  	defer cleanup()
   609  	// This newIters is cheaper than in practice since it does not do
   610  	// tableCacheShard.findNode.
   611  	newIters := func(
   612  		_ context.Context, file *manifest.FileMetadata, opts *IterOptions, _ internalIterOpts,
   613  	) (internalIterator, keyspan.FragmentIterator, error) {
   614  		iter, err := readers[file.FileNum].NewIter(
   615  			opts.LowerBound, opts.UpperBound)
   616  		return iter, nil, err
   617  	}
   618  
   619  	for _, skip := range []int{1, 2, 4, 8, 16} {
   620  		for _, useNext := range []bool{false, true} {
   621  			b.Run(fmt.Sprintf("skip=%d/use-next=%t", skip, useNext),
   622  				func(b *testing.B) {
   623  					l := newLevelIter(context.Background(), IterOptions{}, testkeys.Comparer, newIters, metas.Iter(), manifest.Level(level), internalIterOpts{})
   624  					// Fake up the range deletion initialization, to resemble the usage
   625  					// in a mergingIter.
   626  					l.initRangeDel(new(keyspan.FragmentIterator))
   627  					keyCount := len(keys)
   628  					pos := 0
   629  					l.SeekPrefixGE(keys[pos], keys[pos], base.SeekGEFlagsNone)
   630  					b.ResetTimer()
   631  					for i := 0; i < b.N; i++ {
   632  						pos += skip
   633  						var flags base.SeekGEFlags
   634  						if useNext {
   635  							flags = flags.EnableTrySeekUsingNext()
   636  						}
   637  						if pos >= keyCount {
   638  							pos = 0
   639  							flags = flags.DisableTrySeekUsingNext()
   640  						}
   641  						// SeekPrefixGE will return keys[pos].
   642  						l.SeekPrefixGE(keys[pos], keys[pos], flags)
   643  					}
   644  					b.StopTimer()
   645  					l.Close()
   646  				})
   647  		}
   648  	}
   649  }
   650  
   651  func BenchmarkLevelIterNext(b *testing.B) {
   652  	const blockSize = 32 << 10
   653  
   654  	for _, restartInterval := range []int{16} {
   655  		b.Run(fmt.Sprintf("restart=%d", restartInterval),
   656  			func(b *testing.B) {
   657  				for _, count := range []int{5} {
   658  					b.Run(fmt.Sprintf("count=%d", count),
   659  						func(b *testing.B) {
   660  							readers, metas, _, cleanup := buildLevelIterTables(b, blockSize, restartInterval, count)
   661  							defer cleanup()
   662  							newIters := func(
   663  								_ context.Context, file *manifest.FileMetadata, _ *IterOptions, _ internalIterOpts,
   664  							) (internalIterator, keyspan.FragmentIterator, error) {
   665  								iter, err := readers[file.FileNum].NewIter(nil /* lower */, nil /* upper */)
   666  								return iter, nil, err
   667  							}
   668  							l := newLevelIter(context.Background(), IterOptions{}, testkeys.Comparer, newIters, metas.Iter(), manifest.Level(level), internalIterOpts{})
   669  
   670  							b.ResetTimer()
   671  							for i := 0; i < b.N; i++ {
   672  								key, _ := l.Next()
   673  								if key == nil {
   674  									key, _ = l.First()
   675  								}
   676  								_ = key
   677  							}
   678  							l.Close()
   679  						})
   680  				}
   681  			})
   682  	}
   683  }
   684  
   685  func BenchmarkLevelIterPrev(b *testing.B) {
   686  	const blockSize = 32 << 10
   687  
   688  	for _, restartInterval := range []int{16} {
   689  		b.Run(fmt.Sprintf("restart=%d", restartInterval),
   690  			func(b *testing.B) {
   691  				for _, count := range []int{5} {
   692  					b.Run(fmt.Sprintf("count=%d", count),
   693  						func(b *testing.B) {
   694  							readers, metas, _, cleanup := buildLevelIterTables(b, blockSize, restartInterval, count)
   695  							defer cleanup()
   696  							newIters := func(
   697  								_ context.Context, file *manifest.FileMetadata, _ *IterOptions, _ internalIterOpts,
   698  							) (internalIterator, keyspan.FragmentIterator, error) {
   699  								iter, err := readers[file.FileNum].NewIter(nil /* lower */, nil /* upper */)
   700  								return iter, nil, err
   701  							}
   702  							l := newLevelIter(context.Background(), IterOptions{}, DefaultComparer, newIters, metas.Iter(), manifest.Level(level), internalIterOpts{})
   703  
   704  							b.ResetTimer()
   705  							for i := 0; i < b.N; i++ {
   706  								key, _ := l.Prev()
   707  								if key == nil {
   708  									key, _ = l.Last()
   709  								}
   710  								_ = key
   711  							}
   712  							l.Close()
   713  						})
   714  				}
   715  			})
   716  	}
   717  }