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

     1  /*
     2   * Copyright 2017 Dgraph Labs, Inc. and Contributors
     3   * Modifications copyright (C) 2017 Andy Kimball and Contributors
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   */
    17  
    18  package arenaskl
    19  
    20  import (
    21  	"bytes"
    22  	"encoding/binary"
    23  	"fmt"
    24  	"strconv"
    25  	"sync"
    26  	"sync/atomic"
    27  	"testing"
    28  	"time"
    29  
    30  	"github.com/cockroachdb/pebble/internal/base"
    31  	"github.com/stretchr/testify/require"
    32  	"golang.org/x/exp/rand"
    33  )
    34  
    35  const arenaSize = 1 << 20
    36  
    37  // iterAdapter adapts the new Iterator API which returns the key and value from
    38  // positioning methods (Seek*, First, Last, Next, Prev) to the old API which
    39  // returned a boolean corresponding to Valid. Only used by test code.
    40  type iterAdapter struct {
    41  	*Iterator
    42  	key *base.InternalKey
    43  	val []byte
    44  }
    45  
    46  func newIterAdapter(iter *Iterator) *iterAdapter {
    47  	return &iterAdapter{
    48  		Iterator: iter,
    49  	}
    50  }
    51  
    52  func (i *iterAdapter) update(key *base.InternalKey, val base.LazyValue) bool {
    53  	i.key = key
    54  	i.val = val.InPlaceValue()
    55  	return i.key != nil
    56  }
    57  
    58  func (i *iterAdapter) String() string {
    59  	return "iter-adapter"
    60  }
    61  
    62  func (i *iterAdapter) SeekGE(key []byte, flags base.SeekGEFlags) bool {
    63  	return i.update(i.Iterator.SeekGE(key, flags))
    64  }
    65  
    66  func (i *iterAdapter) SeekPrefixGE(prefix, key []byte, flags base.SeekGEFlags) bool {
    67  	return i.update(i.Iterator.SeekPrefixGE(prefix, key, flags))
    68  }
    69  
    70  func (i *iterAdapter) SeekLT(key []byte, flags base.SeekLTFlags) bool {
    71  	return i.update(i.Iterator.SeekLT(key, flags))
    72  }
    73  
    74  func (i *iterAdapter) First() bool {
    75  	return i.update(i.Iterator.First())
    76  }
    77  
    78  func (i *iterAdapter) Last() bool {
    79  	return i.update(i.Iterator.Last())
    80  }
    81  
    82  func (i *iterAdapter) Next() bool {
    83  	return i.update(i.Iterator.Next())
    84  }
    85  
    86  func (i *iterAdapter) Prev() bool {
    87  	return i.update(i.Iterator.Prev())
    88  }
    89  
    90  func (i *iterAdapter) Key() base.InternalKey {
    91  	return *i.key
    92  }
    93  
    94  func (i *iterAdapter) Value() []byte {
    95  	return i.val
    96  }
    97  
    98  func (i *iterAdapter) Valid() bool {
    99  	return i.key != nil
   100  }
   101  
   102  func makeIntKey(i int) base.InternalKey {
   103  	return base.InternalKey{UserKey: []byte(fmt.Sprintf("%05d", i))}
   104  }
   105  
   106  func makeKey(s string) []byte {
   107  	return []byte(s)
   108  }
   109  
   110  func makeIkey(s string) base.InternalKey {
   111  	return base.InternalKey{UserKey: []byte(s)}
   112  }
   113  
   114  func makeValue(i int) []byte {
   115  	return []byte(fmt.Sprintf("v%05d", i))
   116  }
   117  
   118  func makeInserterAdd(s *Skiplist) func(key base.InternalKey, value []byte) error {
   119  	ins := &Inserter{}
   120  	return func(key base.InternalKey, value []byte) error {
   121  		return ins.Add(s, key, value)
   122  	}
   123  }
   124  
   125  // length iterates over skiplist to give exact size.
   126  func length(s *Skiplist) int {
   127  	count := 0
   128  
   129  	it := newIterAdapter(s.NewIter(nil, nil))
   130  	for valid := it.First(); valid; valid = it.Next() {
   131  		count++
   132  	}
   133  
   134  	return count
   135  }
   136  
   137  // length iterates over skiplist in reverse order to give exact size.
   138  func lengthRev(s *Skiplist) int {
   139  	count := 0
   140  
   141  	it := newIterAdapter(s.NewIter(nil, nil))
   142  	for valid := it.Last(); valid; valid = it.Prev() {
   143  		count++
   144  	}
   145  
   146  	return count
   147  }
   148  
   149  func TestEmpty(t *testing.T) {
   150  	key := makeKey("aaa")
   151  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   152  	it := newIterAdapter(l.NewIter(nil, nil))
   153  
   154  	require.False(t, it.Valid())
   155  
   156  	it.First()
   157  	require.False(t, it.Valid())
   158  
   159  	it.Last()
   160  	require.False(t, it.Valid())
   161  
   162  	require.False(t, it.SeekGE(key, base.SeekGEFlagsNone))
   163  	require.False(t, it.Valid())
   164  }
   165  
   166  func TestFull(t *testing.T) {
   167  	l := NewSkiplist(newArena(1000), bytes.Compare)
   168  
   169  	foundArenaFull := false
   170  	for i := 0; i < 100; i++ {
   171  		err := l.Add(makeIntKey(i), makeValue(i))
   172  		if err == ErrArenaFull {
   173  			foundArenaFull = true
   174  			break
   175  		}
   176  	}
   177  
   178  	require.True(t, foundArenaFull)
   179  
   180  	err := l.Add(makeIkey("someval"), nil)
   181  	require.Equal(t, ErrArenaFull, err)
   182  }
   183  
   184  // TestBasic tests single-threaded seeks and adds.
   185  func TestBasic(t *testing.T) {
   186  	for _, inserter := range []bool{false, true} {
   187  		t.Run(fmt.Sprintf("inserter=%t", inserter), func(t *testing.T) {
   188  			l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   189  			it := newIterAdapter(l.NewIter(nil, nil))
   190  
   191  			add := l.Add
   192  			if inserter {
   193  				add = makeInserterAdd(l)
   194  			}
   195  
   196  			// Try adding values.
   197  			add(makeIkey("key1"), makeValue(1))
   198  			add(makeIkey("key3"), makeValue(3))
   199  			add(makeIkey("key2"), makeValue(2))
   200  
   201  			require.True(t, it.SeekGE(makeKey("key"), base.SeekGEFlagsNone))
   202  			require.True(t, it.Valid())
   203  			require.NotEqual(t, "key", it.Key().UserKey)
   204  
   205  			require.True(t, it.SeekGE(makeKey("key1"), base.SeekGEFlagsNone))
   206  			require.EqualValues(t, "key1", it.Key().UserKey)
   207  			require.EqualValues(t, makeValue(1), it.Value())
   208  
   209  			require.True(t, it.SeekGE(makeKey("key2"), base.SeekGEFlagsNone))
   210  			require.EqualValues(t, "key2", it.Key().UserKey)
   211  			require.EqualValues(t, makeValue(2), it.Value())
   212  
   213  			require.True(t, it.SeekGE(makeKey("key3"), base.SeekGEFlagsNone))
   214  			require.EqualValues(t, "key3", it.Key().UserKey)
   215  			require.EqualValues(t, makeValue(3), it.Value())
   216  
   217  			key := makeIkey("a")
   218  			key.SetSeqNum(1)
   219  			add(key, nil)
   220  			key.SetSeqNum(2)
   221  			add(key, nil)
   222  
   223  			require.True(t, it.SeekGE(makeKey("a"), base.SeekGEFlagsNone))
   224  			require.True(t, it.Valid())
   225  			require.EqualValues(t, "a", it.Key().UserKey)
   226  			require.EqualValues(t, 2, it.Key().SeqNum())
   227  
   228  			require.True(t, it.Next())
   229  			require.True(t, it.Valid())
   230  			require.EqualValues(t, "a", it.Key().UserKey)
   231  			require.EqualValues(t, 1, it.Key().SeqNum())
   232  
   233  			key = makeIkey("b")
   234  			key.SetSeqNum(2)
   235  			add(key, nil)
   236  			key.SetSeqNum(1)
   237  			add(key, nil)
   238  
   239  			require.True(t, it.SeekGE(makeKey("b"), base.SeekGEFlagsNone))
   240  			require.True(t, it.Valid())
   241  			require.EqualValues(t, "b", it.Key().UserKey)
   242  			require.EqualValues(t, 2, it.Key().SeqNum())
   243  
   244  			require.True(t, it.Next())
   245  			require.True(t, it.Valid())
   246  			require.EqualValues(t, "b", it.Key().UserKey)
   247  			require.EqualValues(t, 1, it.Key().SeqNum())
   248  		})
   249  	}
   250  }
   251  
   252  // TestConcurrentBasic tests concurrent writes followed by concurrent reads.
   253  func TestConcurrentBasic(t *testing.T) {
   254  	const n = 1000
   255  
   256  	for _, inserter := range []bool{false, true} {
   257  		t.Run(fmt.Sprintf("inserter=%t", inserter), func(t *testing.T) {
   258  			// Set testing flag to make it easier to trigger unusual race conditions.
   259  			l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   260  			l.testing = true
   261  
   262  			var wg sync.WaitGroup
   263  			for i := 0; i < n; i++ {
   264  				wg.Add(1)
   265  				go func(i int) {
   266  					defer wg.Done()
   267  
   268  					if inserter {
   269  						var ins Inserter
   270  						ins.Add(l, makeIntKey(i), makeValue(i))
   271  					} else {
   272  						l.Add(makeIntKey(i), makeValue(i))
   273  					}
   274  				}(i)
   275  			}
   276  			wg.Wait()
   277  
   278  			// Check values. Concurrent reads.
   279  			for i := 0; i < n; i++ {
   280  				wg.Add(1)
   281  				go func(i int) {
   282  					defer wg.Done()
   283  
   284  					it := newIterAdapter(l.NewIter(nil, nil))
   285  					require.True(t, it.SeekGE(makeKey(fmt.Sprintf("%05d", i)), base.SeekGEFlagsNone))
   286  					require.EqualValues(t, fmt.Sprintf("%05d", i), it.Key().UserKey)
   287  				}(i)
   288  			}
   289  			wg.Wait()
   290  			require.Equal(t, n, length(l))
   291  			require.Equal(t, n, lengthRev(l))
   292  		})
   293  	}
   294  }
   295  
   296  // TestConcurrentOneKey will read while writing to one single key.
   297  func TestConcurrentOneKey(t *testing.T) {
   298  	const n = 100
   299  	key := makeKey("thekey")
   300  	ikey := makeIkey("thekey")
   301  
   302  	for _, inserter := range []bool{false, true} {
   303  		t.Run(fmt.Sprintf("inserter=%t", inserter), func(t *testing.T) {
   304  			// Set testing flag to make it easier to trigger unusual race conditions.
   305  			l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   306  			l.testing = true
   307  
   308  			var wg sync.WaitGroup
   309  			writeDone := make(chan struct{}, 1)
   310  			for i := 0; i < n; i++ {
   311  				wg.Add(1)
   312  				go func(i int) {
   313  					defer func() {
   314  						wg.Done()
   315  						select {
   316  						case writeDone <- struct{}{}:
   317  						default:
   318  						}
   319  					}()
   320  
   321  					if inserter {
   322  						var ins Inserter
   323  						ins.Add(l, ikey, makeValue(i))
   324  					} else {
   325  						l.Add(ikey, makeValue(i))
   326  					}
   327  				}(i)
   328  			}
   329  			// Wait until at least some write made it such that reads return a value.
   330  			<-writeDone
   331  			var sawValue atomic.Int32
   332  			for i := 0; i < n; i++ {
   333  				wg.Add(1)
   334  				go func() {
   335  					defer wg.Done()
   336  
   337  					it := newIterAdapter(l.NewIter(nil, nil))
   338  					it.SeekGE(key, base.SeekGEFlagsNone)
   339  					require.True(t, it.Valid())
   340  					require.True(t, bytes.Equal(key, it.Key().UserKey))
   341  
   342  					sawValue.Add(1)
   343  					v, err := strconv.Atoi(string(it.Value()[1:]))
   344  					require.NoError(t, err)
   345  					require.True(t, 0 <= v && v < n)
   346  				}()
   347  			}
   348  			wg.Wait()
   349  			require.Equal(t, int32(n), sawValue.Load())
   350  			require.Equal(t, 1, length(l))
   351  			require.Equal(t, 1, lengthRev(l))
   352  		})
   353  	}
   354  }
   355  
   356  func TestSkiplistAdd(t *testing.T) {
   357  	for _, inserter := range []bool{false, true} {
   358  		t.Run(fmt.Sprintf("inserter=%t", inserter), func(t *testing.T) {
   359  			l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   360  			it := newIterAdapter(l.NewIter(nil, nil))
   361  
   362  			add := l.Add
   363  			if inserter {
   364  				add = makeInserterAdd(l)
   365  			}
   366  
   367  			// Add nil key and value (treated same as empty).
   368  			err := add(base.InternalKey{}, nil)
   369  			require.Nil(t, err)
   370  			require.True(t, it.SeekGE([]byte{}, base.SeekGEFlagsNone))
   371  			require.EqualValues(t, []byte{}, it.Key().UserKey)
   372  			require.EqualValues(t, []byte{}, it.Value())
   373  
   374  			l = NewSkiplist(newArena(arenaSize), bytes.Compare)
   375  			it = newIterAdapter(l.NewIter(nil, nil))
   376  
   377  			add = l.Add
   378  			if inserter {
   379  				add = makeInserterAdd(l)
   380  			}
   381  
   382  			// Add empty key and value (treated same as nil).
   383  			err = add(makeIkey(""), []byte{})
   384  			require.Nil(t, err)
   385  			require.True(t, it.SeekGE([]byte{}, base.SeekGEFlagsNone))
   386  			require.EqualValues(t, []byte{}, it.Key().UserKey)
   387  			require.EqualValues(t, []byte{}, it.Value())
   388  
   389  			// Add to empty list.
   390  			err = add(makeIntKey(2), makeValue(2))
   391  			require.Nil(t, err)
   392  			require.True(t, it.SeekGE(makeKey("00002"), base.SeekGEFlagsNone))
   393  			require.EqualValues(t, "00002", it.Key().UserKey)
   394  			require.EqualValues(t, makeValue(2), it.Value())
   395  
   396  			// Add first element in non-empty list.
   397  			err = add(makeIntKey(1), makeValue(1))
   398  			require.Nil(t, err)
   399  			require.True(t, it.SeekGE(makeKey("00001"), base.SeekGEFlagsNone))
   400  			require.EqualValues(t, "00001", it.Key().UserKey)
   401  			require.EqualValues(t, makeValue(1), it.Value())
   402  
   403  			// Add last element in non-empty list.
   404  			err = add(makeIntKey(4), makeValue(4))
   405  			require.Nil(t, err)
   406  			require.True(t, it.SeekGE(makeKey("00004"), base.SeekGEFlagsNone))
   407  			require.EqualValues(t, "00004", it.Key().UserKey)
   408  			require.EqualValues(t, makeValue(4), it.Value())
   409  
   410  			// Add element in middle of list.
   411  			err = add(makeIntKey(3), makeValue(3))
   412  			require.Nil(t, err)
   413  			require.True(t, it.SeekGE(makeKey("00003"), base.SeekGEFlagsNone))
   414  			require.EqualValues(t, "00003", it.Key().UserKey)
   415  			require.EqualValues(t, makeValue(3), it.Value())
   416  
   417  			// Try to add element that already exists.
   418  			err = add(makeIntKey(2), nil)
   419  			require.Equal(t, ErrRecordExists, err)
   420  			require.EqualValues(t, "00003", it.Key().UserKey)
   421  			require.EqualValues(t, makeValue(3), it.Value())
   422  
   423  			require.Equal(t, 5, length(l))
   424  			require.Equal(t, 5, lengthRev(l))
   425  		})
   426  	}
   427  }
   428  
   429  // TestConcurrentAdd races between adding same nodes.
   430  func TestConcurrentAdd(t *testing.T) {
   431  	for _, inserter := range []bool{false, true} {
   432  		t.Run(fmt.Sprintf("inserter=%t", inserter), func(t *testing.T) {
   433  			const n = 100
   434  
   435  			// Set testing flag to make it easier to trigger unusual race conditions.
   436  			l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   437  			l.testing = true
   438  
   439  			start := make([]sync.WaitGroup, n)
   440  			end := make([]sync.WaitGroup, n)
   441  
   442  			for i := 0; i < n; i++ {
   443  				start[i].Add(1)
   444  				end[i].Add(2)
   445  			}
   446  
   447  			for f := 0; f < 2; f++ {
   448  				go func(f int) {
   449  					it := newIterAdapter(l.NewIter(nil, nil))
   450  					add := l.Add
   451  					if inserter {
   452  						add = makeInserterAdd(l)
   453  					}
   454  
   455  					for i := 0; i < n; i++ {
   456  						start[i].Wait()
   457  
   458  						key := makeIntKey(i)
   459  						if add(key, nil) == nil {
   460  							require.True(t, it.SeekGE(key.UserKey, base.SeekGEFlagsNone))
   461  							require.EqualValues(t, key, it.Key())
   462  						}
   463  
   464  						end[i].Done()
   465  					}
   466  				}(f)
   467  			}
   468  
   469  			for i := 0; i < n; i++ {
   470  				start[i].Done()
   471  				end[i].Wait()
   472  			}
   473  
   474  			require.Equal(t, n, length(l))
   475  			require.Equal(t, n, lengthRev(l))
   476  		})
   477  	}
   478  }
   479  
   480  // TestIteratorNext tests a basic iteration over all nodes from the beginning.
   481  func TestIteratorNext(t *testing.T) {
   482  	const n = 100
   483  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   484  	it := newIterAdapter(l.NewIter(nil, nil))
   485  
   486  	require.False(t, it.Valid())
   487  
   488  	it.First()
   489  	require.False(t, it.Valid())
   490  
   491  	for i := n - 1; i >= 0; i-- {
   492  		l.Add(makeIntKey(i), makeValue(i))
   493  	}
   494  
   495  	it.First()
   496  	for i := 0; i < n; i++ {
   497  		require.True(t, it.Valid())
   498  		require.EqualValues(t, makeIntKey(i), it.Key())
   499  		require.EqualValues(t, makeValue(i), it.Value())
   500  		it.Next()
   501  	}
   502  	require.False(t, it.Valid())
   503  }
   504  
   505  // TestIteratorPrev tests a basic iteration over all nodes from the end.
   506  func TestIteratorPrev(t *testing.T) {
   507  	const n = 100
   508  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   509  	it := newIterAdapter(l.NewIter(nil, nil))
   510  
   511  	require.False(t, it.Valid())
   512  
   513  	it.Last()
   514  	require.False(t, it.Valid())
   515  
   516  	var ins Inserter
   517  	for i := 0; i < n; i++ {
   518  		ins.Add(l, makeIntKey(i), makeValue(i))
   519  	}
   520  
   521  	it.Last()
   522  	for i := n - 1; i >= 0; i-- {
   523  		require.True(t, it.Valid())
   524  		require.EqualValues(t, makeIntKey(i), it.Key())
   525  		require.EqualValues(t, makeValue(i), it.Value())
   526  		it.Prev()
   527  	}
   528  	require.False(t, it.Valid())
   529  }
   530  
   531  func TestIteratorSeekGEAndSeekPrefixGE(t *testing.T) {
   532  	const n = 100
   533  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   534  	it := newIterAdapter(l.NewIter(nil, nil))
   535  
   536  	require.False(t, it.Valid())
   537  	it.First()
   538  	require.False(t, it.Valid())
   539  	// 1000, 1010, 1020, ..., 1990.
   540  
   541  	var ins Inserter
   542  	for i := n - 1; i >= 0; i-- {
   543  		v := i*10 + 1000
   544  		ins.Add(l, makeIntKey(v), makeValue(v))
   545  	}
   546  
   547  	require.True(t, it.SeekGE(makeKey(""), base.SeekGEFlagsNone))
   548  	require.True(t, it.Valid())
   549  	require.EqualValues(t, "01000", it.Key().UserKey)
   550  	require.EqualValues(t, "v01000", it.Value())
   551  
   552  	require.True(t, it.SeekGE(makeKey("01000"), base.SeekGEFlagsNone))
   553  	require.True(t, it.Valid())
   554  	require.EqualValues(t, "01000", it.Key().UserKey)
   555  	require.EqualValues(t, "v01000", it.Value())
   556  
   557  	require.True(t, it.SeekGE(makeKey("01005"), base.SeekGEFlagsNone))
   558  	require.True(t, it.Valid())
   559  	require.EqualValues(t, "01010", it.Key().UserKey)
   560  	require.EqualValues(t, "v01010", it.Value())
   561  
   562  	require.True(t, it.SeekGE(makeKey("01010"), base.SeekGEFlagsNone))
   563  	require.True(t, it.Valid())
   564  	require.EqualValues(t, "01010", it.Key().UserKey)
   565  	require.EqualValues(t, "v01010", it.Value())
   566  
   567  	require.False(t, it.SeekGE(makeKey("99999"), base.SeekGEFlagsNone))
   568  	require.False(t, it.Valid())
   569  
   570  	// Test SeekGE with trySeekUsingNext optimization.
   571  	{
   572  		require.True(t, it.SeekGE(makeKey("01000"), base.SeekGEFlagsNone))
   573  		require.True(t, it.Valid())
   574  		require.EqualValues(t, "01000", it.Key().UserKey)
   575  		require.EqualValues(t, "v01000", it.Value())
   576  
   577  		// Seeking to the same key.
   578  		require.True(t, it.SeekGE(makeKey("01000"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   579  		require.True(t, it.Valid())
   580  		require.EqualValues(t, "01000", it.Key().UserKey)
   581  		require.EqualValues(t, "v01000", it.Value())
   582  
   583  		// Seeking to a nearby key that can be reached using Next.
   584  		require.True(t, it.SeekGE(makeKey("01020"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   585  		require.True(t, it.Valid())
   586  		require.EqualValues(t, "01020", it.Key().UserKey)
   587  		require.EqualValues(t, "v01020", it.Value())
   588  
   589  		// Seeking to a key that cannot be reached using Next.
   590  		require.True(t, it.SeekGE(makeKey("01200"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   591  		require.True(t, it.Valid())
   592  		require.EqualValues(t, "01200", it.Key().UserKey)
   593  		require.EqualValues(t, "v01200", it.Value())
   594  
   595  		// Seeking to an earlier key, but the caller lies. Incorrect result.
   596  		require.True(t, it.SeekGE(makeKey("01100"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   597  		require.True(t, it.Valid())
   598  		require.EqualValues(t, "01200", it.Key().UserKey)
   599  		require.EqualValues(t, "v01200", it.Value())
   600  
   601  		// Telling the truth works.
   602  		require.True(t, it.SeekGE(makeKey("01100"), base.SeekGEFlagsNone))
   603  		require.True(t, it.Valid())
   604  		require.EqualValues(t, "01100", it.Key().UserKey)
   605  		require.EqualValues(t, "v01100", it.Value())
   606  	}
   607  
   608  	// Test SeekPrefixGE with trySeekUsingNext optimization.
   609  	{
   610  		require.True(t, it.SeekPrefixGE(makeKey("01000"), makeKey("01000"), base.SeekGEFlagsNone))
   611  		require.True(t, it.Valid())
   612  		require.EqualValues(t, "01000", it.Key().UserKey)
   613  		require.EqualValues(t, "v01000", it.Value())
   614  
   615  		// Seeking to the same key.
   616  		require.True(t, it.SeekPrefixGE(makeKey("01000"), makeKey("01000"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   617  		require.True(t, it.Valid())
   618  		require.EqualValues(t, "01000", it.Key().UserKey)
   619  		require.EqualValues(t, "v01000", it.Value())
   620  
   621  		// Seeking to a nearby key that can be reached using Next.
   622  		require.True(t, it.SeekPrefixGE(makeKey("01020"), makeKey("01020"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   623  		require.True(t, it.Valid())
   624  		require.EqualValues(t, "01020", it.Key().UserKey)
   625  		require.EqualValues(t, "v01020", it.Value())
   626  
   627  		// Seeking to a key that cannot be reached using Next.
   628  		require.True(t, it.SeekPrefixGE(makeKey("01200"), makeKey("01200"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   629  		require.True(t, it.Valid())
   630  		require.EqualValues(t, "01200", it.Key().UserKey)
   631  		require.EqualValues(t, "v01200", it.Value())
   632  
   633  		// Seeking to an earlier key, but the caller lies. Incorrect result.
   634  		require.True(t, it.SeekPrefixGE(makeKey("01100"), makeKey("01100"), base.SeekGEFlagsNone.EnableTrySeekUsingNext()))
   635  		require.True(t, it.Valid())
   636  		require.EqualValues(t, "01200", it.Key().UserKey)
   637  		require.EqualValues(t, "v01200", it.Value())
   638  
   639  		// Telling the truth works.
   640  		require.True(t, it.SeekPrefixGE(makeKey("01100"), makeKey("01100"), base.SeekGEFlagsNone))
   641  		require.True(t, it.Valid())
   642  		require.EqualValues(t, "01100", it.Key().UserKey)
   643  		require.EqualValues(t, "v01100", it.Value())
   644  	}
   645  
   646  	// Test seek for empty key.
   647  	ins.Add(l, base.InternalKey{}, nil)
   648  	require.True(t, it.SeekGE([]byte{}, base.SeekGEFlagsNone))
   649  	require.True(t, it.Valid())
   650  	require.EqualValues(t, "", it.Key().UserKey)
   651  
   652  	require.True(t, it.SeekGE(makeKey(""), base.SeekGEFlagsNone))
   653  	require.True(t, it.Valid())
   654  	require.EqualValues(t, "", it.Key().UserKey)
   655  }
   656  
   657  func TestIteratorSeekLT(t *testing.T) {
   658  	const n = 100
   659  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   660  	it := newIterAdapter(l.NewIter(nil, nil))
   661  
   662  	require.False(t, it.Valid())
   663  	it.First()
   664  	require.False(t, it.Valid())
   665  	// 1000, 1010, 1020, ..., 1990.
   666  	var ins Inserter
   667  	for i := n - 1; i >= 0; i-- {
   668  		v := i*10 + 1000
   669  		ins.Add(l, makeIntKey(v), makeValue(v))
   670  	}
   671  
   672  	require.False(t, it.SeekLT(makeKey(""), base.SeekLTFlagsNone))
   673  	require.False(t, it.Valid())
   674  
   675  	require.False(t, it.SeekLT(makeKey("01000"), base.SeekLTFlagsNone))
   676  	require.False(t, it.Valid())
   677  
   678  	require.True(t, it.SeekLT(makeKey("01001"), base.SeekLTFlagsNone))
   679  	require.True(t, it.Valid())
   680  	require.EqualValues(t, "01000", it.Key().UserKey)
   681  	require.EqualValues(t, "v01000", it.Value())
   682  
   683  	require.True(t, it.SeekLT(makeKey("01005"), base.SeekLTFlagsNone))
   684  	require.True(t, it.Valid())
   685  	require.EqualValues(t, "01000", it.Key().UserKey)
   686  	require.EqualValues(t, "v01000", it.Value())
   687  
   688  	require.True(t, it.SeekLT(makeKey("01991"), base.SeekLTFlagsNone))
   689  	require.True(t, it.Valid())
   690  	require.EqualValues(t, "01990", it.Key().UserKey)
   691  	require.EqualValues(t, "v01990", it.Value())
   692  
   693  	require.True(t, it.SeekLT(makeKey("99999"), base.SeekLTFlagsNone))
   694  	require.True(t, it.Valid())
   695  	require.EqualValues(t, "01990", it.Key().UserKey)
   696  	require.EqualValues(t, "v01990", it.Value())
   697  
   698  	// Test seek for empty key.
   699  	ins.Add(l, base.InternalKey{}, nil)
   700  	require.False(t, it.SeekLT([]byte{}, base.SeekLTFlagsNone))
   701  	require.False(t, it.Valid())
   702  
   703  	require.True(t, it.SeekLT(makeKey("\x01"), base.SeekLTFlagsNone))
   704  	require.True(t, it.Valid())
   705  	require.EqualValues(t, "", it.Key().UserKey)
   706  }
   707  
   708  // TODO(peter): test First and Last.
   709  func TestIteratorBounds(t *testing.T) {
   710  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   711  	for i := 1; i < 10; i++ {
   712  		require.NoError(t, l.Add(makeIntKey(i), makeValue(i)))
   713  	}
   714  
   715  	key := func(i int) []byte {
   716  		return makeIntKey(i).UserKey
   717  	}
   718  
   719  	it := newIterAdapter(l.NewIter(key(3), key(7)))
   720  
   721  	// SeekGE within the lower and upper bound succeeds.
   722  	for i := 3; i <= 6; i++ {
   723  		k := key(i)
   724  		require.True(t, it.SeekGE(k, base.SeekGEFlagsNone))
   725  		require.EqualValues(t, string(k), string(it.Key().UserKey))
   726  	}
   727  
   728  	// SeekGE before the lower bound still succeeds (only the upper bound is
   729  	// checked).
   730  	for i := 1; i < 3; i++ {
   731  		k := key(i)
   732  		require.True(t, it.SeekGE(k, base.SeekGEFlagsNone))
   733  		require.EqualValues(t, string(k), string(it.Key().UserKey))
   734  	}
   735  
   736  	// SeekGE beyond the upper bound fails.
   737  	for i := 7; i < 10; i++ {
   738  		require.False(t, it.SeekGE(key(i), base.SeekGEFlagsNone))
   739  	}
   740  
   741  	require.True(t, it.SeekGE(key(6), base.SeekGEFlagsNone))
   742  	require.EqualValues(t, "00006", it.Key().UserKey)
   743  	require.EqualValues(t, "v00006", it.Value())
   744  
   745  	// Next into the upper bound fails.
   746  	require.False(t, it.Next())
   747  
   748  	// SeekLT within the lower and upper bound succeeds.
   749  	for i := 4; i <= 7; i++ {
   750  		require.True(t, it.SeekLT(key(i), base.SeekLTFlagsNone))
   751  		require.EqualValues(t, string(key(i-1)), string(it.Key().UserKey))
   752  	}
   753  
   754  	// SeekLT beyond the upper bound still succeeds (only the lower bound is
   755  	// checked).
   756  	for i := 8; i < 9; i++ {
   757  		require.True(t, it.SeekLT(key(8), base.SeekLTFlagsNone))
   758  		require.EqualValues(t, string(key(i-1)), string(it.Key().UserKey))
   759  	}
   760  
   761  	// SeekLT before the lower bound fails.
   762  	for i := 1; i < 4; i++ {
   763  		require.False(t, it.SeekLT(key(i), base.SeekLTFlagsNone))
   764  	}
   765  
   766  	require.True(t, it.SeekLT(key(4), base.SeekLTFlagsNone))
   767  	require.EqualValues(t, "00003", it.Key().UserKey)
   768  	require.EqualValues(t, "v00003", it.Value())
   769  
   770  	// Prev into the lower bound fails.
   771  	require.False(t, it.Prev())
   772  }
   773  
   774  func TestBytesIterated(t *testing.T) {
   775  	l := NewSkiplist(newArena(arenaSize), bytes.Compare)
   776  	emptySize := l.arena.Size()
   777  	for i := 0; i < 200; i++ {
   778  		bytesIterated := l.bytesIterated(t)
   779  		expected := uint64(l.arena.Size() - emptySize)
   780  		if bytesIterated != expected {
   781  			t.Fatalf("bytesIterated: got %d, want %d", bytesIterated, expected)
   782  		}
   783  		l.Add(base.InternalKey{UserKey: []byte{byte(i)}}, nil)
   784  	}
   785  }
   786  
   787  // bytesIterated returns the number of bytes iterated in the skiplist.
   788  func (s *Skiplist) bytesIterated(t *testing.T) (bytesIterated uint64) {
   789  	x := s.NewFlushIter(&bytesIterated)
   790  	var prevIterated uint64
   791  	for key, _ := x.First(); key != nil; key, _ = x.Next() {
   792  		if bytesIterated < prevIterated {
   793  			t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated)
   794  		}
   795  		prevIterated = bytesIterated
   796  	}
   797  	if x.Close() != nil {
   798  		return 0
   799  	}
   800  	return bytesIterated
   801  }
   802  
   803  func randomKey(rng *rand.Rand, b []byte) base.InternalKey {
   804  	key := rng.Uint32()
   805  	key2 := rng.Uint32()
   806  	binary.LittleEndian.PutUint32(b, key)
   807  	binary.LittleEndian.PutUint32(b[4:], key2)
   808  	return base.InternalKey{UserKey: b}
   809  }
   810  
   811  // Standard test. Some fraction is read. Some fraction is write. Writes have
   812  // to go through mutex lock.
   813  func BenchmarkReadWrite(b *testing.B) {
   814  	for i := 0; i <= 10; i++ {
   815  		readFrac := float32(i) / 10.0
   816  		b.Run(fmt.Sprintf("frac_%d", i*10), func(b *testing.B) {
   817  			l := NewSkiplist(newArena(uint32((b.N+2)*maxNodeSize)), bytes.Compare)
   818  			b.ResetTimer()
   819  			var count int
   820  			b.RunParallel(func(pb *testing.PB) {
   821  				it := l.NewIter(nil, nil)
   822  				rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
   823  				buf := make([]byte, 8)
   824  
   825  				for pb.Next() {
   826  					if rng.Float32() < readFrac {
   827  						key, _ := it.SeekGE(randomKey(rng, buf).UserKey, base.SeekGEFlagsNone)
   828  						if key != nil {
   829  							_ = key
   830  							count++
   831  						}
   832  					} else {
   833  						_ = l.Add(randomKey(rng, buf), nil)
   834  					}
   835  				}
   836  			})
   837  		})
   838  	}
   839  }
   840  
   841  func BenchmarkOrderedWrite(b *testing.B) {
   842  	l := NewSkiplist(newArena(8<<20), bytes.Compare)
   843  	var ins Inserter
   844  	buf := make([]byte, 8)
   845  
   846  	b.ResetTimer()
   847  	for i := 0; i < b.N; i++ {
   848  		binary.BigEndian.PutUint64(buf, uint64(i))
   849  		if err := ins.Add(l, base.InternalKey{UserKey: buf}, nil); err == ErrArenaFull {
   850  			b.StopTimer()
   851  			l = NewSkiplist(newArena(uint32((b.N+2)*maxNodeSize)), bytes.Compare)
   852  			ins = Inserter{}
   853  			b.StartTimer()
   854  		}
   855  	}
   856  }
   857  
   858  func BenchmarkIterNext(b *testing.B) {
   859  	l := NewSkiplist(newArena(64<<10), bytes.Compare)
   860  	rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
   861  	buf := make([]byte, 8)
   862  	for {
   863  		if err := l.Add(randomKey(rng, buf), nil); err == ErrArenaFull {
   864  			break
   865  		}
   866  	}
   867  
   868  	it := l.NewIter(nil, nil)
   869  	b.ResetTimer()
   870  	for i := 0; i < b.N; i++ {
   871  		key, _ := it.Next()
   872  		if key == nil {
   873  			key, _ = it.First()
   874  		}
   875  		_ = key
   876  	}
   877  }
   878  
   879  func BenchmarkIterPrev(b *testing.B) {
   880  	l := NewSkiplist(newArena(64<<10), bytes.Compare)
   881  	rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
   882  	buf := make([]byte, 8)
   883  	for {
   884  		if err := l.Add(randomKey(rng, buf), nil); err == ErrArenaFull {
   885  			break
   886  		}
   887  	}
   888  
   889  	it := l.NewIter(nil, nil)
   890  	_, _ = it.Last()
   891  	b.ResetTimer()
   892  	for i := 0; i < b.N; i++ {
   893  		key, _ := it.Prev()
   894  		if key == nil {
   895  			key, _ = it.Last()
   896  		}
   897  		_ = key
   898  	}
   899  }
   900  
   901  // BenchmarkSeekPrefixGE looks at the performance of repeated calls to
   902  // SeekPrefixGE, with different skip distances and different settings of
   903  // trySeekUsingNext.
   904  func BenchmarkSeekPrefixGE(b *testing.B) {
   905  	l := NewSkiplist(newArena(64<<10), bytes.Compare)
   906  	var count int
   907  	// count was measured to be 1279.
   908  	for count = 0; ; count++ {
   909  		if err := l.Add(makeIntKey(count), makeValue(count)); err == ErrArenaFull {
   910  			break
   911  		}
   912  	}
   913  	for _, skip := range []int{1, 2, 4, 8, 16} {
   914  		for _, useNext := range []bool{false, true} {
   915  			b.Run(fmt.Sprintf("skip=%d/use-next=%t", skip, useNext), func(b *testing.B) {
   916  				it := l.NewIter(nil, nil)
   917  				j := 0
   918  				var k []byte
   919  				makeKey := func() {
   920  					k = []byte(fmt.Sprintf("%05d", j))
   921  				}
   922  				makeKey()
   923  				it.SeekPrefixGE(k, k, base.SeekGEFlagsNone)
   924  				b.ResetTimer()
   925  				for i := 0; i < b.N; i++ {
   926  					j += skip
   927  					var flags base.SeekGEFlags
   928  					if useNext {
   929  						flags = flags.EnableTrySeekUsingNext()
   930  					}
   931  					if j >= count {
   932  						j = 0
   933  						flags = flags.DisableTrySeekUsingNext()
   934  					}
   935  					makeKey()
   936  					it.SeekPrefixGE(k, k, flags)
   937  				}
   938  			})
   939  		}
   940  	}
   941  }
   942  
   943  // Standard test. Some fraction is read. Some fraction is write. Writes have
   944  // to go through mutex lock.
   945  // func BenchmarkReadWriteMap(b *testing.B) {
   946  // 	for i := 0; i <= 10; i++ {
   947  // 		readFrac := float32(i) / 10.0
   948  // 		b.Run(fmt.Sprintf("frac_%d", i*10), func(b *testing.B) {
   949  // 			m := make(map[string]struct{})
   950  // 			var mutex sync.RWMutex
   951  // 			b.ResetTimer()
   952  // 			var count int
   953  // 			b.RunParallel(func(pb *testing.PB) {
   954  // 				rng := rand.New(rand.NewSource(time.Now().UnixNano()))
   955  // 				for pb.Next() {
   956  // 					if rng.Float32() < readFrac {
   957  // 						mutex.RLock()
   958  // 						_, ok := m[string(randomKey(rng))]
   959  // 						mutex.RUnlock()
   960  // 						if ok {
   961  // 							count++
   962  // 						}
   963  // 					} else {
   964  // 						mutex.Lock()
   965  // 						m[string(randomKey(rng))] = struct{}{}
   966  // 						mutex.Unlock()
   967  // 					}
   968  // 				}
   969  // 			})
   970  // 		})
   971  // 	}
   972  // }