github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/batch_test.go (about) 1 // Copyright 2012 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 "encoding/binary" 10 "encoding/hex" 11 "fmt" 12 "io" 13 "math" 14 "math/rand" 15 "strconv" 16 "strings" 17 "sync" 18 "testing" 19 "time" 20 "unicode" 21 22 "github.com/cockroachdb/datadriven" 23 "github.com/cockroachdb/errors" 24 "github.com/cockroachdb/pebble/internal/base" 25 "github.com/cockroachdb/pebble/internal/batchskl" 26 "github.com/cockroachdb/pebble/internal/keyspan" 27 "github.com/cockroachdb/pebble/internal/testkeys" 28 "github.com/cockroachdb/pebble/vfs" 29 "github.com/stretchr/testify/require" 30 ) 31 32 func TestBatch(t *testing.T) { 33 testBatch(t, 0) 34 testBatch(t, batchInitialSize) 35 } 36 37 func testBatch(t *testing.T, size int) { 38 type testCase struct { 39 kind InternalKeyKind 40 key, value string 41 valueInt uint32 42 } 43 44 verifyTestCases := func(b *Batch, testCases []testCase, indexedPointKindsOnly bool) { 45 r := b.Reader() 46 47 for _, tc := range testCases { 48 if indexedPointKindsOnly && (tc.kind == InternalKeyKindLogData || tc.kind == InternalKeyKindIngestSST || 49 tc.kind == InternalKeyKindRangeKeyUnset || tc.kind == InternalKeyKindRangeKeySet || 50 tc.kind == InternalKeyKindRangeKeyDelete || tc.kind == InternalKeyKindRangeDelete) { 51 continue 52 } 53 kind, k, v, ok, err := r.Next() 54 if !ok { 55 if err != nil { 56 t.Fatal(err) 57 } 58 t.Fatalf("next returned !ok: test case = %v", tc) 59 } 60 key, value := string(k), string(v) 61 if kind != tc.kind || key != tc.key || value != tc.value { 62 t.Errorf("got (%d, %q, %q), want (%d, %q, %q)", 63 kind, key, value, tc.kind, tc.key, tc.value) 64 } 65 } 66 if len(r) != 0 { 67 t.Errorf("reader was not exhausted: remaining bytes = %q", r) 68 } 69 } 70 71 encodeFileNum := func(n base.FileNum) string { 72 return string(binary.AppendUvarint(nil, uint64(n))) 73 } 74 decodeFileNum := func(d []byte) base.FileNum { 75 val, n := binary.Uvarint(d) 76 if n <= 0 { 77 t.Fatalf("invalid filenum encoding") 78 } 79 return base.FileNum(val) 80 } 81 82 // RangeKeySet and RangeKeyUnset are untested here because they don't expose 83 // deferred variants. This is a consequence of these keys' more complex 84 // value encodings. 85 testCases := []testCase{ 86 {InternalKeyKindIngestSST, encodeFileNum(1), "", 0}, 87 {InternalKeyKindSet, "roses", "red", 0}, 88 {InternalKeyKindSet, "violets", "blue", 0}, 89 {InternalKeyKindDelete, "roses", "", 0}, 90 {InternalKeyKindSingleDelete, "roses", "", 0}, 91 {InternalKeyKindSet, "", "", 0}, 92 {InternalKeyKindSet, "", "non-empty", 0}, 93 {InternalKeyKindDelete, "", "", 0}, 94 {InternalKeyKindSingleDelete, "", "", 0}, 95 {InternalKeyKindSet, "grass", "green", 0}, 96 {InternalKeyKindSet, "grass", "greener", 0}, 97 {InternalKeyKindSet, "eleventy", strings.Repeat("!!11!", 100), 0}, 98 {InternalKeyKindDelete, "nosuchkey", "", 0}, 99 {InternalKeyKindDeleteSized, "eleventy", string(binary.AppendUvarint([]byte(nil), 508)), 500}, 100 {InternalKeyKindSingleDelete, "nosuchkey", "", 0}, 101 {InternalKeyKindSet, "binarydata", "\x00", 0}, 102 {InternalKeyKindSet, "binarydata", "\xff", 0}, 103 {InternalKeyKindMerge, "merge", "mergedata", 0}, 104 {InternalKeyKindMerge, "merge", "", 0}, 105 {InternalKeyKindMerge, "", "", 0}, 106 {InternalKeyKindRangeDelete, "a", "b", 0}, 107 {InternalKeyKindRangeDelete, "", "", 0}, 108 {InternalKeyKindLogData, "logdata", "", 0}, 109 {InternalKeyKindLogData, "", "", 0}, 110 {InternalKeyKindRangeKeyDelete, "grass", "green", 0}, 111 {InternalKeyKindRangeKeyDelete, "", "", 0}, 112 {InternalKeyKindDeleteSized, "nosuchkey", string(binary.AppendUvarint([]byte(nil), 11)), 2}, 113 } 114 b := newBatchWithSize(nil, size) 115 for _, tc := range testCases { 116 switch tc.kind { 117 case InternalKeyKindSet: 118 _ = b.Set([]byte(tc.key), []byte(tc.value), nil) 119 case InternalKeyKindMerge: 120 _ = b.Merge([]byte(tc.key), []byte(tc.value), nil) 121 case InternalKeyKindDelete: 122 _ = b.Delete([]byte(tc.key), nil) 123 case InternalKeyKindDeleteSized: 124 _ = b.DeleteSized([]byte(tc.key), tc.valueInt, nil) 125 case InternalKeyKindSingleDelete: 126 _ = b.SingleDelete([]byte(tc.key), nil) 127 case InternalKeyKindRangeDelete: 128 _ = b.DeleteRange([]byte(tc.key), []byte(tc.value), nil) 129 case InternalKeyKindLogData: 130 _ = b.LogData([]byte(tc.key), nil) 131 case InternalKeyKindRangeKeyDelete: 132 _ = b.RangeKeyDelete([]byte(tc.key), []byte(tc.value), nil) 133 case InternalKeyKindIngestSST: 134 b.ingestSST(decodeFileNum([]byte(tc.key))) 135 } 136 } 137 verifyTestCases(b, testCases, false /* indexedKindsOnly */) 138 139 b.Reset() 140 // Run the same operations, this time using the Deferred variants of each 141 // operation (eg. SetDeferred). 142 for _, tc := range testCases { 143 key := []byte(tc.key) 144 value := []byte(tc.value) 145 switch tc.kind { 146 case InternalKeyKindSet: 147 d := b.SetDeferred(len(key), len(value)) 148 copy(d.Key, key) 149 copy(d.Value, value) 150 d.Finish() 151 case InternalKeyKindMerge: 152 d := b.MergeDeferred(len(key), len(value)) 153 copy(d.Key, key) 154 copy(d.Value, value) 155 d.Finish() 156 case InternalKeyKindDelete: 157 d := b.DeleteDeferred(len(key)) 158 copy(d.Key, key) 159 copy(d.Value, value) 160 d.Finish() 161 case InternalKeyKindDeleteSized: 162 d := b.DeleteSizedDeferred(len(tc.key), tc.valueInt) 163 copy(d.Key, key) 164 d.Finish() 165 case InternalKeyKindSingleDelete: 166 d := b.SingleDeleteDeferred(len(key)) 167 copy(d.Key, key) 168 copy(d.Value, value) 169 d.Finish() 170 case InternalKeyKindRangeDelete: 171 d := b.DeleteRangeDeferred(len(key), len(value)) 172 copy(d.Key, key) 173 copy(d.Value, value) 174 d.Finish() 175 case InternalKeyKindLogData: 176 _ = b.LogData([]byte(tc.key), nil) 177 case InternalKeyKindIngestSST: 178 b.ingestSST(decodeFileNum([]byte(tc.key))) 179 case InternalKeyKindRangeKeyDelete: 180 d := b.RangeKeyDeleteDeferred(len(key), len(value)) 181 copy(d.Key, key) 182 copy(d.Value, value) 183 d.Finish() 184 } 185 } 186 verifyTestCases(b, testCases, false /* indexedKindsOnly */) 187 188 b.Reset() 189 // Run the same operations, this time using AddInternalKey instead of the 190 // Kind-specific methods. 191 for _, tc := range testCases { 192 if tc.kind == InternalKeyKindLogData || tc.kind == InternalKeyKindIngestSST || 193 tc.kind == InternalKeyKindRangeKeyUnset || tc.kind == InternalKeyKindRangeKeySet || 194 tc.kind == InternalKeyKindRangeKeyDelete || tc.kind == InternalKeyKindRangeDelete { 195 continue 196 } 197 key := []byte(tc.key) 198 value := []byte(tc.value) 199 b.AddInternalKey(&InternalKey{UserKey: key, Trailer: base.MakeTrailer(0, tc.kind)}, value, nil) 200 } 201 verifyTestCases(b, testCases, true /* indexedKindsOnly */) 202 } 203 204 func TestBatchPreAlloc(t *testing.T) { 205 var cases = []struct { 206 size int 207 exp int 208 }{ 209 {0, batchInitialSize}, 210 {batchInitialSize, batchInitialSize}, 211 {2 * batchInitialSize, 2 * batchInitialSize}, 212 } 213 for _, c := range cases { 214 b := newBatchWithSize(nil, c.size) 215 b.Set([]byte{0x1}, []byte{0x2}, nil) 216 if cap(b.data) != c.exp { 217 t.Errorf("Unexpected memory space, required: %d, got: %d", c.exp, cap(b.data)) 218 } 219 } 220 } 221 222 func TestBatchIngestSST(t *testing.T) { 223 // Verify that Batch.IngestSST has the correct batch count and memtable 224 // size. 225 var b Batch 226 b.ingestSST(1) 227 require.Equal(t, int(b.Count()), 1) 228 b.ingestSST(2) 229 require.Equal(t, int(b.Count()), 2) 230 require.Equal(t, int(b.memTableSize), 0) 231 require.Equal(t, b.ingestedSSTBatch, true) 232 } 233 234 func TestBatchLen(t *testing.T) { 235 var b Batch 236 237 requireLenAndReprEq := func(size int) { 238 require.Equal(t, size, b.Len()) 239 require.Equal(t, size, len(b.Repr())) 240 } 241 242 requireLenAndReprEq(batchHeaderLen) 243 244 key := "test-key" 245 value := "test-value" 246 247 err := b.Set([]byte(key), []byte(value), nil) 248 require.NoError(t, err) 249 250 requireLenAndReprEq(33) 251 252 err = b.Delete([]byte(key), nil) 253 require.NoError(t, err) 254 255 requireLenAndReprEq(43) 256 } 257 258 func TestBatchEmpty(t *testing.T) { 259 testBatchEmpty(t, 0) 260 testBatchEmpty(t, batchInitialSize) 261 } 262 263 func testBatchEmpty(t *testing.T, size int) { 264 b := newBatchWithSize(nil, size) 265 require.True(t, b.Empty()) 266 267 ops := []func(*Batch) error{ 268 func(b *Batch) error { return b.Set(nil, nil, nil) }, 269 func(b *Batch) error { return b.Merge(nil, nil, nil) }, 270 func(b *Batch) error { return b.Delete(nil, nil) }, 271 func(b *Batch) error { return b.DeleteRange(nil, nil, nil) }, 272 func(b *Batch) error { return b.LogData(nil, nil) }, 273 func(b *Batch) error { return b.RangeKeySet(nil, nil, nil, nil, nil) }, 274 func(b *Batch) error { return b.RangeKeyUnset(nil, nil, nil, nil) }, 275 func(b *Batch) error { return b.RangeKeyDelete(nil, nil, nil) }, 276 } 277 278 for _, op := range ops { 279 require.NoError(t, op(b)) 280 require.False(t, b.Empty()) 281 b.Reset() 282 require.True(t, b.Empty()) 283 // Reset may choose to reuse b.data, so clear it to the zero value in 284 // order to test the lazy initialization of b.data. 285 b = newBatchWithSize(nil, size) 286 } 287 288 _ = b.Reader() 289 require.True(t, b.Empty()) 290 b.Reset() 291 require.True(t, b.Empty()) 292 b = newBatchWithSize(nil, size) 293 294 require.Equal(t, uint64(0), b.SeqNum()) 295 require.True(t, b.Empty()) 296 b.Reset() 297 require.True(t, b.Empty()) 298 b = &Batch{} 299 300 d, err := Open("", &Options{ 301 FS: vfs.NewMem(), 302 }) 303 require.NoError(t, err) 304 defer d.Close() 305 ib := newIndexedBatch(d, DefaultComparer) 306 iter, _ := ib.NewIter(nil) 307 require.False(t, iter.First()) 308 iter2, err := iter.Clone(CloneOptions{}) 309 require.NoError(t, err) 310 require.NoError(t, iter.Close()) 311 _, err = iter.Clone(CloneOptions{}) 312 require.True(t, err != nil) 313 require.False(t, iter2.First()) 314 require.NoError(t, iter2.Close()) 315 } 316 317 func TestBatchApplyNoSyncWait(t *testing.T) { 318 db, err := Open("", &Options{ 319 FS: vfs.NewMem(), 320 }) 321 require.NoError(t, err) 322 defer db.Close() 323 var batches []*Batch 324 options := &WriteOptions{Sync: true} 325 for i := 0; i < 10000; i++ { 326 b := db.NewBatch() 327 str := fmt.Sprintf("a%d", i) 328 require.NoError(t, b.Set([]byte(str), []byte(str), nil)) 329 require.NoError(t, db.ApplyNoSyncWait(b, options)) 330 // k-v pair is visible even if not yet synced. 331 val, closer, err := db.Get([]byte(str)) 332 require.NoError(t, err) 333 require.Equal(t, str, string(val)) 334 closer.Close() 335 batches = append(batches, b) 336 } 337 for _, b := range batches { 338 require.NoError(t, b.SyncWait()) 339 b.Close() 340 } 341 } 342 343 func TestBatchReset(t *testing.T) { 344 db, err := Open("", &Options{ 345 FS: vfs.NewMem(), 346 }) 347 require.NoError(t, err) 348 defer db.Close() 349 key := "test-key" 350 value := "test-value" 351 b := db.NewBatch() 352 require.NoError(t, b.Set([]byte(key), []byte(value), nil)) 353 dd := b.DeleteRangeDeferred(len(key), len(value)) 354 copy(dd.Key, key) 355 copy(dd.Value, value) 356 dd.Finish() 357 358 require.NoError(t, b.RangeKeySet([]byte(key), []byte(value), []byte(value), []byte(value), nil)) 359 360 b.setSeqNum(100) 361 b.applied.Store(true) 362 b.commitErr = errors.New("test-error") 363 b.commit.Add(1) 364 b.fsyncWait.Add(1) 365 require.Equal(t, uint32(3), b.Count()) 366 require.Equal(t, uint64(1), b.countRangeDels) 367 require.Equal(t, uint64(1), b.countRangeKeys) 368 require.True(t, len(b.data) > 0) 369 require.True(t, b.SeqNum() > 0) 370 require.True(t, b.memTableSize > 0) 371 require.NotEqual(t, b.deferredOp, DeferredBatchOp{}) 372 // At this point b.data has not been modified since the db.NewBatch() and is 373 // either nil or contains a byte slice of length batchHeaderLen, with a 0 374 // seqnum encoded in data[0:8] and an arbitrary count encoded in data[8:12]. 375 // The following commented code will often fail. 376 // count := binary.LittleEndian.Uint32(b.countData()) 377 // if count != 0 && count != 3 { 378 // t.Fatalf("count: %d", count) 379 // } 380 // If we simply called b.Reset now and later used b.data to initialize 381 // expected, the count in expected will also be arbitrary. So we fix the 382 // count in b.data now by calling b.Repr(). This call isn't essential, since 383 // we will call b.Repr() again, and just shows that it fixes the count in 384 // b.data. 385 _ = b.Repr() 386 require.Equal(t, uint32(3), binary.LittleEndian.Uint32(b.countData())) 387 388 b.Reset() 389 require.Equal(t, db, b.db) 390 require.Equal(t, false, b.applied.Load()) 391 require.Nil(t, b.commitErr) 392 require.Equal(t, uint32(0), b.Count()) 393 require.Equal(t, uint64(0), b.countRangeDels) 394 require.Equal(t, uint64(0), b.countRangeKeys) 395 require.Equal(t, batchHeaderLen, len(b.data)) 396 require.Equal(t, uint64(0), b.SeqNum()) 397 require.Equal(t, uint64(0), b.memTableSize) 398 require.Equal(t, FormatMajorVersion(0x00), b.minimumFormatMajorVersion) 399 require.Equal(t, b.deferredOp, DeferredBatchOp{}) 400 _ = b.Repr() 401 402 var expected Batch 403 require.NoError(t, expected.SetRepr(b.data)) 404 expected.db = db 405 require.Equal(t, &expected, b) 406 407 // Reset batch can be used to write and commit a new record. 408 b.Set([]byte(key), []byte(value), nil) 409 require.NoError(t, db.Apply(b, nil)) 410 v, closer, err := db.Get([]byte(key)) 411 require.NoError(t, err) 412 defer closer.Close() 413 require.Equal(t, v, []byte(value)) 414 } 415 416 func TestIndexedBatchReset(t *testing.T) { 417 indexCount := func(sl *batchskl.Skiplist) int { 418 count := 0 419 iter := sl.NewIter(nil, nil) 420 defer iter.Close() 421 for iter.First(); iter.Valid(); iter.Next() { 422 count++ 423 } 424 return count 425 } 426 db, err := Open("", &Options{ 427 FS: vfs.NewMem(), 428 }) 429 require.NoError(t, err) 430 defer db.Close() 431 b := newIndexedBatch(db, DefaultComparer) 432 start := "start-key" 433 end := "end-key" 434 key := "test-key" 435 value := "test-value" 436 b.DeleteRange([]byte(start), []byte(end), nil) 437 b.Set([]byte(key), []byte(value), nil) 438 require.NoError(t, b. 439 RangeKeySet([]byte(start), []byte(end), []byte("suffix"), []byte(value), nil)) 440 require.NotNil(t, b.rangeKeyIndex) 441 require.NotNil(t, b.rangeDelIndex) 442 require.NotNil(t, b.index) 443 require.Equal(t, 1, indexCount(b.index)) 444 445 b.Reset() 446 require.NotNil(t, b.cmp) 447 require.NotNil(t, b.formatKey) 448 require.NotNil(t, b.abbreviatedKey) 449 require.NotNil(t, b.index) 450 require.Nil(t, b.rangeDelIndex) 451 require.Nil(t, b.rangeKeyIndex) 452 453 count := func(ib *Batch) int { 454 iter, _ := ib.NewIter(nil) 455 defer iter.Close() 456 iter2, err := iter.Clone(CloneOptions{}) 457 require.NoError(t, err) 458 defer iter2.Close() 459 var count [2]int 460 for i, it := range []*Iterator{iter, iter2} { 461 for it.First(); it.Valid(); it.Next() { 462 count[i]++ 463 } 464 } 465 require.Equal(t, count[0], count[1]) 466 return count[0] 467 } 468 contains := func(ib *Batch, key, value string) bool { 469 iter, _ := ib.NewIter(nil) 470 defer iter.Close() 471 iter2, err := iter.Clone(CloneOptions{}) 472 require.NoError(t, err) 473 defer iter2.Close() 474 var found [2]bool 475 for i, it := range []*Iterator{iter, iter2} { 476 for it.First(); it.Valid(); it.Next() { 477 if string(it.Key()) == key && 478 string(it.Value()) == value { 479 found[i] = true 480 } 481 } 482 } 483 require.Equal(t, found[0], found[1]) 484 return found[0] 485 } 486 // Set a key and check whether the key-value pair is visible. 487 b.Set([]byte(key), []byte(value), nil) 488 require.Equal(t, 1, indexCount(b.index)) 489 require.Equal(t, 1, count(b)) 490 require.True(t, contains(b, key, value)) 491 492 // Use range delete to delete the above inserted key-value pair. 493 b.DeleteRange([]byte(key), []byte(value), nil) 494 require.NotNil(t, b.rangeDelIndex) 495 require.Equal(t, 1, indexCount(b.rangeDelIndex)) 496 require.Equal(t, 0, count(b)) 497 require.False(t, contains(b, key, value)) 498 } 499 500 // TestIndexedBatchMutation tests mutating an indexed batch with an open 501 // iterator. 502 func TestIndexedBatchMutation(t *testing.T) { 503 opts := &Options{ 504 Comparer: testkeys.Comparer, 505 FS: vfs.NewMem(), 506 FormatMajorVersion: internalFormatNewest, 507 } 508 d, err := Open("", opts) 509 require.NoError(t, err) 510 defer func() { d.Close() }() 511 512 b := newIndexedBatch(d, DefaultComparer) 513 iters := map[string]*Iterator{} 514 defer func() { 515 for _, iter := range iters { 516 require.NoError(t, iter.Close()) 517 } 518 }() 519 520 datadriven.RunTest(t, "testdata/indexed_batch_mutation", func(t *testing.T, td *datadriven.TestData) string { 521 switch td.Cmd { 522 case "batch": 523 writeBatch := newBatch(d) 524 if err := runBatchDefineCmd(td, writeBatch); err != nil { 525 return err.Error() 526 } 527 if err := writeBatch.Commit(nil); err != nil { 528 return err.Error() 529 } 530 return "" 531 case "new-batch-iter": 532 name := td.CmdArgs[0].String() 533 iters[name], _ = b.NewIter(&IterOptions{ 534 KeyTypes: IterKeyTypePointsAndRanges, 535 }) 536 return "" 537 case "new-db-iter": 538 name := td.CmdArgs[0].String() 539 iters[name], _ = d.NewIter(&IterOptions{ 540 KeyTypes: IterKeyTypePointsAndRanges, 541 }) 542 return "" 543 case "new-batch": 544 if b != nil { 545 require.NoError(t, b.Close()) 546 } 547 b = newIndexedBatch(d, opts.Comparer) 548 if err := runBatchDefineCmd(td, b); err != nil { 549 return err.Error() 550 } 551 return "" 552 case "flush": 553 require.NoError(t, d.Flush()) 554 return "" 555 case "iter": 556 var iter string 557 td.ScanArgs(t, "iter", &iter) 558 return runIterCmd(td, iters[iter], false /* closeIter */) 559 case "mutate": 560 mut := newBatch(d) 561 if err := runBatchDefineCmd(td, mut); err != nil { 562 return err.Error() 563 } 564 if err := b.Apply(mut, nil); err != nil { 565 return err.Error() 566 } 567 return "" 568 case "clone": 569 var from, to string 570 var refreshBatchView bool 571 td.ScanArgs(t, "from", &from) 572 td.ScanArgs(t, "to", &to) 573 td.ScanArgs(t, "refresh-batch", &refreshBatchView) 574 var err error 575 iters[to], err = iters[from].Clone(CloneOptions{RefreshBatchView: refreshBatchView}) 576 if err != nil { 577 return err.Error() 578 } 579 return "" 580 case "reset": 581 for key, iter := range iters { 582 if err := iter.Close(); err != nil { 583 return err.Error() 584 } 585 delete(iters, key) 586 } 587 if d != nil { 588 if err := d.Close(); err != nil { 589 return err.Error() 590 } 591 } 592 opts.FS = vfs.NewMem() 593 d, err = Open("", opts) 594 require.NoError(t, err) 595 return "" 596 default: 597 return fmt.Sprintf("unrecognized command %q", td.Cmd) 598 } 599 }) 600 } 601 602 func TestIndexedBatch_GlobalVisibility(t *testing.T) { 603 opts := &Options{ 604 FS: vfs.NewMem(), 605 FormatMajorVersion: internalFormatNewest, 606 Comparer: testkeys.Comparer, 607 } 608 d, err := Open("", opts) 609 require.NoError(t, err) 610 defer d.Close() 611 612 require.NoError(t, d.Set([]byte("foo"), []byte("foo"), nil)) 613 614 // Create an iterator over an empty indexed batch. 615 b := newIndexedBatch(d, DefaultComparer) 616 iterOpts := IterOptions{KeyTypes: IterKeyTypePointsAndRanges} 617 iter, _ := b.NewIter(&iterOpts) 618 defer iter.Close() 619 620 // Mutate the database's committed state. 621 mut := newBatch(d) 622 require.NoError(t, mut.Set([]byte("bar"), []byte("bar"), nil)) 623 require.NoError(t, mut.DeleteRange([]byte("e"), []byte("g"), nil)) 624 require.NoError(t, mut.RangeKeySet([]byte("a"), []byte("c"), []byte("@1"), []byte("v"), nil)) 625 require.NoError(t, mut.Commit(nil)) 626 627 scanIter := func() string { 628 var buf bytes.Buffer 629 for valid := iter.First(); valid; valid = iter.Next() { 630 fmt.Fprintf(&buf, "%s: (", iter.Key()) 631 hasPoint, hasRange := iter.HasPointAndRange() 632 if hasPoint { 633 fmt.Fprintf(&buf, "%s,", iter.Value()) 634 } else { 635 fmt.Fprintf(&buf, ".,") 636 } 637 if hasRange { 638 start, end := iter.RangeBounds() 639 fmt.Fprintf(&buf, "[%s-%s)", start, end) 640 writeRangeKeys(&buf, iter) 641 } else { 642 fmt.Fprintf(&buf, ".") 643 } 644 fmt.Fprintln(&buf, ")") 645 } 646 return strings.TrimSpace(buf.String()) 647 } 648 // Scanning the iterator should only see the point key written before the 649 // iterator was constructed. 650 require.Equal(t, `foo: (foo,.)`, scanIter()) 651 652 // After calling SetOptions, the iterator should still only see the point 653 // key written before the iterator was constructed. SetOptions refreshes the 654 // iterator's view of its own indexed batch, but not committed state. 655 iter.SetOptions(&iterOpts) 656 require.Equal(t, `foo: (foo,.)`, scanIter()) 657 } 658 659 func TestFlushableBatchReset(t *testing.T) { 660 var b Batch 661 var err error 662 b.flushable, err = newFlushableBatch(&b, DefaultComparer) 663 require.NoError(t, err) 664 665 b.Reset() 666 require.Nil(t, b.flushable) 667 } 668 669 func TestBatchIncrement(t *testing.T) { 670 testCases := []uint32{ 671 0x00000000, 672 0x00000001, 673 0x00000002, 674 0x0000007f, 675 0x00000080, 676 0x000000fe, 677 0x000000ff, 678 0x00000100, 679 0x00000101, 680 0x000001ff, 681 0x00000200, 682 0x00000fff, 683 0x00001234, 684 0x0000fffe, 685 0x0000ffff, 686 0x00010000, 687 0x00010001, 688 0x000100fe, 689 0x000100ff, 690 0x00020100, 691 0x03fffffe, 692 0x03ffffff, 693 0x04000000, 694 0x04000001, 695 0x7fffffff, 696 0xfffffffe, 697 } 698 for _, tc := range testCases { 699 var buf [batchHeaderLen]byte 700 binary.LittleEndian.PutUint32(buf[8:12], tc) 701 var b Batch 702 b.SetRepr(buf[:]) 703 b.count++ 704 got := binary.LittleEndian.Uint32(b.Repr()[8:12]) 705 want := tc + 1 706 if got != want { 707 t.Errorf("input=%d: got %d, want %d", tc, got, want) 708 } 709 _, count := ReadBatch(b.Repr()) 710 if got != want { 711 t.Errorf("input=%d: got %d, want %d", tc, count, want) 712 } 713 } 714 715 err := func() (err error) { 716 defer func() { 717 if v := recover(); v != nil { 718 if verr, ok := v.(error); ok { 719 err = verr 720 } 721 } 722 }() 723 var buf [batchHeaderLen]byte 724 binary.LittleEndian.PutUint32(buf[8:12], 0xffffffff) 725 var b Batch 726 b.SetRepr(buf[:]) 727 b.count++ 728 b.Repr() 729 return nil 730 }() 731 if err != ErrInvalidBatch { 732 t.Fatalf("expected %v, but found %v", ErrInvalidBatch, err) 733 } 734 } 735 736 func TestBatchOpDoesIncrement(t *testing.T) { 737 var b Batch 738 key := []byte("foo") 739 value := []byte("bar") 740 741 if b.Count() != 0 { 742 t.Fatalf("new batch has a nonzero count: %d", b.Count()) 743 } 744 745 // Should increment count by 1 746 _ = b.Set(key, value, nil) 747 if b.Count() != 1 { 748 t.Fatalf("expected count: %d, got %d", 1, b.Count()) 749 } 750 751 var b2 Batch 752 // Should increment count by 1 each 753 _ = b2.Set(key, value, nil) 754 _ = b2.Delete(key, nil) 755 if b2.Count() != 2 { 756 t.Fatalf("expected count: %d, got %d", 2, b2.Count()) 757 } 758 759 // Should increment count by b2.count() 760 _ = b.Apply(&b2, nil) 761 if b.Count() != 3 { 762 t.Fatalf("expected count: %d, got %d", 3, b.Count()) 763 } 764 765 // Should increment count by 1 766 _ = b.Merge(key, value, nil) 767 if b.Count() != 4 { 768 t.Fatalf("expected count: %d, got %d", 4, b.Count()) 769 } 770 771 // Should NOT increment count. 772 _ = b.LogData([]byte("foobarbaz"), nil) 773 if b.Count() != 4 { 774 t.Fatalf("expected count: %d, got %d", 4, b.Count()) 775 } 776 } 777 778 func TestBatchGet(t *testing.T) { 779 testCases := []struct { 780 method string 781 memTableSize uint64 782 }{ 783 {"build", 64 << 20}, 784 {"build", 2 << 10}, 785 {"apply", 64 << 20}, 786 } 787 788 for _, c := range testCases { 789 t.Run(fmt.Sprintf("%s,mem=%d", c.method, c.memTableSize), func(t *testing.T) { 790 d, err := Open("", &Options{ 791 FS: vfs.NewMem(), 792 MemTableSize: c.memTableSize, 793 }) 794 if err != nil { 795 t.Fatalf("Open: %v", err) 796 } 797 defer d.Close() 798 var b *Batch 799 800 datadriven.RunTest(t, "testdata/batch_get", func(t *testing.T, td *datadriven.TestData) string { 801 switch td.Cmd { 802 case "define": 803 switch c.method { 804 case "build": 805 b = d.NewIndexedBatch() 806 case "apply": 807 b = d.NewBatch() 808 } 809 810 if err := runBatchDefineCmd(td, b); err != nil { 811 return err.Error() 812 } 813 814 switch c.method { 815 case "apply": 816 tmp := d.NewIndexedBatch() 817 tmp.Apply(b, nil) 818 b = tmp 819 } 820 return "" 821 822 case "commit": 823 if err := b.Commit(nil); err != nil { 824 return err.Error() 825 } 826 b = nil 827 return "" 828 829 case "get": 830 if len(td.CmdArgs) != 1 { 831 return fmt.Sprintf("%s expects 1 argument", td.Cmd) 832 } 833 v, closer, err := b.Get([]byte(td.CmdArgs[0].String())) 834 if err != nil { 835 return err.Error() 836 } 837 defer closer.Close() 838 return string(v) 839 840 default: 841 return fmt.Sprintf("unknown command: %s", td.Cmd) 842 } 843 }) 844 }) 845 } 846 } 847 848 func TestBatchIter(t *testing.T) { 849 var b *Batch 850 851 for _, method := range []string{"build", "apply"} { 852 for _, testdata := range []string{ 853 "testdata/internal_iter_next", "testdata/internal_iter_bounds"} { 854 t.Run(method, func(t *testing.T) { 855 datadriven.RunTest(t, testdata, func(t *testing.T, d *datadriven.TestData) string { 856 switch d.Cmd { 857 case "define": 858 switch method { 859 case "build": 860 b = newIndexedBatch(nil, DefaultComparer) 861 case "apply": 862 b = newBatch(nil) 863 } 864 865 for _, key := range strings.Split(d.Input, "\n") { 866 j := strings.Index(key, ":") 867 ikey := base.ParseInternalKey(key[:j]) 868 value := []byte(key[j+1:]) 869 b.Set(ikey.UserKey, value, nil) 870 } 871 872 switch method { 873 case "apply": 874 tmp := newIndexedBatch(nil, DefaultComparer) 875 tmp.Apply(b, nil) 876 b = tmp 877 } 878 return "" 879 880 case "iter": 881 var options IterOptions 882 for _, arg := range d.CmdArgs { 883 switch arg.Key { 884 case "lower": 885 if len(arg.Vals) != 1 { 886 return fmt.Sprintf( 887 "%s expects at most 1 value for lower", d.Cmd) 888 } 889 options.LowerBound = []byte(arg.Vals[0]) 890 case "upper": 891 if len(arg.Vals) != 1 { 892 return fmt.Sprintf( 893 "%s expects at most 1 value for upper", d.Cmd) 894 } 895 options.UpperBound = []byte(arg.Vals[0]) 896 default: 897 return fmt.Sprintf("unknown arg: %s", arg.Key) 898 } 899 } 900 iter := b.newInternalIter(&options) 901 defer iter.Close() 902 return runInternalIterCmd(t, d, iter) 903 904 default: 905 return fmt.Sprintf("unknown command: %s", d.Cmd) 906 } 907 }) 908 }) 909 } 910 } 911 } 912 913 func TestBatchRangeOps(t *testing.T) { 914 var b *Batch 915 916 datadriven.RunTest(t, "testdata/batch_range_ops", func(t *testing.T, td *datadriven.TestData) string { 917 switch td.Cmd { 918 case "clear": 919 b = nil 920 return "" 921 922 case "apply": 923 if b == nil { 924 b = newIndexedBatch(nil, DefaultComparer) 925 } 926 t := newBatch(nil) 927 if err := runBatchDefineCmd(td, t); err != nil { 928 return err.Error() 929 } 930 if err := b.Apply(t, nil); err != nil { 931 return err.Error() 932 } 933 return "" 934 935 case "define": 936 if b == nil { 937 b = newIndexedBatch(nil, DefaultComparer) 938 } 939 if err := runBatchDefineCmd(td, b); err != nil { 940 return err.Error() 941 } 942 return "" 943 944 case "scan": 945 if len(td.CmdArgs) > 1 { 946 return fmt.Sprintf("%s expects at most 1 argument", td.Cmd) 947 } 948 var fragmentIter keyspan.FragmentIterator 949 var internalIter base.InternalIterator 950 switch { 951 case td.HasArg("range-del"): 952 fragmentIter = b.newRangeDelIter(nil, math.MaxUint64) 953 defer fragmentIter.Close() 954 case td.HasArg("range-key"): 955 fragmentIter = b.newRangeKeyIter(nil, math.MaxUint64) 956 defer fragmentIter.Close() 957 default: 958 internalIter = b.newInternalIter(nil) 959 defer internalIter.Close() 960 } 961 962 var buf bytes.Buffer 963 if fragmentIter != nil { 964 for s := fragmentIter.First(); s != nil; s = fragmentIter.Next() { 965 for i := range s.Keys { 966 s.Keys[i].Trailer = base.MakeTrailer( 967 s.Keys[i].SeqNum()&^base.InternalKeySeqNumBatch, 968 s.Keys[i].Kind(), 969 ) 970 } 971 fmt.Fprintln(&buf, s) 972 } 973 } else { 974 for k, v := internalIter.First(); k != nil; k, v = internalIter.Next() { 975 k.SetSeqNum(k.SeqNum() &^ InternalKeySeqNumBatch) 976 fmt.Fprintf(&buf, "%s:%s\n", k, v.InPlaceValue()) 977 } 978 } 979 return buf.String() 980 981 default: 982 return fmt.Sprintf("unknown command: %s", td.Cmd) 983 } 984 }) 985 } 986 987 func TestBatchTooLarge(t *testing.T) { 988 var b Batch 989 var result interface{} 990 func() { 991 defer func() { 992 if r := recover(); r != nil { 993 result = r 994 } 995 }() 996 b.grow(maxBatchSize) 997 }() 998 require.EqualValues(t, ErrBatchTooLarge, result) 999 } 1000 1001 func TestFlushableBatchIter(t *testing.T) { 1002 var b *flushableBatch 1003 datadriven.RunTest(t, "testdata/internal_iter_next", func(t *testing.T, d *datadriven.TestData) string { 1004 switch d.Cmd { 1005 case "define": 1006 batch := newBatch(nil) 1007 for _, key := range strings.Split(d.Input, "\n") { 1008 j := strings.Index(key, ":") 1009 ikey := base.ParseInternalKey(key[:j]) 1010 value := []byte(fmt.Sprint(ikey.SeqNum())) 1011 batch.Set(ikey.UserKey, value, nil) 1012 } 1013 var err error 1014 b, err = newFlushableBatch(batch, DefaultComparer) 1015 require.NoError(t, err) 1016 return "" 1017 1018 case "iter": 1019 iter := b.newIter(nil) 1020 defer iter.Close() 1021 return runInternalIterCmd(t, d, iter) 1022 1023 default: 1024 return fmt.Sprintf("unknown command: %s", d.Cmd) 1025 } 1026 }) 1027 } 1028 1029 func TestFlushableBatch(t *testing.T) { 1030 var b *flushableBatch 1031 datadriven.RunTest(t, "testdata/flushable_batch", func(t *testing.T, d *datadriven.TestData) string { 1032 switch d.Cmd { 1033 case "define": 1034 batch := newBatch(nil) 1035 for _, key := range strings.Split(d.Input, "\n") { 1036 j := strings.Index(key, ":") 1037 ikey := base.ParseInternalKey(key[:j]) 1038 value := []byte(fmt.Sprint(ikey.SeqNum())) 1039 switch ikey.Kind() { 1040 case InternalKeyKindDelete: 1041 require.NoError(t, batch.Delete(ikey.UserKey, nil)) 1042 case InternalKeyKindSet: 1043 require.NoError(t, batch.Set(ikey.UserKey, value, nil)) 1044 case InternalKeyKindMerge: 1045 require.NoError(t, batch.Merge(ikey.UserKey, value, nil)) 1046 case InternalKeyKindRangeDelete: 1047 require.NoError(t, batch.DeleteRange(ikey.UserKey, value, nil)) 1048 case InternalKeyKindRangeKeyDelete: 1049 require.NoError(t, batch.RangeKeyDelete(ikey.UserKey, value, nil)) 1050 case InternalKeyKindRangeKeySet: 1051 require.NoError(t, batch.RangeKeySet(ikey.UserKey, value, value, value, nil)) 1052 case InternalKeyKindRangeKeyUnset: 1053 require.NoError(t, batch.RangeKeyUnset(ikey.UserKey, value, value, nil)) 1054 } 1055 } 1056 var err error 1057 b, err = newFlushableBatch(batch, DefaultComparer) 1058 require.NoError(t, err) 1059 return "" 1060 1061 case "iter": 1062 var opts IterOptions 1063 for _, arg := range d.CmdArgs { 1064 if len(arg.Vals) != 1 { 1065 return fmt.Sprintf("%s: %s=<value>", d.Cmd, arg.Key) 1066 } 1067 switch arg.Key { 1068 case "lower": 1069 opts.LowerBound = []byte(arg.Vals[0]) 1070 case "upper": 1071 opts.UpperBound = []byte(arg.Vals[0]) 1072 default: 1073 return fmt.Sprintf("%s: unknown arg: %s", d.Cmd, arg.Key) 1074 } 1075 } 1076 1077 iter := b.newIter(&opts) 1078 defer iter.Close() 1079 return runInternalIterCmd(t, d, iter) 1080 1081 case "dump": 1082 if len(d.CmdArgs) != 1 || len(d.CmdArgs[0].Vals) != 1 || d.CmdArgs[0].Key != "seq" { 1083 return "dump seq=<value>\n" 1084 } 1085 seqNum, err := strconv.Atoi(d.CmdArgs[0].Vals[0]) 1086 if err != nil { 1087 return err.Error() 1088 } 1089 b.setSeqNum(uint64(seqNum)) 1090 1091 var buf bytes.Buffer 1092 1093 iter := newInternalIterAdapter(b.newIter(nil)) 1094 for valid := iter.First(); valid; valid = iter.Next() { 1095 fmt.Fprintf(&buf, "%s:%s\n", iter.Key(), iter.Value()) 1096 } 1097 iter.Close() 1098 1099 if rangeDelIter := b.newRangeDelIter(nil); rangeDelIter != nil { 1100 scanKeyspanIterator(&buf, rangeDelIter) 1101 rangeDelIter.Close() 1102 } 1103 if rangeKeyIter := b.newRangeKeyIter(nil); rangeKeyIter != nil { 1104 scanKeyspanIterator(&buf, rangeKeyIter) 1105 rangeKeyIter.Close() 1106 } 1107 return buf.String() 1108 1109 default: 1110 return fmt.Sprintf("unknown command: %s", d.Cmd) 1111 } 1112 }) 1113 } 1114 1115 func TestFlushableBatchDeleteRange(t *testing.T) { 1116 var fb *flushableBatch 1117 var input string 1118 1119 datadriven.RunTest(t, "testdata/delete_range", func(t *testing.T, td *datadriven.TestData) string { 1120 switch td.Cmd { 1121 case "clear": 1122 input = "" 1123 return "" 1124 1125 case "define": 1126 b := newBatch(nil) 1127 // NB: We can't actually add to the flushable batch as we can to a 1128 // memtable (which shares the "testdata/delete_range" data), so we fake 1129 // it by concatenating the input and rebuilding the flushable batch from 1130 // scratch. 1131 input += "\n" + td.Input 1132 td.Input = input 1133 if err := runBatchDefineCmd(td, b); err != nil { 1134 return err.Error() 1135 } 1136 var err error 1137 fb, err = newFlushableBatch(b, DefaultComparer) 1138 require.NoError(t, err) 1139 return "" 1140 1141 case "scan": 1142 var buf bytes.Buffer 1143 if td.HasArg("range-del") { 1144 fi := fb.newRangeDelIter(nil) 1145 defer fi.Close() 1146 scanKeyspanIterator(&buf, fi) 1147 } else { 1148 ii := fb.newIter(nil) 1149 defer ii.Close() 1150 scanInternalIter(&buf, ii) 1151 } 1152 return buf.String() 1153 1154 default: 1155 return fmt.Sprintf("unknown command: %s", td.Cmd) 1156 } 1157 }) 1158 } 1159 1160 func scanInternalIter(w io.Writer, ii internalIterator) { 1161 for k, v := ii.First(); k != nil; k, v = ii.Next() { 1162 fmt.Fprintf(w, "%s:%s\n", k, v.InPlaceValue()) 1163 } 1164 } 1165 1166 func scanKeyspanIterator(w io.Writer, ki keyspan.FragmentIterator) { 1167 for s := ki.First(); s != nil; s = ki.Next() { 1168 fmt.Fprintln(w, s) 1169 } 1170 } 1171 1172 func TestFlushableBatchBytesIterated(t *testing.T) { 1173 batch := newBatch(nil) 1174 for j := 0; j < 1000; j++ { 1175 key := make([]byte, 8+j%3) 1176 value := make([]byte, 7+j%5) 1177 batch.Set(key, value, nil) 1178 1179 fb, err := newFlushableBatch(batch, DefaultComparer) 1180 require.NoError(t, err) 1181 1182 var bytesIterated uint64 1183 it := fb.newFlushIter(nil, &bytesIterated) 1184 1185 var prevIterated uint64 1186 for key, _ := it.First(); key != nil; key, _ = it.Next() { 1187 if bytesIterated < prevIterated { 1188 t.Fatalf("bytesIterated moved backward: %d < %d", bytesIterated, prevIterated) 1189 } 1190 prevIterated = bytesIterated 1191 } 1192 1193 expected := fb.inuseBytes() 1194 if bytesIterated != expected { 1195 t.Fatalf("bytesIterated: got %d, want %d", bytesIterated, expected) 1196 } 1197 } 1198 } 1199 1200 func TestEmptyFlushableBatch(t *testing.T) { 1201 // Verify that we can create a flushable batch on an empty batch. 1202 fb, err := newFlushableBatch(newBatch(nil), DefaultComparer) 1203 require.NoError(t, err) 1204 it := newInternalIterAdapter(fb.newIter(nil)) 1205 require.False(t, it.First()) 1206 } 1207 1208 func TestBatchCommitStats(t *testing.T) { 1209 testFunc := func() error { 1210 db, err := Open("", &Options{ 1211 FS: vfs.NewMem(), 1212 }) 1213 require.NoError(t, err) 1214 defer db.Close() 1215 b := db.NewBatch() 1216 defer b.Close() 1217 stats := b.CommitStats() 1218 require.Equal(t, BatchCommitStats{}, stats) 1219 1220 // The stall code peers into the internals, instead of adding general 1221 // purpose hooks, to avoid changing production code. We can revisit this 1222 // choice if it becomes hard to maintain. 1223 1224 // Commit semaphore stall funcs. 1225 var unstallCommitSemaphore func() 1226 stallCommitSemaphore := func() { 1227 commitPipeline := db.commit 1228 commitSemaphoreReserved := 0 1229 done := false 1230 for !done { 1231 select { 1232 case commitPipeline.commitQueueSem <- struct{}{}: 1233 commitSemaphoreReserved++ 1234 default: 1235 done = true 1236 } 1237 if done { 1238 break 1239 } 1240 } 1241 unstallCommitSemaphore = func() { 1242 for i := 0; i < commitSemaphoreReserved; i++ { 1243 <-commitPipeline.commitQueueSem 1244 } 1245 } 1246 } 1247 1248 // Memstable stall funcs. 1249 var unstallMemtable func() 1250 stallMemtable := func() { 1251 db.mu.Lock() 1252 defer db.mu.Unlock() 1253 prev := db.opts.MemTableStopWritesThreshold 1254 db.opts.MemTableStopWritesThreshold = 0 1255 unstallMemtable = func() { 1256 db.mu.Lock() 1257 defer db.mu.Unlock() 1258 db.opts.MemTableStopWritesThreshold = prev 1259 db.mu.compact.cond.Broadcast() 1260 } 1261 } 1262 1263 // L0 read-amp stall funcs. 1264 var unstallL0ReadAmp func() 1265 stallL0ReadAmp := func() { 1266 db.mu.Lock() 1267 defer db.mu.Unlock() 1268 prev := db.opts.L0StopWritesThreshold 1269 db.opts.L0StopWritesThreshold = 0 1270 unstallL0ReadAmp = func() { 1271 db.mu.Lock() 1272 defer db.mu.Unlock() 1273 db.opts.L0StopWritesThreshold = prev 1274 db.mu.compact.cond.Broadcast() 1275 } 1276 } 1277 1278 // Commit wait stall funcs. 1279 var unstallCommitWait func() 1280 stallCommitWait := func() { 1281 b.commit.Add(1) 1282 unstallCommitWait = func() { 1283 b.commit.Done() 1284 } 1285 } 1286 1287 // Stall everything. 1288 stallCommitSemaphore() 1289 stallMemtable() 1290 stallL0ReadAmp() 1291 stallCommitWait() 1292 1293 // Exceed initialMemTableSize -- this is needed to make stallMemtable work. 1294 require.NoError(t, b.Set(make([]byte, initialMemTableSize), nil, nil)) 1295 1296 var commitWG sync.WaitGroup 1297 commitWG.Add(1) 1298 go func() { 1299 require.NoError(t, db.Apply(b, &WriteOptions{Sync: true})) 1300 commitWG.Done() 1301 }() 1302 // Unstall things in the order that the stalls will happen. 1303 sleepDuration := 10 * time.Millisecond 1304 time.Sleep(sleepDuration) 1305 unstallCommitSemaphore() 1306 time.Sleep(sleepDuration) 1307 unstallMemtable() 1308 time.Sleep(sleepDuration) 1309 unstallL0ReadAmp() 1310 time.Sleep(sleepDuration) 1311 unstallCommitWait() 1312 1313 // Wait for Apply to return. 1314 commitWG.Wait() 1315 stats = b.CommitStats() 1316 expectedDuration := (2 * sleepDuration) / 3 1317 if expectedDuration > stats.SemaphoreWaitDuration { 1318 return errors.Errorf("SemaphoreWaitDuration %s is too low", 1319 stats.SemaphoreWaitDuration.String()) 1320 } 1321 if expectedDuration > stats.MemTableWriteStallDuration { 1322 return errors.Errorf("MemTableWriteStallDuration %s is too low", 1323 stats.MemTableWriteStallDuration.String()) 1324 } 1325 if expectedDuration > stats.L0ReadAmpWriteStallDuration { 1326 return errors.Errorf("L0ReadAmpWriteStallDuration %s is too low", 1327 stats.L0ReadAmpWriteStallDuration) 1328 } 1329 if expectedDuration > stats.CommitWaitDuration { 1330 return errors.Errorf("CommitWaitDuration %s is too low", 1331 stats.CommitWaitDuration) 1332 } 1333 if 5*expectedDuration > stats.TotalDuration { 1334 return errors.Errorf("TotalDuration %s is too low", 1335 stats.TotalDuration) 1336 } 1337 return nil 1338 } 1339 // Try a few times, and succeed if one of them succeeds. 1340 var err error 1341 for i := 0; i < 5; i++ { 1342 err = testFunc() 1343 if err == nil { 1344 break 1345 } 1346 } 1347 require.NoError(t, err) 1348 } 1349 1350 func TestBatchReader(t *testing.T) { 1351 datadriven.RunTest(t, "testdata/batch_reader", func(t *testing.T, td *datadriven.TestData) string { 1352 switch td.Cmd { 1353 case "scan": 1354 var repr bytes.Buffer 1355 for i, l := range strings.Split(td.Input, "\n") { 1356 // Remove any trailing comments behind #. 1357 if i := strings.IndexRune(l, '#'); i >= 0 { 1358 l = l[:i] 1359 } 1360 // Strip all whitespace from the line. 1361 l = strings.Map(func(r rune) rune { 1362 if unicode.IsSpace(r) { 1363 return -1 1364 } 1365 return r 1366 }, l) 1367 b, err := hex.DecodeString(l) 1368 if err != nil { 1369 return fmt.Sprintf("failed to decode hex; line %d", i) 1370 } 1371 repr.Write(b) 1372 } 1373 r, count := ReadBatch(repr.Bytes()) 1374 var out strings.Builder 1375 fmt.Fprintf(&out, "Count: %d\n", count) 1376 for { 1377 kind, ukey, value, ok, err := r.Next() 1378 if !ok { 1379 if err != nil { 1380 fmt.Fprintf(&out, "err: %s\n", err) 1381 } else { 1382 fmt.Fprint(&out, "eof") 1383 } 1384 break 1385 } 1386 fmt.Fprintf(&out, "%s: %q: %q\n", kind, ukey, value) 1387 } 1388 return out.String() 1389 1390 default: 1391 return fmt.Sprintf("unrecognized command %q", td.Cmd) 1392 } 1393 }) 1394 } 1395 1396 func BenchmarkBatchSet(b *testing.B) { 1397 value := make([]byte, 10) 1398 for i := range value { 1399 value[i] = byte(i) 1400 } 1401 key := make([]byte, 8) 1402 batch := newBatch(nil) 1403 1404 b.ResetTimer() 1405 1406 const batchSize = 1000 1407 for i := 0; i < b.N; i += batchSize { 1408 end := i + batchSize 1409 if end > b.N { 1410 end = b.N 1411 } 1412 1413 for j := i; j < end; j++ { 1414 binary.BigEndian.PutUint64(key, uint64(j)) 1415 batch.Set(key, value, nil) 1416 } 1417 batch.Reset() 1418 } 1419 1420 b.StopTimer() 1421 } 1422 1423 func BenchmarkIndexedBatchSet(b *testing.B) { 1424 value := make([]byte, 10) 1425 for i := range value { 1426 value[i] = byte(i) 1427 } 1428 key := make([]byte, 8) 1429 batch := newIndexedBatch(nil, DefaultComparer) 1430 1431 b.ResetTimer() 1432 1433 const batchSize = 1000 1434 for i := 0; i < b.N; i += batchSize { 1435 end := i + batchSize 1436 if end > b.N { 1437 end = b.N 1438 } 1439 1440 for j := i; j < end; j++ { 1441 binary.BigEndian.PutUint64(key, uint64(j)) 1442 batch.Set(key, value, nil) 1443 } 1444 batch.Reset() 1445 } 1446 1447 b.StopTimer() 1448 } 1449 1450 func BenchmarkBatchSetDeferred(b *testing.B) { 1451 value := make([]byte, 10) 1452 for i := range value { 1453 value[i] = byte(i) 1454 } 1455 key := make([]byte, 8) 1456 batch := newBatch(nil) 1457 1458 b.ResetTimer() 1459 1460 const batchSize = 1000 1461 for i := 0; i < b.N; i += batchSize { 1462 end := i + batchSize 1463 if end > b.N { 1464 end = b.N 1465 } 1466 1467 for j := i; j < end; j++ { 1468 binary.BigEndian.PutUint64(key, uint64(j)) 1469 deferredOp := batch.SetDeferred(len(key), len(value)) 1470 1471 copy(deferredOp.Key, key) 1472 copy(deferredOp.Value, value) 1473 1474 deferredOp.Finish() 1475 } 1476 batch.Reset() 1477 } 1478 1479 b.StopTimer() 1480 } 1481 1482 func BenchmarkIndexedBatchSetDeferred(b *testing.B) { 1483 value := make([]byte, 10) 1484 for i := range value { 1485 value[i] = byte(i) 1486 } 1487 key := make([]byte, 8) 1488 batch := newIndexedBatch(nil, DefaultComparer) 1489 1490 b.ResetTimer() 1491 1492 const batchSize = 1000 1493 for i := 0; i < b.N; i += batchSize { 1494 end := i + batchSize 1495 if end > b.N { 1496 end = b.N 1497 } 1498 1499 for j := i; j < end; j++ { 1500 binary.BigEndian.PutUint64(key, uint64(j)) 1501 deferredOp := batch.SetDeferred(len(key), len(value)) 1502 1503 copy(deferredOp.Key, key) 1504 copy(deferredOp.Value, value) 1505 1506 deferredOp.Finish() 1507 } 1508 batch.Reset() 1509 } 1510 1511 b.StopTimer() 1512 } 1513 1514 func TestBatchMemTableSizeOverflow(t *testing.T) { 1515 opts := &Options{ 1516 FS: vfs.NewMem(), 1517 } 1518 opts.EnsureDefaults() 1519 d, err := Open("", opts) 1520 require.NoError(t, err) 1521 1522 bigValue := make([]byte, 1000) 1523 b := d.NewBatch() 1524 1525 // memTableSize can overflow as a uint32. 1526 b.memTableSize = math.MaxUint32 - 50 1527 for i := 0; i < 10; i++ { 1528 k := fmt.Sprintf("key-%05d", i) 1529 require.NoError(t, b.Set([]byte(k), bigValue, nil)) 1530 } 1531 require.Greater(t, b.memTableSize, uint64(math.MaxUint32)) 1532 require.NoError(t, b.Close()) 1533 require.NoError(t, d.Close()) 1534 } 1535 1536 // TestBatchSpanCaching stress tests the caching of keyspan.Spans for range 1537 // tombstones and range keys. 1538 func TestBatchSpanCaching(t *testing.T) { 1539 opts := &Options{ 1540 Comparer: testkeys.Comparer, 1541 FS: vfs.NewMem(), 1542 FormatMajorVersion: internalFormatNewest, 1543 } 1544 d, err := Open("", opts) 1545 require.NoError(t, err) 1546 defer d.Close() 1547 1548 ks := testkeys.Alpha(1) 1549 b := d.NewIndexedBatch() 1550 for i := int64(0); i < ks.Count(); i++ { 1551 k := testkeys.Key(ks, i) 1552 require.NoError(t, b.Set(k, k, nil)) 1553 } 1554 1555 seed := int64(time.Now().UnixNano()) 1556 t.Logf("seed = %d", seed) 1557 rng := rand.New(rand.NewSource(seed)) 1558 iters := make([][]*Iterator, ks.Count()) 1559 defer func() { 1560 for _, keyIters := range iters { 1561 for _, iter := range keyIters { 1562 _ = iter.Close() 1563 } 1564 } 1565 }() 1566 1567 // This test begins with one point key for every letter of the alphabet. 1568 // Over the course of the test, point keys are 'replaced' with range keys 1569 // with narrow bounds from left to right. Iterators are created at random, 1570 // sometimes from the batch and sometimes by cloning existing iterators. 1571 1572 checkIter := func(iter *Iterator, nextKey int64) { 1573 var i int64 1574 for valid := iter.First(); valid; valid = iter.Next() { 1575 hasPoint, hasRange := iter.HasPointAndRange() 1576 require.Equal(t, testkeys.Key(ks, i), iter.Key()) 1577 if i < nextKey { 1578 // This key should not exist as a point key, just a range key. 1579 require.False(t, hasPoint) 1580 require.True(t, hasRange) 1581 } else { 1582 require.True(t, hasPoint) 1583 require.False(t, hasRange) 1584 } 1585 i++ 1586 } 1587 require.Equal(t, ks.Count(), i) 1588 } 1589 1590 // Each iteration of the below loop either reads or writes. 1591 // 1592 // A write iteration writes a new RANGEDEL and RANGEKEYSET into the batch, 1593 // covering a single point key seeded above. Writing these two span keys 1594 // together 'replaces' the point key with a range key. Each write iteration 1595 // ratchets nextWriteKey so the next write iteration will write the next 1596 // key. 1597 // 1598 // A read iteration creates a new iterator and ensures its state is 1599 // expected: some prefix of only point keys, followed by a suffix of only 1600 // range keys. Iterators created through Clone should observe the point keys 1601 // that existed when the cloned iterator was created. 1602 for nextWriteKey := int64(0); nextWriteKey < ks.Count(); { 1603 p := rng.Float64() 1604 switch { 1605 case p < .10: /* 10 % */ 1606 // Write a new range deletion and range key. 1607 start := testkeys.Key(ks, nextWriteKey) 1608 end := append(start, 0x00) 1609 require.NoError(t, b.DeleteRange(start, end, nil)) 1610 require.NoError(t, b.RangeKeySet(start, end, nil, []byte("foo"), nil)) 1611 nextWriteKey++ 1612 case p < .55: /* 45 % */ 1613 // Create a new iterator directly from the batch and check that it 1614 // observes the correct state. 1615 iter, _ := b.NewIter(&IterOptions{KeyTypes: IterKeyTypePointsAndRanges}) 1616 checkIter(iter, nextWriteKey) 1617 iters[nextWriteKey] = append(iters[nextWriteKey], iter) 1618 default: /* 45 % */ 1619 // Create a new iterator through cloning a random existing iterator 1620 // and check that it observes the right state. 1621 readKey := rng.Int63n(nextWriteKey + 1) 1622 itersForReadKey := iters[readKey] 1623 if len(itersForReadKey) == 0 { 1624 continue 1625 } 1626 iter, err := itersForReadKey[rng.Intn(len(itersForReadKey))].Clone(CloneOptions{}) 1627 require.NoError(t, err) 1628 checkIter(iter, readKey) 1629 iters[readKey] = append(iters[readKey], iter) 1630 } 1631 } 1632 }