github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/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 "context" 10 "fmt" 11 "strconv" 12 "strings" 13 "sync/atomic" 14 "testing" 15 "time" 16 17 "github.com/cockroachdb/datadriven" 18 "github.com/cockroachdb/errors" 19 "github.com/cockroachdb/pebble/internal/arenaskl" 20 "github.com/cockroachdb/pebble/internal/base" 21 "github.com/cockroachdb/pebble/internal/rangekey" 22 "github.com/stretchr/testify/require" 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, InternalKeyKindDeleteSized: 40 return nil, ErrNotFound 41 default: 42 return val.InPlaceValue(), 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(t *testing.T, 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(t, 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(t *testing.T, 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 var buf bytes.Buffer 326 if td.HasArg("range-del") { 327 iter := mem.newRangeDelIter(nil) 328 defer iter.Close() 329 scanKeyspanIterator(&buf, iter) 330 } else { 331 iter := mem.newIter(nil) 332 defer iter.Close() 333 scanInternalIter(&buf, iter) 334 } 335 return buf.String() 336 337 default: 338 return fmt.Sprintf("unknown command: %s", td.Cmd) 339 } 340 }) 341 } 342 343 func TestMemTableConcurrentDeleteRange(t *testing.T) { 344 // Concurrently write and read range tombstones. Workers add range 345 // tombstones, and then immediately retrieve them verifying that the 346 // tombstones they've added are all present. 347 348 m := newMemTable(memTableOptions{Options: &Options{MemTableSize: 64 << 20}}) 349 350 const workers = 10 351 eg, _ := errgroup.WithContext(context.Background()) 352 var seqNum atomic.Uint64 353 seqNum.Store(1) 354 for i := 0; i < workers; i++ { 355 i := i 356 eg.Go(func() error { 357 start := ([]byte)(fmt.Sprintf("%03d", i)) 358 end := ([]byte)(fmt.Sprintf("%03d", i+1)) 359 for j := 0; j < 100; j++ { 360 b := newBatch(nil) 361 b.DeleteRange(start, end, nil) 362 n := seqNum.Add(1) - 1 363 require.NoError(t, m.apply(b, n)) 364 b.release() 365 366 var count int 367 it := m.newRangeDelIter(nil) 368 for s := it.SeekGE(start); s != nil; s = it.Next() { 369 if m.cmp(s.Start, end) >= 0 { 370 break 371 } 372 count += len(s.Keys) 373 } 374 if j+1 != count { 375 return errors.Errorf("%d: expected %d tombstones, but found %d", i, j+1, count) 376 } 377 } 378 return nil 379 }) 380 } 381 err := eg.Wait() 382 if err != nil { 383 t.Error(err) 384 } 385 } 386 387 func TestMemTableReserved(t *testing.T) { 388 m := newMemTable(memTableOptions{size: 5000}) 389 // Increase to 2 references. 390 m.writerRef() 391 // The initial reservation accounts for the already allocated bytes from the 392 // arena. 393 require.Equal(t, m.reserved, m.skl.Arena().Size()) 394 b := newBatch(nil) 395 b.Set([]byte("blueberry"), []byte("pie"), nil) 396 require.NotEqual(t, 0, int(b.memTableSize)) 397 prevReserved := m.reserved 398 m.prepare(b) 399 require.Equal(t, int(m.reserved), int(b.memTableSize)+int(prevReserved)) 400 } 401 402 func buildMemTable(b *testing.B) (*memTable, [][]byte) { 403 m := newMemTable(memTableOptions{}) 404 var keys [][]byte 405 var ikey InternalKey 406 for i := 0; ; i++ { 407 key := []byte(fmt.Sprintf("%08d", i)) 408 keys = append(keys, key) 409 ikey = base.MakeInternalKey(key, 0, InternalKeyKindSet) 410 if m.set(ikey, nil) == arenaskl.ErrArenaFull { 411 break 412 } 413 } 414 return m, keys 415 } 416 417 func BenchmarkMemTableIterSeekGE(b *testing.B) { 418 m, keys := buildMemTable(b) 419 iter := m.newIter(nil) 420 rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano()))) 421 422 b.ResetTimer() 423 for i := 0; i < b.N; i++ { 424 iter.SeekGE(keys[rng.Intn(len(keys))], base.SeekGEFlagsNone) 425 } 426 } 427 428 func BenchmarkMemTableIterNext(b *testing.B) { 429 m, _ := buildMemTable(b) 430 iter := m.newIter(nil) 431 _, _ = iter.First() 432 b.ResetTimer() 433 for i := 0; i < b.N; i++ { 434 key, _ := iter.Next() 435 if key == nil { 436 key, _ = iter.First() 437 } 438 _ = key 439 } 440 } 441 442 func BenchmarkMemTableIterPrev(b *testing.B) { 443 m, _ := buildMemTable(b) 444 iter := m.newIter(nil) 445 _, _ = iter.Last() 446 b.ResetTimer() 447 for i := 0; i < b.N; i++ { 448 key, _ := iter.Prev() 449 if key == nil { 450 key, _ = iter.Last() 451 } 452 _ = key 453 } 454 }