github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/mem_table_test.go (about)

     1  // Copyright 2011 The LevelDB-Go and Pebble and Bitalostored 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 bitalostable
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"fmt"
    11  	"strconv"
    12  	"strings"
    13  	"sync/atomic"
    14  	"testing"
    15  	"time"
    16  
    17  	"github.com/cockroachdb/errors"
    18  	"github.com/stretchr/testify/require"
    19  	"github.com/zuoyebang/bitalostable/internal/arenaskl"
    20  	"github.com/zuoyebang/bitalostable/internal/base"
    21  	"github.com/zuoyebang/bitalostable/internal/datadriven"
    22  	"github.com/zuoyebang/bitalostable/internal/rangekey"
    23  	"golang.org/x/exp/rand"
    24  	"golang.org/x/sync/errgroup"
    25  )
    26  
    27  // get gets the value for the given key. It returns ErrNotFound if the DB does
    28  // not contain the key.
    29  func (m *memTable) get(key []byte) (value []byte, err error) {
    30  	it := m.skl.NewIter(nil, nil)
    31  	ikey, val := it.SeekGE(key, base.SeekGEFlagsNone)
    32  	if ikey == nil {
    33  		return nil, ErrNotFound
    34  	}
    35  	if !m.equal(key, ikey.UserKey) {
    36  		return nil, ErrNotFound
    37  	}
    38  	switch ikey.Kind() {
    39  	case InternalKeyKindDelete, InternalKeyKindSingleDelete:
    40  		return nil, ErrNotFound
    41  	default:
    42  		return val, nil
    43  	}
    44  }
    45  
    46  // Set sets the value for the given key. It overwrites any previous value for
    47  // that key; a DB is not a multi-map. NB: this might have unexpected
    48  // interaction with prepare/apply. Caveat emptor!
    49  func (m *memTable) set(key InternalKey, value []byte) error {
    50  	if key.Kind() == InternalKeyKindRangeDelete {
    51  		if err := m.rangeDelSkl.Add(key, value); err != nil {
    52  			return err
    53  		}
    54  		m.tombstones.invalidate(1)
    55  		return nil
    56  	}
    57  	if rangekey.IsRangeKey(key.Kind()) {
    58  		if err := m.rangeKeySkl.Add(key, value); err != nil {
    59  			return err
    60  		}
    61  		m.rangeKeys.invalidate(1)
    62  		return nil
    63  	}
    64  	return m.skl.Add(key, value)
    65  }
    66  
    67  // count returns the number of entries in a DB.
    68  func (m *memTable) count() (n int) {
    69  	x := newInternalIterAdapter(m.newIter(nil))
    70  	for valid := x.First(); valid; valid = x.Next() {
    71  		n++
    72  	}
    73  	if x.Close() != nil {
    74  		return -1
    75  	}
    76  	return n
    77  }
    78  
    79  // bytesIterated returns the number of bytes iterated in a DB.
    80  func (m *memTable) bytesIterated(t *testing.T) (bytesIterated uint64) {
    81  	x := newInternalIterAdapter(m.newFlushIter(nil, &bytesIterated))
    82  	var prevIterated uint64
    83  	for valid := x.First(); valid; valid = x.Next() {
    84  		if bytesIterated < prevIterated {
    85  			t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated)
    86  		}
    87  		prevIterated = bytesIterated
    88  	}
    89  	if x.Close() != nil {
    90  		return 0
    91  	}
    92  	return bytesIterated
    93  }
    94  
    95  func ikey(s string) InternalKey {
    96  	return base.MakeInternalKey([]byte(s), 0, InternalKeyKindSet)
    97  }
    98  
    99  func TestMemTableBasic(t *testing.T) {
   100  	// Check the empty DB.
   101  	m := newMemTable(memTableOptions{})
   102  	if got, want := m.count(), 0; got != want {
   103  		t.Fatalf("0.count: got %v, want %v", got, want)
   104  	}
   105  	v, err := m.get([]byte("cherry"))
   106  	if string(v) != "" || err != ErrNotFound {
   107  		t.Fatalf("1.get: got (%q, %v), want (%q, %v)", v, err, "", ErrNotFound)
   108  	}
   109  	// Add some key/value pairs.
   110  	m.set(ikey("cherry"), []byte("red"))
   111  	m.set(ikey("peach"), []byte("yellow"))
   112  	m.set(ikey("grape"), []byte("red"))
   113  	m.set(ikey("grape"), []byte("green"))
   114  	m.set(ikey("plum"), []byte("purple"))
   115  	if got, want := m.count(), 4; got != want {
   116  		t.Fatalf("2.count: got %v, want %v", got, want)
   117  	}
   118  	// Get keys that are and aren't in the DB.
   119  	v, err = m.get([]byte("plum"))
   120  	if string(v) != "purple" || err != nil {
   121  		t.Fatalf("6.get: got (%q, %v), want (%q, %v)", v, err, "purple", error(nil))
   122  	}
   123  	v, err = m.get([]byte("lychee"))
   124  	if string(v) != "" || err != ErrNotFound {
   125  		t.Fatalf("7.get: got (%q, %v), want (%q, %v)", v, err, "", ErrNotFound)
   126  	}
   127  	// Check an iterator.
   128  	s, x := "", newInternalIterAdapter(m.newIter(nil))
   129  	for valid := x.SeekGE([]byte("mango"), base.SeekGEFlagsNone); valid; valid = x.Next() {
   130  		s += fmt.Sprintf("%s/%s.", x.Key().UserKey, x.Value())
   131  	}
   132  	if want := "peach/yellow.plum/purple."; s != want {
   133  		t.Fatalf("8.iter: got %q, want %q", s, want)
   134  	}
   135  	if err = x.Close(); err != nil {
   136  		t.Fatalf("9.close: %v", err)
   137  	}
   138  	// Check some more sets and deletes.
   139  	if err := m.set(ikey("apricot"), []byte("orange")); err != nil {
   140  		t.Fatalf("12.set: %v", err)
   141  	}
   142  	if got, want := m.count(), 5; got != want {
   143  		t.Fatalf("13.count: got %v, want %v", got, want)
   144  	}
   145  }
   146  
   147  func TestMemTableCount(t *testing.T) {
   148  	m := newMemTable(memTableOptions{})
   149  	for i := 0; i < 200; i++ {
   150  		if j := m.count(); j != i {
   151  			t.Fatalf("count: got %d, want %d", j, i)
   152  		}
   153  		m.set(InternalKey{UserKey: []byte{byte(i)}}, nil)
   154  	}
   155  }
   156  
   157  func TestMemTableBytesIterated(t *testing.T) {
   158  	m := newMemTable(memTableOptions{})
   159  	for i := 0; i < 200; i++ {
   160  		bytesIterated := m.bytesIterated(t)
   161  		expected := m.inuseBytes()
   162  		if bytesIterated != expected {
   163  			t.Fatalf("bytesIterated: got %d, want %d", bytesIterated, expected)
   164  		}
   165  		m.set(InternalKey{UserKey: []byte{byte(i)}}, nil)
   166  	}
   167  }
   168  
   169  func TestMemTableEmpty(t *testing.T) {
   170  	m := newMemTable(memTableOptions{})
   171  	if !m.empty() {
   172  		t.Errorf("got !empty, want empty")
   173  	}
   174  	// Add one key/value pair with an empty key and empty value.
   175  	m.set(InternalKey{}, nil)
   176  	if m.empty() {
   177  		t.Errorf("got empty, want !empty")
   178  	}
   179  }
   180  
   181  func TestMemTable1000Entries(t *testing.T) {
   182  	// Initialize the DB.
   183  	const N = 1000
   184  	m0 := newMemTable(memTableOptions{})
   185  	for i := 0; i < N; i++ {
   186  		k := ikey(strconv.Itoa(i))
   187  		v := []byte(strings.Repeat("x", i))
   188  		m0.set(k, v)
   189  	}
   190  	// Check the DB count.
   191  	if got, want := m0.count(), 1000; got != want {
   192  		t.Fatalf("count: got %v, want %v", got, want)
   193  	}
   194  	// Check random-access lookup.
   195  	r := rand.New(rand.NewSource(0))
   196  	for i := 0; i < 3*N; i++ {
   197  		j := r.Intn(N)
   198  		k := []byte(strconv.Itoa(j))
   199  		v, err := m0.get(k)
   200  		require.NoError(t, err)
   201  		if len(v) != cap(v) {
   202  			t.Fatalf("get: j=%d, got len(v)=%d, cap(v)=%d", j, len(v), cap(v))
   203  		}
   204  		var c uint8
   205  		if len(v) != 0 {
   206  			c = v[0]
   207  		} else {
   208  			c = 'x'
   209  		}
   210  		if len(v) != j || c != 'x' {
   211  			t.Fatalf("get: j=%d, got len(v)=%d,c=%c, want %d,%c", j, len(v), c, j, 'x')
   212  		}
   213  	}
   214  	// Check that iterating through the middle of the DB looks OK.
   215  	// Keys are in lexicographic order, not numerical order.
   216  	// Multiples of 3 are not present.
   217  	wants := []string{
   218  		"499",
   219  		"5",
   220  		"50",
   221  		"500",
   222  		"501",
   223  		"502",
   224  		"503",
   225  		"504",
   226  		"505",
   227  		"506",
   228  		"507",
   229  	}
   230  	x := newInternalIterAdapter(m0.newIter(nil))
   231  	x.SeekGE([]byte(wants[0]), base.SeekGEFlagsNone)
   232  	for _, want := range wants {
   233  		if !x.Valid() {
   234  			t.Fatalf("iter: next failed, want=%q", want)
   235  		}
   236  		if got := string(x.Key().UserKey); got != want {
   237  			t.Fatalf("iter: got %q, want %q", got, want)
   238  		}
   239  		if k := x.Key().UserKey; len(k) != cap(k) {
   240  			t.Fatalf("iter: len(k)=%d, cap(k)=%d", len(k), cap(k))
   241  		}
   242  		if v := x.Value(); len(v) != cap(v) {
   243  			t.Fatalf("iter: len(v)=%d, cap(v)=%d", len(v), cap(v))
   244  		}
   245  		x.Next()
   246  	}
   247  	if err := x.Close(); err != nil {
   248  		t.Fatalf("close: %v", err)
   249  	}
   250  }
   251  
   252  func TestMemTableIter(t *testing.T) {
   253  	var mem *memTable
   254  	for _, testdata := range []string{
   255  		"testdata/internal_iter_next", "testdata/internal_iter_bounds"} {
   256  		datadriven.RunTest(t, testdata, func(d *datadriven.TestData) string {
   257  			switch d.Cmd {
   258  			case "define":
   259  				mem = newMemTable(memTableOptions{})
   260  				for _, key := range strings.Split(d.Input, "\n") {
   261  					j := strings.Index(key, ":")
   262  					if err := mem.set(base.ParseInternalKey(key[:j]), []byte(key[j+1:])); err != nil {
   263  						return err.Error()
   264  					}
   265  				}
   266  				return ""
   267  
   268  			case "iter":
   269  				var options IterOptions
   270  				for _, arg := range d.CmdArgs {
   271  					switch arg.Key {
   272  					case "lower":
   273  						if len(arg.Vals) != 1 {
   274  							return fmt.Sprintf(
   275  								"%s expects at most 1 value for lower", d.Cmd)
   276  						}
   277  						options.LowerBound = []byte(arg.Vals[0])
   278  					case "upper":
   279  						if len(arg.Vals) != 1 {
   280  							return fmt.Sprintf(
   281  								"%s expects at most 1 value for upper", d.Cmd)
   282  						}
   283  						options.UpperBound = []byte(arg.Vals[0])
   284  					default:
   285  						return fmt.Sprintf("unknown arg: %s", arg.Key)
   286  					}
   287  				}
   288  				iter := mem.newIter(&options)
   289  				defer iter.Close()
   290  				return runInternalIterCmd(d, iter)
   291  
   292  			default:
   293  				return fmt.Sprintf("unknown command: %s", d.Cmd)
   294  			}
   295  		})
   296  	}
   297  }
   298  
   299  func TestMemTableDeleteRange(t *testing.T) {
   300  	var mem *memTable
   301  	var seqNum uint64
   302  
   303  	datadriven.RunTest(t, "testdata/delete_range", func(td *datadriven.TestData) string {
   304  		switch td.Cmd {
   305  		case "clear":
   306  			mem = nil
   307  			seqNum = 0
   308  			return ""
   309  
   310  		case "define":
   311  			b := newBatch(nil)
   312  			if err := runBatchDefineCmd(td, b); err != nil {
   313  				return err.Error()
   314  			}
   315  			if mem == nil {
   316  				mem = newMemTable(memTableOptions{})
   317  			}
   318  			if err := mem.apply(b, seqNum); err != nil {
   319  				return err.Error()
   320  			}
   321  			seqNum += uint64(b.Count())
   322  			return ""
   323  
   324  		case "scan":
   325  			if len(td.CmdArgs) > 1 {
   326  				return fmt.Sprintf("%s expects at most 1 argument", td.Cmd)
   327  			}
   328  			var buf bytes.Buffer
   329  			if len(td.CmdArgs) == 1 {
   330  				if td.CmdArgs[0].String() != "range-del" {
   331  					return fmt.Sprintf("%s unknown argument %s", td.Cmd, td.CmdArgs[0])
   332  				}
   333  				iter := mem.newRangeDelIter(nil)
   334  				defer iter.Close()
   335  				scanKeyspanIterator(&buf, iter)
   336  			} else {
   337  				iter := mem.newIter(nil)
   338  				defer iter.Close()
   339  				scanInternalIterator(&buf, iter)
   340  			}
   341  			return buf.String()
   342  
   343  		default:
   344  			return fmt.Sprintf("unknown command: %s", td.Cmd)
   345  		}
   346  	})
   347  }
   348  
   349  func TestMemTableConcurrentDeleteRange(t *testing.T) {
   350  	// Concurrently write and read range tombstones. Workers add range
   351  	// tombstones, and then immediately retrieve them verifying that the
   352  	// tombstones they've added are all present.
   353  
   354  	m := newMemTable(memTableOptions{Options: &Options{MemTableSize: 64 << 20}})
   355  
   356  	const workers = 10
   357  	eg, _ := errgroup.WithContext(context.Background())
   358  	seqNum := uint64(1)
   359  	for i := 0; i < workers; i++ {
   360  		i := i
   361  		eg.Go(func() error {
   362  			start := ([]byte)(fmt.Sprintf("%03d", i))
   363  			end := ([]byte)(fmt.Sprintf("%03d", i+1))
   364  			for j := 0; j < 100; j++ {
   365  				b := newBatch(nil)
   366  				b.DeleteRange(start, end, nil)
   367  				n := atomic.AddUint64(&seqNum, 1) - 1
   368  				require.NoError(t, m.apply(b, n))
   369  				b.release()
   370  
   371  				var count int
   372  				it := m.newRangeDelIter(nil)
   373  				for s := it.SeekGE(start); s != nil; s = it.Next() {
   374  					if m.cmp(s.Start, end) >= 0 {
   375  						break
   376  					}
   377  					count += len(s.Keys)
   378  				}
   379  				if j+1 != count {
   380  					return errors.Errorf("%d: expected %d tombstones, but found %d", i, j+1, count)
   381  				}
   382  			}
   383  			return nil
   384  		})
   385  	}
   386  	err := eg.Wait()
   387  	if err != nil {
   388  		t.Error(err)
   389  	}
   390  }
   391  
   392  func buildMemTable(b *testing.B) (*memTable, [][]byte) {
   393  	m := newMemTable(memTableOptions{})
   394  	var keys [][]byte
   395  	var ikey InternalKey
   396  	for i := 0; ; i++ {
   397  		key := []byte(fmt.Sprintf("%08d", i))
   398  		keys = append(keys, key)
   399  		ikey = base.MakeInternalKey(key, 0, InternalKeyKindSet)
   400  		if m.set(ikey, nil) == arenaskl.ErrArenaFull {
   401  			break
   402  		}
   403  	}
   404  	return m, keys
   405  }
   406  
   407  func BenchmarkMemTableIterSeekGE(b *testing.B) {
   408  	m, keys := buildMemTable(b)
   409  	iter := m.newIter(nil)
   410  	rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
   411  
   412  	b.ResetTimer()
   413  	for i := 0; i < b.N; i++ {
   414  		iter.SeekGE(keys[rng.Intn(len(keys))], base.SeekGEFlagsNone)
   415  	}
   416  }
   417  
   418  func BenchmarkMemTableIterNext(b *testing.B) {
   419  	m, _ := buildMemTable(b)
   420  	iter := m.newIter(nil)
   421  	_, _ = iter.First()
   422  	b.ResetTimer()
   423  	for i := 0; i < b.N; i++ {
   424  		key, _ := iter.Next()
   425  		if key == nil {
   426  			key, _ = iter.First()
   427  		}
   428  		_ = key
   429  	}
   430  }
   431  
   432  func BenchmarkMemTableIterPrev(b *testing.B) {
   433  	m, _ := buildMemTable(b)
   434  	iter := m.newIter(nil)
   435  	_, _ = iter.Last()
   436  	b.ResetTimer()
   437  	for i := 0; i < b.N; i++ {
   438  		key, _ := iter.Prev()
   439  		if key == nil {
   440  			key, _ = iter.Last()
   441  		}
   442  		_ = key
   443  	}
   444  }