github.com/zuoyebang/bitalosdb@v1.1.1-0.20240516111551-79a8c4d8ce20/mem_table_test.go (about)

     1  // Copyright 2021 The Bitalosdb author(hustxrb@163.com) and other contributors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package bitalosdb
    16  
    17  import (
    18  	"fmt"
    19  	"strconv"
    20  	"strings"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/stretchr/testify/require"
    25  	"github.com/zuoyebang/bitalosdb/internal/arenaskl"
    26  	"github.com/zuoyebang/bitalosdb/internal/base"
    27  	"golang.org/x/exp/rand"
    28  )
    29  
    30  func (m *memTable) count() (n int) {
    31  	x := newInternalIterAdapter(m.newIter(nil))
    32  	for valid := x.First(); valid; valid = x.Next() {
    33  		n++
    34  	}
    35  	if x.Close() != nil {
    36  		return -1
    37  	}
    38  	return n
    39  }
    40  
    41  func (m *memTable) bytesIterated(t *testing.T) (bytesIterated uint64) {
    42  	x := newInternalIterAdapter(m.newFlushIter(nil, &bytesIterated))
    43  	var prevIterated uint64
    44  	for valid := x.First(); valid; valid = x.Next() {
    45  		if bytesIterated < prevIterated {
    46  			t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated)
    47  		}
    48  		prevIterated = bytesIterated
    49  	}
    50  	if x.Close() != nil {
    51  		return 0
    52  	}
    53  	return bytesIterated
    54  }
    55  
    56  func ikey(s string) InternalKey {
    57  	return base.MakeInternalKey([]byte(s), 0, InternalKeyKindSet)
    58  }
    59  
    60  func testNewMemTable() *memTable {
    61  	opts := &Options{}
    62  	opts = opts.Clone().EnsureDefaults()
    63  	return newMemTable(memTableOptions{Options: opts})
    64  }
    65  
    66  func TestMemTableBasic(t *testing.T) {
    67  	m := testNewMemTable()
    68  	if got, want := m.count(), 0; got != want {
    69  		t.Fatalf("0.count: got %v, want %v", got, want)
    70  	}
    71  
    72  	getKey := func(key []byte) ([]byte, error) {
    73  		v, exist, kind := m.get(key)
    74  		if exist && kind == base.InternalKeyKindSet {
    75  			return v, nil
    76  		}
    77  		return nil, ErrNotFound
    78  	}
    79  
    80  	v, err := getKey([]byte("cherry"))
    81  	if string(v) != "" || err != ErrNotFound {
    82  		t.Fatalf("1.get: got (%q, %v), want (%q, %v)", v, err, "", ErrNotFound)
    83  	}
    84  	m.set(ikey("cherry"), []byte("red"))
    85  	m.set(ikey("peach"), []byte("yellow"))
    86  	m.set(ikey("grape"), []byte("red"))
    87  	m.set(ikey("grape"), []byte("green"))
    88  	m.set(ikey("plum"), []byte("purple"))
    89  	if got, want := m.count(), 4; got != want {
    90  		t.Fatalf("2.count: got %v, want %v", got, want)
    91  	}
    92  	v, err = getKey([]byte("plum"))
    93  	if string(v) != "purple" || err != nil {
    94  		t.Fatalf("6.get: got (%q, %v), want (%q, %v)", v, err, "purple", error(nil))
    95  	}
    96  	v, err = getKey([]byte("lychee"))
    97  	if string(v) != "" || err != ErrNotFound {
    98  		t.Fatalf("7.get: got (%q, %v), want (%q, %v)", v, err, "", ErrNotFound)
    99  	}
   100  	s, x := "", newInternalIterAdapter(m.newIter(nil))
   101  	for valid := x.SeekGE([]byte("mango")); valid; valid = x.Next() {
   102  		s += fmt.Sprintf("%s/%s.", x.Key().UserKey, x.Value())
   103  	}
   104  	if want := "peach/yellow.plum/purple."; s != want {
   105  		t.Fatalf("8.iter: got %q, want %q", s, want)
   106  	}
   107  	if err = x.Close(); err != nil {
   108  		t.Fatalf("9.close: %v", err)
   109  	}
   110  	if err := m.set(ikey("apricot"), []byte("orange")); err != nil {
   111  		t.Fatalf("12.set: %v", err)
   112  	}
   113  	if got, want := m.count(), 5; got != want {
   114  		t.Fatalf("13.count: got %v, want %v", got, want)
   115  	}
   116  	if err := m.close(); err != nil {
   117  		t.Fatalf("14.close: %v", err)
   118  	}
   119  }
   120  
   121  func TestMemTableCount(t *testing.T) {
   122  	m := testNewMemTable()
   123  	for i := 0; i < 200; i++ {
   124  		if j := m.count(); j != i {
   125  			t.Fatalf("count: got %d, want %d", j, i)
   126  		}
   127  		m.set(InternalKey{UserKey: []byte{byte(i)}}, nil)
   128  	}
   129  	require.NoError(t, m.close())
   130  }
   131  
   132  func TestMemTableBytesIterated(t *testing.T) {
   133  	m := testNewMemTable()
   134  	for i := 0; i < 200; i++ {
   135  		bytesIterated := m.bytesIterated(t)
   136  		expected := m.inuseBytes()
   137  		if bytesIterated != expected {
   138  			t.Fatalf("bytesIterated: got %d, want %d", bytesIterated, expected)
   139  		}
   140  		m.set(InternalKey{UserKey: []byte{byte(i)}}, nil)
   141  	}
   142  	require.NoError(t, m.close())
   143  }
   144  
   145  func TestMemTableEmpty(t *testing.T) {
   146  	m := testNewMemTable()
   147  	if !m.empty() {
   148  		t.Errorf("got !empty, want empty")
   149  	}
   150  	m.set(InternalKey{}, nil)
   151  	if m.empty() {
   152  		t.Errorf("got empty, want !empty")
   153  	}
   154  }
   155  
   156  func TestMemTable1000Entries(t *testing.T) {
   157  	const N = 1000
   158  	m0 := testNewMemTable()
   159  	for i := 0; i < N; i++ {
   160  		k := ikey(strconv.Itoa(i))
   161  		v := []byte(strings.Repeat("x", i))
   162  		m0.set(k, v)
   163  	}
   164  	if got, want := m0.count(), 1000; got != want {
   165  		t.Fatalf("count: got %v, want %v", got, want)
   166  	}
   167  	r := rand.New(rand.NewSource(0))
   168  	for i := 0; i < 3*N; i++ {
   169  		j := r.Intn(N)
   170  		k := []byte(strconv.Itoa(j))
   171  		v, vexist, vkind := m0.get(k)
   172  		require.Equal(t, true, vexist)
   173  		require.Equal(t, base.InternalKeyKindSet, vkind)
   174  		if len(v) != cap(v) {
   175  			t.Fatalf("get: j=%d, got len(v)=%d, cap(v)=%d", j, len(v), cap(v))
   176  		}
   177  		var c uint8
   178  		if len(v) != 0 {
   179  			c = v[0]
   180  		} else {
   181  			c = 'x'
   182  		}
   183  		if len(v) != j || c != 'x' {
   184  			t.Fatalf("get: j=%d, got len(v)=%d,c=%c, want %d,%c", j, len(v), c, j, 'x')
   185  		}
   186  	}
   187  	wants := []string{
   188  		"499",
   189  		"5",
   190  		"50",
   191  		"500",
   192  		"501",
   193  		"502",
   194  		"503",
   195  		"504",
   196  		"505",
   197  		"506",
   198  		"507",
   199  	}
   200  	x := newInternalIterAdapter(m0.newIter(nil))
   201  	x.SeekGE([]byte(wants[0]))
   202  	for _, want := range wants {
   203  		if !x.Valid() {
   204  			t.Fatalf("iter: next failed, want=%q", want)
   205  		}
   206  		if got := string(x.Key().UserKey); got != want {
   207  			t.Fatalf("iter: got %q, want %q", got, want)
   208  		}
   209  		if k := x.Key().UserKey; len(k) != cap(k) {
   210  			t.Fatalf("iter: len(k)=%d, cap(k)=%d", len(k), cap(k))
   211  		}
   212  		if v := x.Value(); len(v) != cap(v) {
   213  			t.Fatalf("iter: len(v)=%d, cap(v)=%d", len(v), cap(v))
   214  		}
   215  		x.Next()
   216  	}
   217  	if err := x.Close(); err != nil {
   218  		t.Fatalf("close: %v", err)
   219  	}
   220  
   221  	if err := m0.close(); err != nil {
   222  		t.Fatalf("close: %v", err)
   223  	}
   224  }
   225  
   226  func buildMemTable(b *testing.B) (*memTable, [][]byte) {
   227  	m := testNewMemTable()
   228  	var keys [][]byte
   229  	var ikey InternalKey
   230  	for i := 0; ; i++ {
   231  		key := []byte(fmt.Sprintf("%08d", i))
   232  		keys = append(keys, key)
   233  		ikey = base.MakeInternalKey(key, 0, InternalKeyKindSet)
   234  		if m.set(ikey, nil) == arenaskl.ErrArenaFull {
   235  			break
   236  		}
   237  	}
   238  	return m, keys
   239  }
   240  
   241  func BenchmarkMemTableIterSeekGE(b *testing.B) {
   242  	m, keys := buildMemTable(b)
   243  	iter := m.newIter(nil)
   244  	rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano())))
   245  
   246  	b.ResetTimer()
   247  	for i := 0; i < b.N; i++ {
   248  		iter.SeekGE(keys[rng.Intn(len(keys))])
   249  	}
   250  }
   251  
   252  func BenchmarkMemTableIterNext(b *testing.B) {
   253  	m, _ := buildMemTable(b)
   254  	iter := m.newIter(nil)
   255  	_, _ = iter.First()
   256  	b.ResetTimer()
   257  	for i := 0; i < b.N; i++ {
   258  		key, _ := iter.Next()
   259  		if key == nil {
   260  			key, _ = iter.First()
   261  		}
   262  		_ = key
   263  	}
   264  }
   265  
   266  func BenchmarkMemTableIterPrev(b *testing.B) {
   267  	m, _ := buildMemTable(b)
   268  	iter := m.newIter(nil)
   269  	_, _ = iter.Last()
   270  	b.ResetTimer()
   271  	for i := 0; i < b.N; i++ {
   272  		key, _ := iter.Prev()
   273  		if key == nil {
   274  			key, _ = iter.Last()
   275  		}
   276  		_ = key
   277  	}
   278  }