github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/db_test.go (about) 1 // Copyright 2012 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 "fmt" 10 "io" 11 "path/filepath" 12 "sort" 13 "strconv" 14 "strings" 15 "sync" 16 "sync/atomic" 17 "testing" 18 "time" 19 20 "github.com/cockroachdb/errors" 21 "github.com/stretchr/testify/require" 22 "github.com/zuoyebang/bitalostable/internal/base" 23 "github.com/zuoyebang/bitalostable/sstable" 24 "github.com/zuoyebang/bitalostable/vfs" 25 "golang.org/x/exp/rand" 26 ) 27 28 // try repeatedly calls f, sleeping between calls with exponential back-off, 29 // until f returns a nil error or the total sleep time is greater than or equal 30 // to maxTotalSleep. It always calls f at least once. 31 func try(initialSleep, maxTotalSleep time.Duration, f func() error) error { 32 totalSleep := time.Duration(0) 33 for d := initialSleep; ; d *= 2 { 34 time.Sleep(d) 35 totalSleep += d 36 if err := f(); err == nil || totalSleep >= maxTotalSleep { 37 return err 38 } 39 } 40 } 41 42 func TestTry(t *testing.T) { 43 c := make(chan struct{}) 44 go func() { 45 time.Sleep(1 * time.Millisecond) 46 close(c) 47 }() 48 49 attemptsMu := sync.Mutex{} 50 attempts := 0 51 52 err := try(100*time.Microsecond, 20*time.Second, func() error { 53 attemptsMu.Lock() 54 attempts++ 55 attemptsMu.Unlock() 56 57 select { 58 default: 59 return errors.New("timed out") 60 case <-c: 61 return nil 62 } 63 }) 64 require.NoError(t, err) 65 66 attemptsMu.Lock() 67 a := attempts 68 attemptsMu.Unlock() 69 70 if a == 0 { 71 t.Fatalf("attempts: got 0, want > 0") 72 } 73 } 74 75 func TestBasicReads(t *testing.T) { 76 testCases := []struct { 77 dirname string 78 wantMap map[string]string 79 }{ 80 { 81 "db-stage-1", 82 map[string]string{ 83 "aaa": "", 84 "bar": "", 85 "baz": "", 86 "foo": "", 87 "quux": "", 88 "zzz": "", 89 }, 90 }, 91 { 92 "db-stage-2", 93 map[string]string{ 94 "aaa": "", 95 "bar": "", 96 "baz": "three", 97 "foo": "four", 98 "quux": "", 99 "zzz": "", 100 }, 101 }, 102 { 103 "db-stage-3", 104 map[string]string{ 105 "aaa": "", 106 "bar": "", 107 "baz": "three", 108 "foo": "four", 109 "quux": "", 110 "zzz": "", 111 }, 112 }, 113 { 114 "db-stage-4", 115 map[string]string{ 116 "aaa": "", 117 "bar": "", 118 "baz": "", 119 "foo": "five", 120 "quux": "six", 121 "zzz": "", 122 }, 123 }, 124 } 125 for _, tc := range testCases { 126 t.Run(tc.dirname, func(t *testing.T) { 127 fs := vfs.NewMem() 128 _, err := vfs.Clone(vfs.Default, fs, filepath.Join("testdata", tc.dirname), tc.dirname) 129 if err != nil { 130 t.Fatalf("%s: cloneFileSystem failed: %v", tc.dirname, err) 131 } 132 d, err := Open(tc.dirname, testingRandomized(&Options{ 133 FS: fs, 134 })) 135 if err != nil { 136 t.Fatalf("%s: Open failed: %v", tc.dirname, err) 137 } 138 for key, want := range tc.wantMap { 139 got, closer, err := d.Get([]byte(key)) 140 if err != nil && err != ErrNotFound { 141 t.Fatalf("%s: Get(%q) failed: %v", tc.dirname, key, err) 142 } 143 if string(got) != string(want) { 144 t.Fatalf("%s: Get(%q): got %q, want %q", tc.dirname, key, got, want) 145 } 146 if closer != nil { 147 closer.Close() 148 } 149 } 150 err = d.Close() 151 if err != nil { 152 t.Fatalf("%s: Close failed: %v", tc.dirname, err) 153 } 154 }) 155 } 156 } 157 158 func TestBasicWrites(t *testing.T) { 159 d, err := Open("", testingRandomized(&Options{ 160 FS: vfs.NewMem(), 161 })) 162 require.NoError(t, err) 163 164 names := []string{ 165 "Alatar", 166 "Gandalf", 167 "Pallando", 168 "Radagast", 169 "Saruman", 170 "Joe", 171 } 172 wantMap := map[string]string{} 173 174 inBatch, batch, pending := false, &Batch{}, [][]string(nil) 175 set0 := func(k, v string) error { 176 return d.Set([]byte(k), []byte(v), nil) 177 } 178 del0 := func(k string) error { 179 return d.Delete([]byte(k), nil) 180 } 181 set1 := func(k, v string) error { 182 batch.Set([]byte(k), []byte(v), nil) 183 return nil 184 } 185 del1 := func(k string) error { 186 batch.Delete([]byte(k), nil) 187 return nil 188 } 189 set, del := set0, del0 190 191 testCases := []string{ 192 "set Gandalf Grey", 193 "set Saruman White", 194 "set Radagast Brown", 195 "delete Saruman", 196 "set Gandalf White", 197 "batch", 198 " set Alatar AliceBlue", 199 "apply", 200 "delete Pallando", 201 "set Alatar AntiqueWhite", 202 "set Pallando PapayaWhip", 203 "batch", 204 "apply", 205 "set Pallando PaleVioletRed", 206 "batch", 207 " delete Alatar", 208 " set Gandalf GhostWhite", 209 " set Saruman Seashell", 210 " delete Saruman", 211 " set Saruman SeaGreen", 212 " set Radagast RosyBrown", 213 " delete Pallando", 214 "apply", 215 "delete Radagast", 216 "delete Radagast", 217 "delete Radagast", 218 "set Gandalf Goldenrod", 219 "set Pallando PeachPuff", 220 "batch", 221 " delete Joe", 222 " delete Saruman", 223 " delete Radagast", 224 " delete Pallando", 225 " delete Gandalf", 226 " delete Alatar", 227 "apply", 228 "set Joe Plumber", 229 } 230 for i, tc := range testCases { 231 s := strings.Split(strings.TrimSpace(tc), " ") 232 switch s[0] { 233 case "set": 234 if err := set(s[1], s[2]); err != nil { 235 t.Fatalf("#%d %s: %v", i, tc, err) 236 } 237 if inBatch { 238 pending = append(pending, s) 239 } else { 240 wantMap[s[1]] = s[2] 241 } 242 case "delete": 243 if err := del(s[1]); err != nil { 244 t.Fatalf("#%d %s: %v", i, tc, err) 245 } 246 if inBatch { 247 pending = append(pending, s) 248 } else { 249 delete(wantMap, s[1]) 250 } 251 case "batch": 252 inBatch, batch, set, del = true, &Batch{}, set1, del1 253 case "apply": 254 if err := d.Apply(batch, nil); err != nil { 255 t.Fatalf("#%d %s: %v", i, tc, err) 256 } 257 for _, p := range pending { 258 switch p[0] { 259 case "set": 260 wantMap[p[1]] = p[2] 261 case "delete": 262 delete(wantMap, p[1]) 263 } 264 } 265 inBatch, pending, set, del = false, nil, set0, del0 266 default: 267 t.Fatalf("#%d %s: bad test case: %q", i, tc, s) 268 } 269 270 fail := false 271 for _, name := range names { 272 g, closer, err := d.Get([]byte(name)) 273 if err != nil && err != ErrNotFound { 274 t.Errorf("#%d %s: Get(%q): %v", i, tc, name, err) 275 fail = true 276 } 277 got, gOK := string(g), err == nil 278 want, wOK := wantMap[name] 279 if got != want || gOK != wOK { 280 t.Errorf("#%d %s: Get(%q): got %q, %t, want %q, %t", 281 i, tc, name, got, gOK, want, wOK) 282 fail = true 283 } 284 if closer != nil { 285 closer.Close() 286 } 287 } 288 if fail { 289 return 290 } 291 } 292 293 require.NoError(t, d.Close()) 294 } 295 296 func TestRandomWrites(t *testing.T) { 297 d, err := Open("", testingRandomized(&Options{ 298 FS: vfs.NewMem(), 299 MemTableSize: 8 * 1024, 300 })) 301 require.NoError(t, err) 302 303 keys := [64][]byte{} 304 wants := [64]int{} 305 for k := range keys { 306 keys[k] = []byte(strconv.Itoa(k)) 307 wants[k] = -1 308 } 309 xxx := bytes.Repeat([]byte("x"), 512) 310 311 rng := rand.New(rand.NewSource(123)) 312 const N = 1000 313 for i := 0; i < N; i++ { 314 k := rng.Intn(len(keys)) 315 if rng.Intn(20) != 0 { 316 wants[k] = rng.Intn(len(xxx) + 1) 317 if err := d.Set(keys[k], xxx[:wants[k]], nil); err != nil { 318 t.Fatalf("i=%d: Set: %v", i, err) 319 } 320 } else { 321 wants[k] = -1 322 if err := d.Delete(keys[k], nil); err != nil { 323 t.Fatalf("i=%d: Delete: %v", i, err) 324 } 325 } 326 327 if i != N-1 || rng.Intn(50) != 0 { 328 continue 329 } 330 for k := range keys { 331 got := -1 332 if v, closer, err := d.Get(keys[k]); err != nil { 333 if err != ErrNotFound { 334 t.Fatalf("Get: %v", err) 335 } 336 } else { 337 got = len(v) 338 closer.Close() 339 } 340 if got != wants[k] { 341 t.Errorf("i=%d, k=%d: got %d, want %d", i, k, got, wants[k]) 342 } 343 } 344 } 345 346 require.NoError(t, d.Close()) 347 } 348 349 func TestLargeBatch(t *testing.T) { 350 d, err := Open("", testingRandomized(&Options{ 351 FS: vfs.NewMem(), 352 MemTableSize: 1400, 353 MemTableStopWritesThreshold: 100, 354 })) 355 require.NoError(t, err) 356 357 verifyLSM := func(expected string) func() error { 358 return func() error { 359 d.mu.Lock() 360 s := d.mu.versions.currentVersion().String() 361 d.mu.Unlock() 362 if expected != s { 363 if testing.Verbose() { 364 fmt.Println(strings.TrimSpace(s)) 365 } 366 return errors.Errorf("expected %s, but found %s", expected, s) 367 } 368 return nil 369 } 370 } 371 372 logNum := func() FileNum { 373 d.mu.Lock() 374 defer d.mu.Unlock() 375 return d.mu.log.queue[len(d.mu.log.queue)-1].fileNum 376 } 377 fileSize := func(fileNum FileNum) int64 { 378 info, err := d.opts.FS.Stat(base.MakeFilepath(d.opts.FS, "", fileTypeLog, fileNum)) 379 require.NoError(t, err) 380 return info.Size() 381 } 382 memTableCreationSeqNum := func() uint64 { 383 d.mu.Lock() 384 defer d.mu.Unlock() 385 return d.mu.mem.mutable.logSeqNum 386 } 387 388 startLogNum := logNum() 389 startLogStartSize := fileSize(startLogNum) 390 startSeqNum := atomic.LoadUint64(&d.mu.versions.atomic.logSeqNum) 391 392 // Write a key with a value larger than the memtable size. 393 require.NoError(t, d.Set([]byte("a"), bytes.Repeat([]byte("a"), 512), nil)) 394 395 // Verify that the large batch was written to the WAL that existed before it 396 // was committed. We verify that WAL rotation occurred, where the large batch 397 // was written to, and that the new WAL is empty. 398 endLogNum := logNum() 399 if startLogNum == endLogNum { 400 t.Fatal("expected WAL rotation") 401 } 402 startLogEndSize := fileSize(startLogNum) 403 if startLogEndSize == startLogStartSize { 404 t.Fatalf("expected large batch to be written to %s.log, but file size unchanged at %d", 405 startLogNum, startLogEndSize) 406 } 407 endLogSize := fileSize(endLogNum) 408 if endLogSize != 0 { 409 t.Fatalf("expected %s.log to be empty, but found %d", endLogNum, endLogSize) 410 } 411 if creationSeqNum := memTableCreationSeqNum(); creationSeqNum <= startSeqNum { 412 t.Fatalf("expected memTable.logSeqNum=%d > largeBatch.seqNum=%d", creationSeqNum, startSeqNum) 413 } 414 415 // Verify this results in one L0 table being created. 416 require.NoError(t, try(100*time.Microsecond, 20*time.Second, 417 verifyLSM("0.0:\n 000005:[a#1,SET-a#1,SET]\n"))) 418 419 require.NoError(t, d.Set([]byte("b"), bytes.Repeat([]byte("b"), 512), nil)) 420 421 // Verify this results in a second L0 table being created. 422 require.NoError(t, try(100*time.Microsecond, 20*time.Second, 423 verifyLSM("0.0:\n 000005:[a#1,SET-a#1,SET]\n 000007:[b#2,SET-b#2,SET]\n"))) 424 425 // Allocate a bunch of batches to exhaust the batchPool. None of these 426 // batches should have a non-zero count. 427 for i := 0; i < 10; i++ { 428 b := d.NewBatch() 429 require.EqualValues(t, 0, b.Count()) 430 } 431 432 require.NoError(t, d.Close()) 433 } 434 435 func TestGetNoCache(t *testing.T) { 436 cache := NewCache(0) 437 defer cache.Unref() 438 439 d, err := Open("", testingRandomized(&Options{ 440 Cache: cache, 441 FS: vfs.NewMem(), 442 })) 443 require.NoError(t, err) 444 445 require.NoError(t, d.Set([]byte("a"), []byte("aa"), nil)) 446 require.NoError(t, d.Flush()) 447 verifyGet(t, d, []byte("a"), []byte("aa")) 448 449 require.NoError(t, d.Close()) 450 } 451 452 func TestGetMerge(t *testing.T) { 453 d, err := Open("", testingRandomized(&Options{ 454 FS: vfs.NewMem(), 455 })) 456 require.NoError(t, err) 457 458 key := []byte("a") 459 verify := func(expected string) { 460 val, closer, err := d.Get(key) 461 require.NoError(t, err) 462 463 if expected != string(val) { 464 t.Fatalf("expected %s, but got %s", expected, val) 465 } 466 closer.Close() 467 } 468 469 const val = "1" 470 for i := 1; i <= 3; i++ { 471 require.NoError(t, d.Merge(key, []byte(val), nil)) 472 473 expected := strings.Repeat(val, i) 474 verify(expected) 475 476 require.NoError(t, d.Flush()) 477 verify(expected) 478 } 479 480 require.NoError(t, d.Close()) 481 } 482 483 func TestMergeOrderSameAfterFlush(t *testing.T) { 484 // Ensure compaction iterator (used by flush) and user iterator process merge 485 // operands in the same order 486 d, err := Open("", testingRandomized(&Options{ 487 FS: vfs.NewMem(), 488 })) 489 require.NoError(t, err) 490 491 key := []byte("a") 492 verify := func(expected string) { 493 iter := d.NewIter(nil) 494 if !iter.SeekGE([]byte("a")) { 495 t.Fatal("expected one value, but got empty iterator") 496 } 497 if expected != string(iter.Value()) { 498 t.Fatalf("expected %s, but got %s", expected, string(iter.Value())) 499 } 500 if !iter.SeekLT([]byte("b")) { 501 t.Fatal("expected one value, but got empty iterator") 502 } 503 if expected != string(iter.Value()) { 504 t.Fatalf("expected %s, but got %s", expected, string(iter.Value())) 505 } 506 require.NoError(t, iter.Close()) 507 } 508 509 require.NoError(t, d.Merge(key, []byte("0"), nil)) 510 require.NoError(t, d.Merge(key, []byte("1"), nil)) 511 512 verify("01") 513 require.NoError(t, d.Flush()) 514 verify("01") 515 516 require.NoError(t, d.Close()) 517 } 518 519 type closableMerger struct { 520 lastBuf []byte 521 closed bool 522 } 523 524 func (m *closableMerger) MergeNewer(value []byte) error { 525 m.lastBuf = append(m.lastBuf[:0], value...) 526 return nil 527 } 528 529 func (m *closableMerger) MergeOlder(value []byte) error { 530 m.lastBuf = append(m.lastBuf[:0], value...) 531 return nil 532 } 533 534 func (m *closableMerger) Finish(includesBase bool) ([]byte, io.Closer, error) { 535 return m.lastBuf, m, nil 536 } 537 538 func (m *closableMerger) Close() error { 539 m.closed = true 540 return nil 541 } 542 543 func TestMergerClosing(t *testing.T) { 544 m := &closableMerger{} 545 546 d, err := Open("", testingRandomized(&Options{ 547 FS: vfs.NewMem(), 548 Merger: &Merger{ 549 Merge: func(key, value []byte) (base.ValueMerger, error) { 550 return m, m.MergeNewer(value) 551 }, 552 }, 553 })) 554 require.NoError(t, err) 555 556 defer func() { 557 require.NoError(t, d.Close()) 558 }() 559 560 err = d.Merge([]byte("a"), []byte("b"), nil) 561 require.NoError(t, err) 562 require.False(t, m.closed) 563 564 val, closer, err := d.Get([]byte("a")) 565 require.NoError(t, err) 566 require.Equal(t, []byte("b"), val) 567 require.NotNil(t, closer) 568 require.False(t, m.closed) 569 _ = closer.Close() 570 require.True(t, m.closed) 571 } 572 573 func TestLogData(t *testing.T) { 574 d, err := Open("", testingRandomized(&Options{ 575 FS: vfs.NewMem(), 576 })) 577 require.NoError(t, err) 578 579 defer func() { 580 require.NoError(t, d.Close()) 581 }() 582 583 require.NoError(t, d.LogData([]byte("foo"), Sync)) 584 require.NoError(t, d.LogData([]byte("bar"), Sync)) 585 // TODO(itsbilal): Confirm that we wrote some bytes to the WAL. 586 // For now, LogData proceeding ahead without a panic is good enough. 587 } 588 589 func TestSingleDeleteGet(t *testing.T) { 590 d, err := Open("", testingRandomized(&Options{ 591 FS: vfs.NewMem(), 592 })) 593 require.NoError(t, err) 594 defer func() { 595 require.NoError(t, d.Close()) 596 }() 597 598 key := []byte("key") 599 val := []byte("val") 600 601 require.NoError(t, d.Set(key, val, nil)) 602 verifyGet(t, d, key, val) 603 604 key2 := []byte("key2") 605 val2 := []byte("val2") 606 607 require.NoError(t, d.Set(key2, val2, nil)) 608 verifyGet(t, d, key2, val2) 609 610 require.NoError(t, d.SingleDelete(key2, nil)) 611 verifyGetNotFound(t, d, key2) 612 } 613 614 func TestSingleDeleteFlush(t *testing.T) { 615 d, err := Open("", testingRandomized(&Options{ 616 FS: vfs.NewMem(), 617 })) 618 require.NoError(t, err) 619 defer func() { 620 require.NoError(t, d.Close()) 621 }() 622 623 key := []byte("key") 624 valFirst := []byte("first") 625 valSecond := []byte("second") 626 key2 := []byte("key2") 627 val2 := []byte("val2") 628 629 require.NoError(t, d.Set(key, valFirst, nil)) 630 require.NoError(t, d.Set(key2, val2, nil)) 631 require.NoError(t, d.Flush()) 632 633 require.NoError(t, d.SingleDelete(key, nil)) 634 require.NoError(t, d.Set(key, valSecond, nil)) 635 require.NoError(t, d.Delete(key2, nil)) 636 require.NoError(t, d.Set(key2, val2, nil)) 637 require.NoError(t, d.Flush()) 638 639 require.NoError(t, d.SingleDelete(key, nil)) 640 require.NoError(t, d.Delete(key2, nil)) 641 require.NoError(t, d.Flush()) 642 643 verifyGetNotFound(t, d, key) 644 verifyGetNotFound(t, d, key2) 645 } 646 647 func TestUnremovableSingleDelete(t *testing.T) { 648 d, err := Open("", testingRandomized(&Options{ 649 FS: vfs.NewMem(), 650 L0CompactionThreshold: 8, 651 })) 652 require.NoError(t, err) 653 defer func() { 654 require.NoError(t, d.Close()) 655 }() 656 657 key := []byte("key") 658 valFirst := []byte("valFirst") 659 valSecond := []byte("valSecond") 660 661 require.NoError(t, d.Set(key, valFirst, nil)) 662 ss := d.NewSnapshot() 663 require.NoError(t, d.SingleDelete(key, nil)) 664 require.NoError(t, d.Set(key, valSecond, nil)) 665 require.NoError(t, d.Flush()) 666 667 verifyGet(t, ss, key, valFirst) 668 verifyGet(t, d, key, valSecond) 669 670 require.NoError(t, d.SingleDelete(key, nil)) 671 672 verifyGet(t, ss, key, valFirst) 673 verifyGetNotFound(t, d, key) 674 675 require.NoError(t, d.Flush()) 676 677 verifyGet(t, ss, key, valFirst) 678 verifyGetNotFound(t, d, key) 679 } 680 681 func TestIterLeak(t *testing.T) { 682 for _, leak := range []bool{true, false} { 683 t.Run(fmt.Sprintf("leak=%t", leak), func(t *testing.T) { 684 for _, flush := range []bool{true, false} { 685 t.Run(fmt.Sprintf("flush=%t", flush), func(t *testing.T) { 686 d, err := Open("", testingRandomized(&Options{ 687 FS: vfs.NewMem(), 688 })) 689 require.NoError(t, err) 690 691 require.NoError(t, d.Set([]byte("a"), []byte("a"), nil)) 692 if flush { 693 require.NoError(t, d.Flush()) 694 } 695 iter := d.NewIter(nil) 696 iter.First() 697 if !leak { 698 require.NoError(t, iter.Close()) 699 require.NoError(t, d.Close()) 700 } else { 701 defer iter.Close() 702 if err := d.Close(); err == nil { 703 t.Fatalf("expected failure, but found success") 704 } else if !strings.HasPrefix(err.Error(), "leaked iterators:") { 705 t.Fatalf("expected leaked iterators, but found %+v", err) 706 } else { 707 t.Log(err.Error()) 708 } 709 } 710 }) 711 } 712 }) 713 } 714 } 715 716 // Make sure that we detect an iter leak when only one DB closes 717 // while the second db still holds a reference to the TableCache. 718 func TestIterLeakSharedCache(t *testing.T) { 719 for _, leak := range []bool{true, false} { 720 t.Run(fmt.Sprintf("leak=%t", leak), func(t *testing.T) { 721 for _, flush := range []bool{true, false} { 722 t.Run(fmt.Sprintf("flush=%t", flush), func(t *testing.T) { 723 d1, err := Open("", &Options{ 724 FS: vfs.NewMem(), 725 }) 726 require.NoError(t, err) 727 728 d2, err := Open("", &Options{ 729 FS: vfs.NewMem(), 730 }) 731 require.NoError(t, err) 732 733 require.NoError(t, d1.Set([]byte("a"), []byte("a"), nil)) 734 if flush { 735 require.NoError(t, d1.Flush()) 736 } 737 738 require.NoError(t, d2.Set([]byte("a"), []byte("a"), nil)) 739 if flush { 740 require.NoError(t, d2.Flush()) 741 } 742 743 // Check if leak detection works with only one db closing. 744 { 745 iter1 := d1.NewIter(nil) 746 iter1.First() 747 if !leak { 748 require.NoError(t, iter1.Close()) 749 require.NoError(t, d1.Close()) 750 } else { 751 defer iter1.Close() 752 if err := d1.Close(); err == nil { 753 t.Fatalf("expected failure, but found success") 754 } else if !strings.HasPrefix(err.Error(), "leaked iterators:") { 755 t.Fatalf("expected leaked iterators, but found %+v", err) 756 } else { 757 t.Log(err.Error()) 758 } 759 } 760 } 761 762 { 763 iter2 := d2.NewIter(nil) 764 iter2.First() 765 if !leak { 766 require.NoError(t, iter2.Close()) 767 require.NoError(t, d2.Close()) 768 } else { 769 defer iter2.Close() 770 if err := d2.Close(); err == nil { 771 t.Fatalf("expected failure, but found success") 772 } else if !strings.HasPrefix(err.Error(), "leaked iterators:") { 773 t.Fatalf("expected leaked iterators, but found %+v", err) 774 } else { 775 t.Log(err.Error()) 776 } 777 } 778 } 779 780 }) 781 } 782 }) 783 } 784 } 785 786 func TestMemTableReservation(t *testing.T) { 787 cache := NewCache(128 << 10 /* 128 KB */) 788 defer cache.Unref() 789 790 opts := &Options{ 791 Cache: cache, 792 MemTableSize: initialMemTableSize, 793 FS: vfs.NewMem(), 794 } 795 opts.testingRandomized() 796 opts.EnsureDefaults() 797 // We're going to be looking at and asserting the global memtable reservation 798 // amount below so we don't want to race with any triggered stats collections. 799 opts.private.disableTableStats = true 800 801 // Add a block to the cache. Note that the memtable size is larger than the 802 // cache size, so opening the DB should cause this block to be evicted. 803 tmpID := opts.Cache.NewID() 804 helloWorld := []byte("hello world") 805 value := opts.Cache.Alloc(len(helloWorld)) 806 copy(value.Buf(), helloWorld) 807 opts.Cache.Set(tmpID, 0, 0, value).Release() 808 809 d, err := Open("", opts) 810 require.NoError(t, err) 811 812 checkReserved := func(expected int64) { 813 t.Helper() 814 if reserved := atomic.LoadInt64(&d.atomic.memTableReserved); expected != reserved { 815 t.Fatalf("expected %d reserved, but found %d", expected, reserved) 816 } 817 } 818 819 checkReserved(int64(opts.MemTableSize)) 820 if refs := atomic.LoadInt32(&d.mu.mem.queue[len(d.mu.mem.queue)-1].readerRefs); refs != 2 { 821 t.Fatalf("expected 2 refs, but found %d", refs) 822 } 823 // Verify the memtable reservation has caused our test block to be evicted. 824 if h := opts.Cache.Get(tmpID, 0, 0); h.Get() != nil { 825 t.Fatalf("expected failure, but found success: %s", h.Get()) 826 } 827 828 // Flush the memtable. The memtable reservation should be unchanged because a 829 // new memtable will be allocated. 830 require.NoError(t, d.Flush()) 831 checkReserved(int64(opts.MemTableSize)) 832 833 // Flush in the presence of an active iterator. The iterator will hold a 834 // reference to a readState which will in turn hold a reader reference to the 835 // memtable. While the iterator is open, there will be the memory for 2 836 // memtables reserved. 837 iter := d.NewIter(nil) 838 require.NoError(t, d.Flush()) 839 checkReserved(2 * int64(opts.MemTableSize)) 840 require.NoError(t, iter.Close()) 841 checkReserved(int64(opts.MemTableSize)) 842 843 require.NoError(t, d.Close()) 844 } 845 846 func TestMemTableReservationLeak(t *testing.T) { 847 d, err := Open("", &Options{FS: vfs.NewMem()}) 848 require.NoError(t, err) 849 850 d.mu.Lock() 851 last := d.mu.mem.queue[len(d.mu.mem.queue)-1] 852 last.readerRef() 853 defer last.readerUnref() 854 d.mu.Unlock() 855 if err := d.Close(); err == nil { 856 t.Fatalf("expected failure, but found success") 857 } else if !strings.HasPrefix(err.Error(), "leaked memtable reservation:") { 858 t.Fatalf("expected leaked memtable reservation, but found %+v", err) 859 } else { 860 t.Log(err.Error()) 861 } 862 } 863 864 func TestCacheEvict(t *testing.T) { 865 cache := NewCache(10 << 20) 866 defer cache.Unref() 867 868 d, err := Open("", &Options{ 869 Cache: cache, 870 FS: vfs.NewMem(), 871 }) 872 require.NoError(t, err) 873 874 for i := 0; i < 1000; i++ { 875 key := []byte(fmt.Sprintf("%04d", i)) 876 require.NoError(t, d.Set(key, key, nil)) 877 } 878 879 require.NoError(t, d.Flush()) 880 iter := d.NewIter(nil) 881 for iter.First(); iter.Valid(); iter.Next() { 882 } 883 require.NoError(t, iter.Close()) 884 885 if size := cache.Size(); size == 0 { 886 t.Fatalf("expected non-zero cache size") 887 } 888 889 for i := 0; i < 1000; i++ { 890 key := []byte(fmt.Sprintf("%04d", i)) 891 require.NoError(t, d.Delete(key, nil)) 892 } 893 894 require.NoError(t, d.Compact([]byte("0"), []byte("1"), false)) 895 896 require.NoError(t, d.Close()) 897 898 if size := cache.Size(); size != 0 { 899 t.Fatalf("expected empty cache, but found %d", size) 900 } 901 } 902 903 func TestFlushEmpty(t *testing.T) { 904 d, err := Open("", testingRandomized(&Options{ 905 FS: vfs.NewMem(), 906 })) 907 require.NoError(t, err) 908 909 // Flushing an empty memtable should not fail. 910 require.NoError(t, d.Flush()) 911 require.NoError(t, d.Close()) 912 } 913 914 func TestRollManifest(t *testing.T) { 915 toPreserve := rand.Int31n(5) + 1 916 opts := &Options{ 917 MaxManifestFileSize: 1, 918 L0CompactionThreshold: 10, 919 FS: vfs.NewMem(), 920 NumPrevManifest: int(toPreserve), 921 } 922 opts.DisableAutomaticCompactions = true 923 opts.testingRandomized() 924 d, err := Open("", opts) 925 require.NoError(t, err) 926 927 manifestFileNumber := func() FileNum { 928 d.mu.Lock() 929 defer d.mu.Unlock() 930 return d.mu.versions.manifestFileNum 931 } 932 933 current := func() string { 934 desc, err := Peek(d.dirname, d.opts.FS) 935 require.NoError(t, err) 936 return desc.ManifestFilename 937 } 938 939 lastManifestNum := manifestFileNumber() 940 manifestNums := []base.FileNum{lastManifestNum} 941 for i := 0; i < 5; i++ { 942 require.NoError(t, d.Set([]byte("a"), nil, nil)) 943 require.NoError(t, d.Flush()) 944 num := manifestFileNumber() 945 if lastManifestNum == num { 946 t.Fatalf("manifest failed to roll: %d == %d", lastManifestNum, num) 947 } 948 949 manifestNums = append(manifestNums, num) 950 lastManifestNum = num 951 952 expectedCurrent := fmt.Sprintf("MANIFEST-%s", lastManifestNum) 953 if v := current(); expectedCurrent != v { 954 t.Fatalf("expected %s, but found %s", expectedCurrent, v) 955 } 956 } 957 958 files, err := d.opts.FS.List("") 959 require.NoError(t, err) 960 961 var manifests []string 962 for _, filename := range files { 963 fileType, _, ok := base.ParseFilename(d.opts.FS, filename) 964 if !ok { 965 continue 966 } 967 if fileType == fileTypeManifest { 968 manifests = append(manifests, filename) 969 } 970 } 971 972 sort.Slice(manifests, func(i, j int) bool { 973 return manifests[i] < manifests[j] 974 }) 975 976 var expected []string 977 for i := len(manifestNums) - int(toPreserve) - 1; i < len(manifestNums); i++ { 978 expected = append( 979 expected, 980 fmt.Sprintf("MANIFEST-%s", manifestNums[i]), 981 ) 982 } 983 require.EqualValues(t, expected, manifests) 984 985 require.NoError(t, d.Close()) 986 } 987 988 func TestDBClosed(t *testing.T) { 989 d, err := Open("", &Options{ 990 FS: vfs.NewMem(), 991 }) 992 require.NoError(t, err) 993 require.NoError(t, d.Close()) 994 995 catch := func(f func()) (err error) { 996 defer func() { 997 if r := recover(); r != nil { 998 err = r.(error) 999 } 1000 }() 1001 f() 1002 return nil 1003 } 1004 1005 require.True(t, errors.Is(catch(func() { _ = d.Close() }), ErrClosed)) 1006 1007 require.True(t, errors.Is(catch(func() { _ = d.Compact(nil, nil, false) }), ErrClosed)) 1008 require.True(t, errors.Is(catch(func() { _ = d.Flush() }), ErrClosed)) 1009 require.True(t, errors.Is(catch(func() { _, _ = d.AsyncFlush() }), ErrClosed)) 1010 1011 require.True(t, errors.Is(catch(func() { _, _, _ = d.Get(nil) }), ErrClosed)) 1012 require.True(t, errors.Is(catch(func() { _ = d.Delete(nil, nil) }), ErrClosed)) 1013 require.True(t, errors.Is(catch(func() { _ = d.DeleteRange(nil, nil, nil) }), ErrClosed)) 1014 require.True(t, errors.Is(catch(func() { _ = d.Ingest(nil) }), ErrClosed)) 1015 require.True(t, errors.Is(catch(func() { _ = d.LogData(nil, nil) }), ErrClosed)) 1016 require.True(t, errors.Is(catch(func() { _ = d.Merge(nil, nil, nil) }), ErrClosed)) 1017 require.True(t, errors.Is(catch(func() { _ = d.RatchetFormatMajorVersion(FormatNewest) }), ErrClosed)) 1018 require.True(t, errors.Is(catch(func() { _ = d.Set(nil, nil, nil) }), ErrClosed)) 1019 1020 require.True(t, errors.Is(catch(func() { _ = d.NewSnapshot() }), ErrClosed)) 1021 1022 b := d.NewIndexedBatch() 1023 require.True(t, errors.Is(catch(func() { _ = b.Commit(nil) }), ErrClosed)) 1024 require.True(t, errors.Is(catch(func() { _ = d.Apply(b, nil) }), ErrClosed)) 1025 require.True(t, errors.Is(catch(func() { _ = b.NewIter(nil) }), ErrClosed)) 1026 } 1027 1028 func TestDBConcurrentCommitCompactFlush(t *testing.T) { 1029 d, err := Open("", testingRandomized(&Options{ 1030 FS: vfs.NewMem(), 1031 })) 1032 require.NoError(t, err) 1033 1034 // Concurrently commit, compact, and flush in order to stress the locking around 1035 // those operations. 1036 const n = 1000 1037 var wg sync.WaitGroup 1038 wg.Add(n) 1039 for i := 0; i < n; i++ { 1040 go func(i int) { 1041 defer wg.Done() 1042 _ = d.Set([]byte(fmt.Sprint(i)), nil, nil) 1043 var err error 1044 switch i % 3 { 1045 case 0: 1046 err = d.Compact(nil, []byte("\xff"), false) 1047 case 1: 1048 err = d.Flush() 1049 case 2: 1050 _, err = d.AsyncFlush() 1051 } 1052 require.NoError(t, err) 1053 }(i) 1054 } 1055 wg.Wait() 1056 1057 require.NoError(t, d.Close()) 1058 } 1059 1060 func TestDBConcurrentCompactClose(t *testing.T) { 1061 // Test closing while a compaction is ongoing. This ensures compaction code 1062 // detects the close and finishes cleanly. 1063 mem := vfs.NewMem() 1064 for i := 0; i < 100; i++ { 1065 opts := &Options{ 1066 FS: mem, 1067 MaxConcurrentCompactions: func() int { 1068 return 2 1069 }, 1070 } 1071 d, err := Open("", testingRandomized(opts)) 1072 require.NoError(t, err) 1073 1074 // Ingest a series of files containing a single key each. As the outer 1075 // loop progresses, these ingestions will build up compaction debt 1076 // causing compactions to be running concurrently with the close below. 1077 for j := 0; j < 10; j++ { 1078 path := fmt.Sprintf("ext%d", j) 1079 f, err := mem.Create(path) 1080 require.NoError(t, err) 1081 w := sstable.NewWriter(f, sstable.WriterOptions{ 1082 TableFormat: d.FormatMajorVersion().MaxTableFormat(), 1083 }) 1084 require.NoError(t, w.Set([]byte(fmt.Sprint(j)), nil)) 1085 require.NoError(t, w.Close()) 1086 require.NoError(t, d.Ingest([]string{path})) 1087 } 1088 1089 require.NoError(t, d.Close()) 1090 } 1091 } 1092 1093 func TestDBApplyBatchNilDB(t *testing.T) { 1094 d, err := Open("", &Options{FS: vfs.NewMem()}) 1095 require.NoError(t, err) 1096 1097 b1 := &Batch{} 1098 b1.Set([]byte("test"), nil, nil) 1099 1100 b2 := &Batch{} 1101 b2.Apply(b1, nil) 1102 if b2.memTableSize != 0 { 1103 t.Fatalf("expected memTableSize to not be set") 1104 } 1105 require.NoError(t, d.Apply(b2, nil)) 1106 if b1.memTableSize != b2.memTableSize { 1107 t.Fatalf("expected memTableSize %d, but found %d", b1.memTableSize, b2.memTableSize) 1108 } 1109 1110 require.NoError(t, d.Close()) 1111 } 1112 1113 func TestDBApplyBatchMismatch(t *testing.T) { 1114 srcDB, err := Open("", &Options{FS: vfs.NewMem()}) 1115 require.NoError(t, err) 1116 1117 applyDB, err := Open("", &Options{FS: vfs.NewMem()}) 1118 require.NoError(t, err) 1119 1120 err = func() (err error) { 1121 defer func() { 1122 if v := recover(); v != nil { 1123 err = errors.Errorf("%v", v) 1124 } 1125 }() 1126 1127 b := srcDB.NewBatch() 1128 b.Set([]byte("test"), nil, nil) 1129 return applyDB.Apply(b, nil) 1130 }() 1131 if err == nil || !strings.Contains(err.Error(), "bitalostable: batch db mismatch:") { 1132 t.Fatalf("expected error, but found %v", err) 1133 } 1134 1135 require.NoError(t, srcDB.Close()) 1136 require.NoError(t, applyDB.Close()) 1137 } 1138 1139 func TestCloseCleanerRace(t *testing.T) { 1140 mem := vfs.NewMem() 1141 for i := 0; i < 20; i++ { 1142 db, err := Open("", testingRandomized(&Options{FS: mem})) 1143 require.NoError(t, err) 1144 require.NoError(t, db.Set([]byte("a"), []byte("something"), Sync)) 1145 require.NoError(t, db.Flush()) 1146 // Ref the sstables so cannot be deleted. 1147 it := db.NewIter(nil) 1148 require.NotNil(t, it) 1149 require.NoError(t, db.DeleteRange([]byte("a"), []byte("b"), Sync)) 1150 require.NoError(t, db.Compact([]byte("a"), []byte("b"), false)) 1151 // Only the iterator is keeping the sstables alive. 1152 files, err := mem.List("/") 1153 require.NoError(t, err) 1154 var found bool 1155 for _, f := range files { 1156 if strings.HasSuffix(f, ".sst") { 1157 found = true 1158 break 1159 } 1160 } 1161 require.True(t, found) 1162 // Close the iterator and the db in succession so file cleaning races with DB.Close() -- 1163 // latter should wait for file cleaning to finish. 1164 require.NoError(t, it.Close()) 1165 require.NoError(t, db.Close()) 1166 files, err = mem.List("/") 1167 require.NoError(t, err) 1168 for _, f := range files { 1169 if strings.HasSuffix(f, ".sst") { 1170 t.Fatalf("found sst: %s", f) 1171 } 1172 } 1173 } 1174 } 1175 1176 func TestSSTables(t *testing.T) { 1177 d, err := Open("", &Options{ 1178 FS: vfs.NewMem(), 1179 }) 1180 require.NoError(t, err) 1181 defer func() { 1182 if d != nil { 1183 require.NoError(t, d.Close()) 1184 } 1185 }() 1186 1187 // Create two sstables. 1188 require.NoError(t, d.Set([]byte("hello"), nil, nil)) 1189 require.NoError(t, d.Flush()) 1190 require.NoError(t, d.Set([]byte("world"), nil, nil)) 1191 require.NoError(t, d.Flush()) 1192 1193 // by default returned table infos should not contain Properties 1194 tableInfos, err := d.SSTables() 1195 require.NoError(t, err) 1196 for _, levelTables := range tableInfos { 1197 for _, info := range levelTables { 1198 require.Nil(t, info.Properties) 1199 } 1200 } 1201 1202 // with opt `WithProperties()` the `Properties` in table info should not be nil 1203 tableInfos, err = d.SSTables(WithProperties()) 1204 require.NoError(t, err) 1205 for _, levelTables := range tableInfos { 1206 for _, info := range levelTables { 1207 require.NotNil(t, info.Properties) 1208 } 1209 } 1210 } 1211 1212 func BenchmarkDelete(b *testing.B) { 1213 rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano()))) 1214 const keyCount = 10000 1215 var keys [keyCount][]byte 1216 for i := 0; i < keyCount; i++ { 1217 keys[i] = []byte(strconv.Itoa(rng.Int())) 1218 } 1219 val := bytes.Repeat([]byte("x"), 10) 1220 1221 benchmark := func(b *testing.B, useSingleDelete bool) { 1222 d, err := Open( 1223 "", 1224 &Options{ 1225 FS: vfs.NewMem(), 1226 }) 1227 if err != nil { 1228 b.Fatal(err) 1229 } 1230 defer func() { 1231 if err := d.Close(); err != nil { 1232 b.Fatal(err) 1233 } 1234 }() 1235 1236 b.StartTimer() 1237 for _, key := range keys { 1238 _ = d.Set(key, val, nil) 1239 if useSingleDelete { 1240 _ = d.SingleDelete(key, nil) 1241 } else { 1242 _ = d.Delete(key, nil) 1243 } 1244 } 1245 // Manually flush as it is flushing/compaction where SingleDelete 1246 // performance shows up. With SingleDelete, we can elide all of the 1247 // SingleDelete and Set records. 1248 if err := d.Flush(); err != nil { 1249 b.Fatal(err) 1250 } 1251 b.StopTimer() 1252 } 1253 1254 b.Run("delete", func(b *testing.B) { 1255 for i := 0; i < b.N; i++ { 1256 benchmark(b, false) 1257 } 1258 }) 1259 1260 b.Run("single-delete", func(b *testing.B) { 1261 for i := 0; i < b.N; i++ { 1262 benchmark(b, true) 1263 } 1264 }) 1265 } 1266 1267 func BenchmarkNewIterReadAmp(b *testing.B) { 1268 for _, readAmp := range []int{10, 100, 1000} { 1269 b.Run(strconv.Itoa(readAmp), func(b *testing.B) { 1270 opts := &Options{ 1271 FS: vfs.NewMem(), 1272 L0StopWritesThreshold: 1000, 1273 } 1274 opts.DisableAutomaticCompactions = true 1275 1276 d, err := Open("", opts) 1277 require.NoError(b, err) 1278 1279 for i := 0; i < readAmp; i++ { 1280 require.NoError(b, d.Set([]byte("a"), []byte("b"), NoSync)) 1281 require.NoError(b, d.Flush()) 1282 } 1283 1284 require.Equal(b, d.Metrics().ReadAmp(), readAmp) 1285 1286 b.StopTimer() 1287 b.ResetTimer() 1288 for i := 0; i < b.N; i++ { 1289 b.StartTimer() 1290 iter := d.NewIter(nil) 1291 b.StopTimer() 1292 require.NoError(b, iter.Close()) 1293 } 1294 1295 require.NoError(b, d.Close()) 1296 }) 1297 } 1298 } 1299 1300 func verifyGet(t *testing.T, r Reader, key, expected []byte) { 1301 val, closer, err := r.Get(key) 1302 require.NoError(t, err) 1303 if !bytes.Equal(expected, val) { 1304 t.Fatalf("expected %s, but got %s", expected, val) 1305 } 1306 closer.Close() 1307 } 1308 1309 func verifyGetNotFound(t *testing.T, r Reader, key []byte) { 1310 val, _, err := r.Get(key) 1311 if err != base.ErrNotFound { 1312 t.Fatalf("expected nil, but got %s", val) 1313 } 1314 }