github.com/zuoyebang/bitalosdb@v1.1.1-0.20240516111551-79a8c4d8ce20/iterator_test.go (about) 1 // Copyright 2021 The Bitalosdb author(hustxrb@163.com) and other contributors. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package bitalosdb 16 17 import ( 18 "fmt" 19 "math" 20 "os" 21 "runtime/debug" 22 "sort" 23 "strconv" 24 "strings" 25 "sync" 26 "testing" 27 28 "github.com/stretchr/testify/require" 29 "github.com/zuoyebang/bitalosdb/internal/base" 30 "github.com/zuoyebang/bitalosdb/internal/consts" 31 "github.com/zuoyebang/bitalosdb/internal/options" 32 "github.com/zuoyebang/bitalosdb/internal/unsafe2" 33 "github.com/zuoyebang/bitalosdb/internal/utils" 34 ) 35 36 func testIter(t *testing.T) { 37 db := openTestDB(testDirname, nil) 38 defer func() { 39 require.NoError(t, db.Close()) 40 }() 41 42 o := &IterOptions{} 43 iter := new(Iterator) 44 cmp := DefaultComparer.Compare 45 46 o = &IterOptions{ 47 LowerBound: makeTestSlotKey([]byte("key_1")), 48 UpperBound: makeTestSlotKey([]byte("key_73")), 49 SlotId: uint32(testSlotId), 50 } 51 iter = db.NewIter(o) 52 for iter.First(); iter.Valid(); iter.Next() { 53 if cmp(iter.Key(), o.GetLowerBound()) < 0 { 54 t.Fatalf("get %s less than %s", iter.Key(), o.GetLowerBound()) 55 } 56 if cmp(iter.Key(), o.GetUpperBound()) >= 0 { 57 t.Fatalf("get %s greater than %s", iter.Key(), o.GetUpperBound()) 58 } 59 } 60 require.NoError(t, iter.Close()) 61 62 o = &IterOptions{ 63 LowerBound: makeTestSlotKey([]byte("key_4")), 64 UpperBound: makeTestSlotKey([]byte("key_82")), 65 SlotId: uint32(testSlotId), 66 } 67 iter = db.NewIter(o) 68 for iter.SeekGE(makeTestSlotKey([]byte("key_6"))); iter.Valid(); iter.Next() { 69 if cmp(iter.Key(), o.GetLowerBound()) < 0 { 70 t.Fatalf("get %s less than %s", iter.Key(), o.GetLowerBound()) 71 } 72 if cmp(iter.Key(), o.GetUpperBound()) >= 0 { 73 t.Fatalf("get %s greater than %s", iter.Key(), o.GetUpperBound()) 74 } 75 } 76 require.NoError(t, iter.Close()) 77 78 o = &IterOptions{ 79 LowerBound: makeTestSlotKey([]byte("key_4")), 80 UpperBound: makeTestSlotKey([]byte("key_54")), 81 SlotId: uint32(testSlotId), 82 } 83 iter = db.NewIter(o) 84 for iter.SeekLT(makeTestSlotKey([]byte("key_6"))); iter.Valid(); iter.Prev() { 85 if cmp(iter.Key(), o.GetLowerBound()) < 0 { 86 t.Fatalf("get %s less than %s", iter.Key(), o.GetLowerBound()) 87 } 88 if cmp(iter.Key(), o.GetUpperBound()) >= 0 { 89 t.Fatalf("get %s greater than %s", iter.Key(), o.GetUpperBound()) 90 } 91 } 92 require.NoError(t, iter.Close()) 93 } 94 95 func TestIter(t *testing.T) { 96 defer os.RemoveAll(testDirname) 97 testBitalosdbWrite(t, 100, true) 98 testIter(t) 99 } 100 101 func testIterSeek(t *testing.T) { 102 db := openTestDB(testDirname, nil) 103 defer func() { 104 require.NoError(t, db.Close()) 105 }() 106 107 iter := new(Iterator) 108 cmp := DefaultComparer.Compare 109 var want []byte 110 111 o := &IterOptions{ 112 SlotId: uint32(testSlotId), 113 } 114 iter = db.NewIter(o) 115 defer iter.Close() 116 iter.First() 117 if iter.Valid() { 118 want = makeTestSlotKey([]byte("key_10")) 119 iter.Next() 120 if cmp(iter.Key(), want) != 0 { 121 t.Fatalf("got %s want %s", iter.Key(), want) 122 } 123 124 want = makeTestSlotKey([]byte("key_71")) 125 iter.SeekGE(want) 126 if cmp(iter.Key(), want) != 0 { 127 t.Fatalf("got %s want %s", iter.Key(), want) 128 } 129 130 want = makeTestSlotKey([]byte("key_72")) 131 iter.Next() 132 if cmp(iter.Key(), want) != 0 { 133 t.Fatalf("got %s want %s", iter.Key(), want) 134 } 135 136 want = makeTestSlotKey([]byte("key_71")) 137 iter.Prev() 138 if cmp(iter.Key(), want) != 0 { 139 t.Fatalf("got %s want %s", iter.Key(), want) 140 } 141 142 want = makeTestSlotKey([]byte("key_52")) 143 iter.SeekLT(makeTestSlotKey([]byte("key_53"))) 144 if cmp(iter.Key(), want) != 0 { 145 t.Fatalf("got %s want %s", iter.Key(), want) 146 } 147 148 want = makeTestSlotKey([]byte("key_51")) 149 iter.Prev() 150 if cmp(iter.Key(), want) != 0 { 151 t.Fatalf("got %s want %s", iter.Key(), want) 152 } 153 154 want = makeTestSlotKey([]byte("key_52")) 155 iter.Next() 156 if cmp(iter.Key(), want) != 0 { 157 t.Fatalf("got %s want %s", iter.Key(), want) 158 } 159 } 160 } 161 162 func TestNewIter(t *testing.T) { 163 defer os.RemoveAll(testDirname) 164 os.RemoveAll(testDirname) 165 db := openTestDB(testDirname, nil) 166 defer func() { 167 require.NoError(t, db.Close()) 168 }() 169 170 it := db.NewIter(&IterOptions{IsAll: true}) 171 require.Equal(t, -1, it.index) 172 require.Equal(t, 8, len(it.alloc.merging.levels)) 173 require.NoError(t, it.Close()) 174 175 it = db.NewIter(nil) 176 require.Equal(t, -1, it.index) 177 require.Equal(t, 8, len(it.alloc.merging.levels)) 178 require.NoError(t, it.Close()) 179 180 it = db.NewIter(&IterOptions{IsAll: false}) 181 require.Equal(t, 0, it.index) 182 require.Equal(t, 1, len(it.alloc.merging.levels)) 183 require.NoError(t, it.Close()) 184 185 it = db.NewIter(&IterOptions{SlotId: 6}) 186 require.Equal(t, 6, it.index) 187 require.Equal(t, 1, len(it.alloc.merging.levels)) 188 require.NoError(t, it.Close()) 189 190 it = db.NewIter(&IterOptions{SlotId: 61}) 191 require.Equal(t, 5, it.index) 192 require.Equal(t, 1, len(it.alloc.merging.levels)) 193 require.NoError(t, it.Close()) 194 195 for i := 0; i < 100; i++ { 196 key := makeTestIntKey(i) 197 require.NoError(t, db.Set(key, key, NoSync)) 198 } 199 200 it = db.NewIter(&IterOptions{IsAll: true}) 201 require.Equal(t, -1, it.index) 202 require.Equal(t, 16, len(it.alloc.merging.levels)) 203 require.NoError(t, it.Close()) 204 205 it = db.NewIter(&IterOptions{SlotId: 6}) 206 require.Equal(t, 6, it.index) 207 require.Equal(t, 2, len(it.alloc.merging.levels)) 208 require.NoError(t, it.Close()) 209 210 it = db.NewIter(&IterOptions{SlotId: 61}) 211 require.Equal(t, 5, it.index) 212 require.Equal(t, 2, len(it.alloc.merging.levels)) 213 require.NoError(t, it.Close()) 214 } 215 216 func TestNewBitreeIter(t *testing.T) { 217 defer os.RemoveAll(testDirname) 218 os.RemoveAll(testDirname) 219 220 testOptsUseBitable = true 221 db := openTestDB(testDirname, nil) 222 iterOpts := &options.IterOptions{SlotId: 2} 223 index := base.GetBitowerIndex(int(iterOpts.GetSlotId())) 224 siters := db.bitowers[index].newBitreeIter(iterOpts) 225 require.Equal(t, 1, len(siters)) 226 for i := range siters { 227 require.NoError(t, siters[i].Close()) 228 } 229 230 newAllBitreeIters := func(o *options.IterOptions) (its []base.InternalIterator) { 231 for i := range db.bitowers { 232 btreeIters := db.bitowers[i].newBitreeIter(o) 233 if len(btreeIters) > 0 { 234 its = append(its, btreeIters...) 235 } 236 } 237 return its 238 } 239 240 iters := newAllBitreeIters(nil) 241 require.Equal(t, 8, len(iters)) 242 for i := range iters { 243 require.NoError(t, iters[i].Close()) 244 } 245 246 db.setFlushedBitable() 247 siters = db.bitowers[index].newBitreeIter(iterOpts) 248 require.Equal(t, 2, len(siters)) 249 for i := range siters { 250 require.NoError(t, siters[i].Close()) 251 } 252 iters = newAllBitreeIters(nil) 253 require.Equal(t, 16, len(iters)) 254 for i := range iters { 255 require.NoError(t, iters[i].Close()) 256 } 257 require.NoError(t, db.Close()) 258 259 db = openTestDB(testDirname, nil) 260 siters = db.bitowers[index].newBitreeIter(iterOpts) 261 require.Equal(t, 1, len(siters)) 262 for i := range siters { 263 require.NoError(t, siters[i].Close()) 264 } 265 iters = newAllBitreeIters(nil) 266 require.Equal(t, 8, len(iters)) 267 for i := range iters { 268 require.NoError(t, iters[i].Close()) 269 } 270 require.NoError(t, db.Close()) 271 } 272 273 func TestIterSeek(t *testing.T) { 274 defer os.RemoveAll(testDirname) 275 os.RemoveAll(testDirname) 276 testBitalosdbWrite(t, 100, true) 277 testIterSeek(t) 278 } 279 280 func TestIterPrefix(t *testing.T) { 281 defer os.RemoveAll(testDirname) 282 os.RemoveAll(testDirname) 283 db := openTestDB(testDirname, nil) 284 defer func() { 285 require.NoError(t, db.Close()) 286 }() 287 288 keyUpperBound := func(b []byte) []byte { 289 end := make([]byte, len(b)) 290 copy(end, b) 291 for i := len(end) - 1; i >= 0; i-- { 292 end[i] = end[i] + 1 293 if end[i] != 0 { 294 return end[:i+1] 295 } 296 } 297 return nil 298 } 299 300 keys := []string{"hello", "world", "hello world"} 301 for _, key := range keys { 302 require.NoError(t, db.Set([]byte(key), []byte(key), Sync)) 303 } 304 305 rangeIter := func() { 306 iterOpts := &IterOptions{ 307 IsAll: true, 308 } 309 iter := db.NewIter(iterOpts) 310 i := 0 311 for iter.First(); iter.Valid(); iter.Next() { 312 if i == 0 { 313 require.Equal(t, []byte("hello"), iter.Key()) 314 } else if i == 1 { 315 require.Equal(t, []byte("hello world"), iter.Key()) 316 } 317 i++ 318 } 319 require.NoError(t, iter.Close()) 320 321 iterOpts = &IterOptions{ 322 IsAll: true, 323 LowerBound: []byte("hello"), 324 UpperBound: keyUpperBound([]byte("hello")), 325 } 326 iter = db.NewIter(iterOpts) 327 i = 0 328 for iter.First(); iter.Valid(); iter.Next() { 329 if i == 0 { 330 require.Equal(t, []byte("hello"), iter.Key()) 331 } else if i == 1 { 332 require.Equal(t, []byte("hello world"), iter.Key()) 333 } 334 i++ 335 } 336 require.NoError(t, iter.Close()) 337 } 338 339 rangeIter() 340 341 require.NoError(t, db.Flush()) 342 for _, key := range keys { 343 require.NoError(t, verifyGet(db, []byte(key), []byte(key))) 344 } 345 346 rangeIter() 347 } 348 349 func TestIterSeekGE(t *testing.T) { 350 defer os.RemoveAll(testDirname) 351 os.RemoveAll(testDirname) 352 db := openTestDB(testDirname, nil) 353 defer func() { 354 require.NoError(t, db.Close()) 355 }() 356 357 keys := []string{"hello", "world", "hello world"} 358 for _, key := range keys { 359 require.NoError(t, db.Set([]byte(key), []byte(key), Sync)) 360 } 361 362 rangeIter := func() { 363 iterOpts := &IterOptions{ 364 IsAll: true, 365 } 366 iter := db.NewIter(iterOpts) 367 368 if iter.SeekGE([]byte("a")); iter.Valid() { 369 require.Equal(t, []byte("hello"), iter.Value()) 370 } 371 if iter.SeekGE([]byte("hello w")); iter.Valid() { 372 require.Equal(t, []byte("hello world"), iter.Value()) 373 } 374 if iter.SeekGE([]byte("w")); iter.Valid() { 375 require.Equal(t, []byte("world"), iter.Value()) 376 } 377 378 require.NoError(t, iter.Close()) 379 } 380 381 rangeIter() 382 require.NoError(t, db.Flush()) 383 rangeIter() 384 } 385 386 func TestIterSeekLT(t *testing.T) { 387 defer os.RemoveAll(testDirname) 388 os.RemoveAll(testDirname) 389 390 db := openTestDB(testDirname, nil) 391 defer func() { 392 require.NoError(t, db.Close()) 393 }() 394 395 batch := db.NewBatchBitower() 396 seek := makeTestSlotKey([]byte(fmt.Sprintf("key_%s", utils.Float64ToByteSort(math.MaxFloat64, nil)))) 397 key0 := makeTestSlotKey([]byte("key_0")) 398 val0 := makeTestSlotKey([]byte("val_0")) 399 key1 := makeTestSlotKey([]byte("key_1")) 400 val1 := makeTestSlotKey([]byte("val_1")) 401 key2 := makeTestSlotKey([]byte("key_2")) 402 val2 := makeTestSlotKey([]byte("val_2")) 403 require.NoError(t, batch.Set(key0, val0, NoSync)) 404 require.NoError(t, batch.Set(key1, val1, NoSync)) 405 require.NoError(t, batch.Set(key2, val2, NoSync)) 406 require.NoError(t, batch.Commit(NoSync)) 407 require.NoError(t, batch.Close()) 408 409 it := db.NewIter(&IterOptions{SlotId: uint32(testSlotId)}) 410 for it.SeekLT(seek); it.Valid(); it.Next() { 411 require.Equal(t, key2, it.Key()) 412 require.Equal(t, val2, it.Value()) 413 } 414 require.NoError(t, it.Close()) 415 require.NoError(t, db.Close()) 416 417 db = openTestDB(testDirname, nil) 418 it = db.NewIter(&IterOptions{SlotId: uint32(testSlotId)}) 419 for it.SeekLT(seek); it.Valid(); it.Next() { 420 require.Equal(t, key2, it.Key()) 421 require.Equal(t, val2, it.Value()) 422 } 423 require.NoError(t, it.Close()) 424 425 it = db.NewIter(&IterOptions{SlotId: uint32(testSlotId)}) 426 cnt := 2 427 for it.Last(); it.Valid(); it.Prev() { 428 require.Equal(t, makeTestSlotKey([]byte(fmt.Sprintf("key_%d", cnt))), it.Key()) 429 require.Equal(t, makeTestSlotKey([]byte(fmt.Sprintf("val_%d", cnt))), it.Value()) 430 cnt-- 431 } 432 require.NoError(t, it.Close()) 433 } 434 435 func TestIterReadAmp(t *testing.T) { 436 for _, isNext := range []bool{true, false} { 437 t.Run(fmt.Sprintf("isNext=%t", isNext), func(t *testing.T) { 438 defer os.RemoveAll(testDirname) 439 os.RemoveAll(testDirname) 440 441 db := openTestDB(testDirname, nil) 442 defer func() { 443 require.NoError(t, db.Close()) 444 }() 445 446 for index := range db.bitowers { 447 it := db.NewIter(&IterOptions{SlotId: uint32(index)}) 448 for it.First(); it.Valid(); it.Next() { 449 } 450 require.NoError(t, it.Close()) 451 } 452 for i := range db.bitowers { 453 require.Equal(t, uint64(0), db.bitowers[i].iterSlowCount.Load()) 454 } 455 456 num := (consts.IterSlowCountThreshold + 10) * len(db.bitowers) 457 for i := 0; i < num; i++ { 458 key := makeTestIntKey(i) 459 require.NoError(t, db.Set(key, key, NoSync)) 460 require.NoError(t, db.Delete(key, NoSync)) 461 } 462 463 loop := consts.IterReadAmplificationThreshold + 1 464 for i := range db.bitowers { 465 for j := 0; j < loop; j++ { 466 it := db.NewIter(&IterOptions{SlotId: uint32(i)}) 467 for it.First(); it.Valid(); it.Last() { 468 } 469 require.NoError(t, it.Close()) 470 require.Equal(t, uint64(1+j), db.bitowers[i].iterSlowCount.Load()) 471 } 472 } 473 474 db.compactIterReadAmplification(1) 475 for i := range db.bitowers { 476 require.Equal(t, uint64(0), db.bitowers[i].iterSlowCount.Load()) 477 } 478 }) 479 } 480 } 481 482 func TestIterPrefixDeleteKey(t *testing.T) { 483 defer os.RemoveAll(testDirname) 484 os.RemoveAll(testDirname) 485 db := openTestDB(testDirname, nil) 486 defer func() { 487 require.NoError(t, db.Close()) 488 }() 489 490 val := testRandBytes(2048) 491 pdKey := makeTestSlotKey([]byte("key_prefix_delete")) 492 493 for i := 0; i < 100; i++ { 494 newKey := makeTestSlotKey([]byte(fmt.Sprintf("key_%d", i))) 495 require.NoError(t, db.Set(newKey, val, NoSync)) 496 } 497 498 batch := db.NewBatch() 499 batch.PrefixDeleteKeySet(pdKey, NoSync) 500 require.NoError(t, batch.Commit(NoSync)) 501 502 testIter := func() { 503 it := db.NewIter(&IterOptions{SlotId: uint32(testSlotId)}) 504 for it.First(); it.Valid(); it.Next() { 505 require.NotEqual(t, pdKey, it.Key()) 506 } 507 isFound := it.SeekGE(pdKey) 508 require.Equal(t, false, isFound) 509 require.NoError(t, it.Close()) 510 it = db.NewIter(&IterOptions{SlotId: uint32(testSlotId)}) 511 for it.Last(); it.Valid(); it.Prev() { 512 require.NotEqual(t, pdKey, it.Key()) 513 } 514 isFound = it.SeekLT(pdKey) 515 require.Equal(t, it.Key(), makeTestSlotKey([]byte(fmt.Sprintf("key_%d", 99)))) 516 require.NoError(t, it.Close()) 517 } 518 519 testIter() 520 require.NoError(t, db.Flush()) 521 testIter() 522 } 523 524 func TestIterMemShard(t *testing.T) { 525 defer os.RemoveAll(testDirname) 526 os.RemoveAll(testDirname) 527 db := openTestDB(testDirname, nil) 528 defer func() { 529 require.NoError(t, db.Close()) 530 }() 531 532 val := testRandBytes(100) 533 534 batch := db.NewBatch() 535 for i := 0; i < 100; i++ { 536 key := makeTestSlotKey([]byte(fmt.Sprintf("key_%d", i))) 537 require.NoError(t, batch.Set(key, val, NoSync)) 538 } 539 require.NoError(t, batch.Commit(NoSync)) 540 require.NoError(t, batch.Close()) 541 542 iterOpts := &IterOptions{ 543 LowerBound: makeTestSlotKey([]byte("key_1")), 544 SlotId: uint32(testSlotId), 545 } 546 it := db.NewIter(iterOpts) 547 defer it.Close() 548 wg := sync.WaitGroup{} 549 for it.SeekGE(iterOpts.LowerBound); it.Valid(); it.Next() { 550 key := utils.CloneBytes(it.Key()) 551 indexStr := strings.Split(string(key), "_")[1] 552 index, _ := strconv.Atoi(indexStr) 553 if index >= 20 && index <= 40 { 554 wg.Add(1) 555 go func(k []byte) { 556 defer wg.Done() 557 b := db.NewBatch() 558 require.NoError(t, b.Delete(k, NoSync)) 559 require.NoError(t, b.Commit(NoSync)) 560 require.NoError(t, b.Close()) 561 _, _, err := db.Get(k) 562 if err != ErrNotFound { 563 t.Fatal("get delete key find ", string(k)) 564 } 565 }(key) 566 } 567 } 568 wg.Wait() 569 } 570 571 func TestIterStat(t *testing.T) { 572 defer os.RemoveAll(testDirname) 573 os.RemoveAll(testDirname) 574 db := openTestDB(testDirname, nil) 575 defer func() { 576 require.NoError(t, db.Close()) 577 }() 578 579 keyList := make([]string, 100) 580 for i := 0; i < 100; i++ { 581 key := makeTestSlotIntKey(i) 582 keyList[i] = unsafe2.String(key) 583 for j := 0; j < 100; j++ { 584 require.NoError(t, db.Set(key, []byte(fmt.Sprintf("%d", j)), NoSync)) 585 } 586 } 587 sort.Strings(keyList) 588 589 iterOpts := &IterOptions{ 590 SlotId: uint32(testSlotId), 591 } 592 it := db.NewIter(iterOpts) 593 defer it.Close() 594 i := 0 595 for it.First(); it.Valid(); it.Next() { 596 require.Equal(t, unsafe2.ByteSlice(keyList[i]), it.Key()) 597 require.Equal(t, []byte("99"), it.Value()) 598 i++ 599 } 600 stats := it.Stats() 601 require.Equal(t, 1, stats.ForwardSeekCount[0]) 602 require.Equal(t, 100, stats.ForwardStepCount[0]) 603 require.Equal(t, 1, stats.ForwardSeekCount[1]) 604 require.Equal(t, 200, stats.ForwardStepCount[1]) 605 } 606 607 func TestIterAll(t *testing.T) { 608 defer os.RemoveAll(testDirname) 609 os.RemoveAll(testDirname) 610 611 testOptsDisableWAL = true 612 db := openTestDB(testDirname, nil) 613 defer func() { 614 if r := recover(); r != nil { 615 fmt.Printf("panic err:%v stack:%s", r, string(debug.Stack())) 616 } 617 require.NoError(t, db.Close()) 618 }() 619 620 keyNum := 10000 621 keyList := make([]string, keyNum) 622 for i := 0; i < keyNum; i++ { 623 key := makeTestIntKey(i) 624 keyList[i] = unsafe2.String(key) 625 require.NoError(t, db.Set(key, key, NoSync)) 626 } 627 sort.Strings(keyList) 628 629 iterOpts := &IterOptions{ 630 LowerBound: nil, 631 UpperBound: nil, 632 IsAll: true, 633 SlotId: 0, 634 } 635 it := db.NewIter(iterOpts) 636 i := 0 637 for it.First(); it.Valid(); it.Next() { 638 require.Equal(t, unsafe2.ByteSlice(keyList[i]), it.Key()) 639 require.Equal(t, it.Key(), it.Value()) 640 i++ 641 } 642 require.NoError(t, it.Close()) 643 }