github.com/petermattis/pebble@v0.0.0-20190905164901-ab51a2166067/mem_table_test.go (about)

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