github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/internal/manifest/btree.go (about) 1 // Copyright 2020 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 manifest 6 7 import ( 8 "bytes" 9 "fmt" 10 "strings" 11 "sync/atomic" 12 "unsafe" 13 14 "github.com/cockroachdb/errors" 15 "github.com/cockroachdb/pebble/internal/invariants" 16 ) 17 18 // The Annotator type defined below is used by other packages to lazily 19 // compute a value over a B-Tree. Each node of the B-Tree stores one 20 // `annotation` per annotator, containing the result of the computation over 21 // the node's subtree. 22 // 23 // An annotation is marked as valid if it's current with the current subtree 24 // state. Annotations are marked as invalid whenever a node will be mutated 25 // (in mut). Annotators may also return `false` from `Accumulate` to signal 26 // that a computation for a file is not stable and may change in the future. 27 // Annotations that include these unstable values are also marked as invalid 28 // on the node, ensuring that future queries for the annotation will recompute 29 // the value. 30 31 // An Annotator defines a computation over a level's FileMetadata. If the 32 // computation is stable and uses inputs that are fixed for the lifetime of 33 // a FileMetadata, the LevelMetadata's internal data structures are annotated 34 // with the intermediary computations. This allows the computation to be 35 // computed incrementally as edits are applied to a level. 36 type Annotator interface { 37 // Zero returns the zero value of an annotation. This value is returned 38 // when a LevelMetadata is empty. The dst argument, if non-nil, is an 39 // obsolete value previously returned by this Annotator and may be 40 // overwritten and reused to avoid a memory allocation. 41 Zero(dst interface{}) (v interface{}) 42 43 // Accumulate computes the annotation for a single file in a level's 44 // metadata. It merges the file's value into dst and returns a bool flag 45 // indicating whether or not the value is stable and okay to cache as an 46 // annotation. If the file's value may change over the life of the file, 47 // the annotator must return false. 48 // 49 // Implementations may modify dst and return it to avoid an allocation. 50 Accumulate(m *FileMetadata, dst interface{}) (v interface{}, cacheOK bool) 51 52 // Merge combines two values src and dst, returning the result. 53 // Implementations may modify dst and return it to avoid an allocation. 54 Merge(src interface{}, dst interface{}) interface{} 55 } 56 57 type btreeCmp func(*FileMetadata, *FileMetadata) int 58 59 func btreeCmpSeqNum(a, b *FileMetadata) int { 60 return a.cmpSeqNum(b) 61 } 62 63 func btreeCmpSmallestKey(cmp Compare) btreeCmp { 64 return func(a, b *FileMetadata) int { 65 return a.cmpSmallestKey(b, cmp) 66 } 67 } 68 69 // btreeCmpSpecificOrder is used in tests to construct a B-Tree with a 70 // specific ordering of FileMetadata within the tree. It's typically used to 71 // test consistency checking code that needs to construct a malformed B-Tree. 72 func btreeCmpSpecificOrder(files []*FileMetadata) btreeCmp { 73 m := map[*FileMetadata]int{} 74 for i, f := range files { 75 m[f] = i 76 } 77 return func(a, b *FileMetadata) int { 78 ai, aok := m[a] 79 bi, bok := m[b] 80 if !aok || !bok { 81 panic("btreeCmpSliceOrder called with unknown files") 82 } 83 switch { 84 case ai < bi: 85 return -1 86 case ai > bi: 87 return +1 88 default: 89 return 0 90 } 91 } 92 } 93 94 const ( 95 degree = 16 96 maxItems = 2*degree - 1 97 minItems = degree - 1 98 ) 99 100 type annotation struct { 101 annotator Annotator 102 // v is an annotation value, the output of either 103 // annotator.Value or annotator.Merge. 104 v interface{} 105 // valid indicates whether future reads of the annotation may use v as-is. 106 // If false, v will be zeroed and recalculated. 107 valid bool 108 } 109 110 type leafNode struct { 111 ref atomic.Int32 112 count int16 113 leaf bool 114 // subtreeCount holds the count of files in the entire subtree formed by 115 // this node. For leaf nodes, subtreeCount is always equal to count. For 116 // non-leaf nodes, it's the sum of count plus all the children's 117 // subtreeCounts. 118 // 119 // NB: We could move this field to the end of the node struct, since leaf => 120 // count=subtreeCount, however the unsafe casting [leafToNode] performs make 121 // it risky and cumbersome. 122 subtreeCount int 123 items [maxItems]*FileMetadata 124 // annot contains one annotation per annotator, merged over the entire 125 // node's files (and all descendants for non-leaf nodes). 126 annot []annotation 127 } 128 129 type node struct { 130 leafNode 131 children [maxItems + 1]*node 132 } 133 134 //go:nocheckptr casts a ptr to a smaller struct to a ptr to a larger struct. 135 func leafToNode(ln *leafNode) *node { 136 return (*node)(unsafe.Pointer(ln)) 137 } 138 139 func newLeafNode() *node { 140 n := leafToNode(new(leafNode)) 141 n.leaf = true 142 n.ref.Store(1) 143 return n 144 } 145 146 func newNode() *node { 147 n := new(node) 148 n.ref.Store(1) 149 return n 150 } 151 152 // mut creates and returns a mutable node reference. If the node is not shared 153 // with any other trees then it can be modified in place. Otherwise, it must be 154 // cloned to ensure unique ownership. In this way, we enforce a copy-on-write 155 // policy which transparently incorporates the idea of local mutations, like 156 // Clojure's transients or Haskell's ST monad, where nodes are only copied 157 // during the first time that they are modified between Clone operations. 158 // 159 // When a node is cloned, the provided pointer will be redirected to the new 160 // mutable node. 161 func mut(n **node) *node { 162 if (*n).ref.Load() == 1 { 163 // Exclusive ownership. Can mutate in place. 164 165 // Whenever a node will be mutated, reset its annotations to be marked 166 // as uncached. This ensures any future calls to (*node).annotation 167 // will recompute annotations on the modified subtree. 168 for i := range (*n).annot { 169 (*n).annot[i].valid = false 170 } 171 return *n 172 } 173 // If we do not have unique ownership over the node then we 174 // clone it to gain unique ownership. After doing so, we can 175 // release our reference to the old node. We pass recursive 176 // as true because even though we just observed the node's 177 // reference count to be greater than 1, we might be racing 178 // with another call to decRef on this node. 179 c := (*n).clone() 180 (*n).decRef(true /* contentsToo */, nil) 181 *n = c 182 // NB: We don't need to clear annotations, because (*node).clone does not 183 // copy them. 184 return *n 185 } 186 187 // incRef acquires a reference to the node. 188 func (n *node) incRef() { 189 n.ref.Add(1) 190 } 191 192 // decRef releases a reference to the node. If requested, the method will unref 193 // its items and recurse into child nodes and decrease their refcounts as well. 194 // Some internal codepaths that manually copy the node's items or children to 195 // new nodes pass contentsToo=false to preserve existing reference counts during 196 // operations that should yield a net-zero change to descendant refcounts. 197 // When a node is released, its contained files are dereferenced. 198 func (n *node) decRef(contentsToo bool, obsolete *[]*FileBacking) { 199 if n.ref.Add(-1) > 0 { 200 // Other references remain. Can't free. 201 return 202 } 203 204 // Dereference the node's metadata and release child references if 205 // requested. Some internal callers may not want to propagate the deref 206 // because they're manually copying the filemetadata and children to other 207 // nodes, and they want to preserve the existing reference count. 208 if contentsToo { 209 for _, f := range n.items[:n.count] { 210 if f.Unref() == 0 { 211 // There are two sources of node dereferences: tree mutations 212 // and Version dereferences. Files should only be made obsolete 213 // during Version dereferences, during which `obsolete` will be 214 // non-nil. 215 if obsolete == nil { 216 panic(fmt.Sprintf("file metadata %s dereferenced to zero during tree mutation", f.FileNum)) 217 } 218 // Reference counting is performed on the FileBacking. In the case 219 // of a virtual sstable, this reference counting is performed on 220 // a FileBacking which is shared by every single virtual sstable 221 // with the same backing sstable. If the reference count hits 0, 222 // then we know that the FileBacking won't be required by any 223 // sstable in Pebble, and that the backing sstable can be deleted. 224 *obsolete = append(*obsolete, f.FileBacking) 225 } 226 } 227 if !n.leaf { 228 for i := int16(0); i <= n.count; i++ { 229 n.children[i].decRef(true /* contentsToo */, obsolete) 230 } 231 } 232 } 233 } 234 235 // clone creates a clone of the receiver with a single reference count. 236 func (n *node) clone() *node { 237 var c *node 238 if n.leaf { 239 c = newLeafNode() 240 } else { 241 c = newNode() 242 } 243 // NB: copy field-by-field without touching n.ref to avoid 244 // triggering the race detector and looking like a data race. 245 c.count = n.count 246 c.items = n.items 247 c.subtreeCount = n.subtreeCount 248 // Increase the refcount of each contained item. 249 for _, f := range n.items[:n.count] { 250 f.Ref() 251 } 252 if !c.leaf { 253 // Copy children and increase each refcount. 254 c.children = n.children 255 for i := int16(0); i <= c.count; i++ { 256 c.children[i].incRef() 257 } 258 } 259 return c 260 } 261 262 // insertAt inserts the provided file and node at the provided index. This 263 // function is for use only as a helper function for internal B-Tree code. 264 // Clients should not invoke it directly. 265 func (n *node) insertAt(index int, item *FileMetadata, nd *node) { 266 if index < int(n.count) { 267 copy(n.items[index+1:n.count+1], n.items[index:n.count]) 268 if !n.leaf { 269 copy(n.children[index+2:n.count+2], n.children[index+1:n.count+1]) 270 } 271 } 272 n.items[index] = item 273 if !n.leaf { 274 n.children[index+1] = nd 275 } 276 n.count++ 277 } 278 279 // pushBack inserts the provided file and node at the tail of the node's items. 280 // This function is for use only as a helper function for internal B-Tree code. 281 // Clients should not invoke it directly. 282 func (n *node) pushBack(item *FileMetadata, nd *node) { 283 n.items[n.count] = item 284 if !n.leaf { 285 n.children[n.count+1] = nd 286 } 287 n.count++ 288 } 289 290 // pushFront inserts the provided file and node at the head of the 291 // node's items. This function is for use only as a helper function for internal B-Tree 292 // code. Clients should not invoke it directly. 293 func (n *node) pushFront(item *FileMetadata, nd *node) { 294 if !n.leaf { 295 copy(n.children[1:n.count+2], n.children[:n.count+1]) 296 n.children[0] = nd 297 } 298 copy(n.items[1:n.count+1], n.items[:n.count]) 299 n.items[0] = item 300 n.count++ 301 } 302 303 // removeAt removes a value at a given index, pulling all subsequent values 304 // back. This function is for use only as a helper function for internal B-Tree 305 // code. Clients should not invoke it directly. 306 func (n *node) removeAt(index int) (*FileMetadata, *node) { 307 var child *node 308 if !n.leaf { 309 child = n.children[index+1] 310 copy(n.children[index+1:n.count], n.children[index+2:n.count+1]) 311 n.children[n.count] = nil 312 } 313 n.count-- 314 out := n.items[index] 315 copy(n.items[index:n.count], n.items[index+1:n.count+1]) 316 n.items[n.count] = nil 317 return out, child 318 } 319 320 // popBack removes and returns the last element in the list. This function is 321 // for use only as a helper function for internal B-Tree code. Clients should 322 // not invoke it directly. 323 func (n *node) popBack() (*FileMetadata, *node) { 324 n.count-- 325 out := n.items[n.count] 326 n.items[n.count] = nil 327 if n.leaf { 328 return out, nil 329 } 330 child := n.children[n.count+1] 331 n.children[n.count+1] = nil 332 return out, child 333 } 334 335 // popFront removes and returns the first element in the list. This function is 336 // for use only as a helper function for internal B-Tree code. Clients should 337 // not invoke it directly. 338 func (n *node) popFront() (*FileMetadata, *node) { 339 n.count-- 340 var child *node 341 if !n.leaf { 342 child = n.children[0] 343 copy(n.children[:n.count+1], n.children[1:n.count+2]) 344 n.children[n.count+1] = nil 345 } 346 out := n.items[0] 347 copy(n.items[:n.count], n.items[1:n.count+1]) 348 n.items[n.count] = nil 349 return out, child 350 } 351 352 // find returns the index where the given item should be inserted into this 353 // list. 'found' is true if the item already exists in the list at the given 354 // index. 355 // 356 // This function is for use only as a helper function for internal B-Tree code. 357 // Clients should not invoke it directly. 358 func (n *node) find(cmp btreeCmp, item *FileMetadata) (index int, found bool) { 359 // Logic copied from sort.Search. Inlining this gave 360 // an 11% speedup on BenchmarkBTreeDeleteInsert. 361 i, j := 0, int(n.count) 362 for i < j { 363 h := int(uint(i+j) >> 1) // avoid overflow when computing h 364 // i ≤ h < j 365 v := cmp(item, n.items[h]) 366 if v == 0 { 367 return h, true 368 } else if v > 0 { 369 i = h + 1 370 } else { 371 j = h 372 } 373 } 374 return i, false 375 } 376 377 // split splits the given node at the given index. The current node shrinks, 378 // and this function returns the item that existed at that index and a new 379 // node containing all items/children after it. 380 // 381 // split is called when we want to perform a transformation like the one 382 // depicted in the following diagram. 383 // 384 // Before: 385 // +-----------+ 386 // n *node | x y z | 387 // +--/-/-\-\--+ 388 // 389 // After: 390 // +-----------+ 391 // | y | n's parent 392 // +----/-\----+ 393 // / \ 394 // v v 395 // +-----------+ +-----------+ 396 // n *node | x | | z | next *node 397 // +-----------+ +-----------+ 398 // 399 // split does not perform the complete transformation; the caller is responsible 400 // for updating the parent appropriately. split splits `n` into two nodes, `n` 401 // and `next`, returning `next` and the file that separates them. In the diagram 402 // above, `n.split` removes y and z from `n`, returning y in the first return 403 // value and `next` in the second return value. The caller is responsible for 404 // updating n's parent to now contain `y` as the separator between nodes `n` and 405 // `next`. 406 // 407 // This function is for use only as a helper function for internal B-Tree code. 408 // Clients should not invoke it directly. 409 func (n *node) split(i int) (*FileMetadata, *node) { 410 out := n.items[i] 411 var next *node 412 if n.leaf { 413 next = newLeafNode() 414 } else { 415 next = newNode() 416 } 417 next.count = n.count - int16(i+1) 418 copy(next.items[:], n.items[i+1:n.count]) 419 for j := int16(i); j < n.count; j++ { 420 n.items[j] = nil 421 } 422 if !n.leaf { 423 copy(next.children[:], n.children[i+1:n.count+1]) 424 descendantsMoved := 0 425 for j := int16(i + 1); j <= n.count; j++ { 426 descendantsMoved += n.children[j].subtreeCount 427 n.children[j] = nil 428 } 429 n.subtreeCount -= descendantsMoved 430 next.subtreeCount += descendantsMoved 431 } 432 n.count = int16(i) 433 // NB: We subtract one more than `next.count` from n's subtreeCount because 434 // the item at index `i` was removed from `n.items`. We'll return the item 435 // at index `i`, and the caller is responsible for updating the subtree 436 // count of whichever node adopts it. 437 n.subtreeCount -= int(next.count) + 1 438 next.subtreeCount += int(next.count) 439 return out, next 440 } 441 442 // Insert inserts a item into the subtree rooted at this node, making sure no 443 // nodes in the subtree exceed maxItems items. 444 func (n *node) Insert(cmp btreeCmp, item *FileMetadata) error { 445 i, found := n.find(cmp, item) 446 if found { 447 // cmp provides a total ordering of the files within a level. 448 // If we're inserting a metadata that's equal to an existing item 449 // in the tree, we're inserting a file into a level twice. 450 return errors.Errorf("files %s and %s collided on sort keys", 451 errors.Safe(item.FileNum), errors.Safe(n.items[i].FileNum)) 452 } 453 if n.leaf { 454 n.insertAt(i, item, nil) 455 n.subtreeCount++ 456 return nil 457 } 458 if n.children[i].count >= maxItems { 459 splitLa, splitNode := mut(&n.children[i]).split(maxItems / 2) 460 n.insertAt(i, splitLa, splitNode) 461 462 switch cmp := cmp(item, n.items[i]); { 463 case cmp < 0: 464 // no change, we want first split node 465 case cmp > 0: 466 i++ // we want second split node 467 default: 468 // cmp provides a total ordering of the files within a level. 469 // If we're inserting a metadata that's equal to an existing item 470 // in the tree, we're inserting a file into a level twice. 471 return errors.Errorf("files %s and %s collided on sort keys", 472 errors.Safe(item.FileNum), errors.Safe(n.items[i].FileNum)) 473 } 474 } 475 476 err := mut(&n.children[i]).Insert(cmp, item) 477 if err == nil { 478 n.subtreeCount++ 479 } 480 return err 481 } 482 483 // removeMax removes and returns the maximum item from the subtree rooted at 484 // this node. This function is for use only as a helper function for internal 485 // B-Tree code. Clients should not invoke it directly. 486 func (n *node) removeMax() *FileMetadata { 487 if n.leaf { 488 n.count-- 489 n.subtreeCount-- 490 out := n.items[n.count] 491 n.items[n.count] = nil 492 return out 493 } 494 child := mut(&n.children[n.count]) 495 if child.count <= minItems { 496 n.rebalanceOrMerge(int(n.count)) 497 return n.removeMax() 498 } 499 n.subtreeCount-- 500 return child.removeMax() 501 } 502 503 // Remove removes a item from the subtree rooted at this node. Returns 504 // the item that was removed or nil if no matching item was found. 505 func (n *node) Remove(cmp btreeCmp, item *FileMetadata) (out *FileMetadata) { 506 i, found := n.find(cmp, item) 507 if n.leaf { 508 if found { 509 out, _ = n.removeAt(i) 510 n.subtreeCount-- 511 return out 512 } 513 return nil 514 } 515 if n.children[i].count <= minItems { 516 // Child not large enough to remove from. 517 n.rebalanceOrMerge(i) 518 return n.Remove(cmp, item) 519 } 520 child := mut(&n.children[i]) 521 if found { 522 // Replace the item being removed with the max item in our left child. 523 out = n.items[i] 524 n.items[i] = child.removeMax() 525 n.subtreeCount-- 526 return out 527 } 528 // File is not in this node and child is large enough to remove from. 529 out = child.Remove(cmp, item) 530 if out != nil { 531 n.subtreeCount-- 532 } 533 return out 534 } 535 536 // rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove a 537 // item from it while keeping it at or above minItems. This function is for use 538 // only as a helper function for internal B-Tree code. Clients should not invoke 539 // it directly. 540 func (n *node) rebalanceOrMerge(i int) { 541 switch { 542 case i > 0 && n.children[i-1].count > minItems: 543 // Rebalance from left sibling. 544 // 545 // +-----------+ 546 // | y | 547 // +----/-\----+ 548 // / \ 549 // v v 550 // +-----------+ +-----------+ 551 // | x | | | 552 // +----------\+ +-----------+ 553 // \ 554 // v 555 // a 556 // 557 // After: 558 // 559 // +-----------+ 560 // | x | 561 // +----/-\----+ 562 // / \ 563 // v v 564 // +-----------+ +-----------+ 565 // | | | y | 566 // +-----------+ +/----------+ 567 // / 568 // v 569 // a 570 // 571 left := mut(&n.children[i-1]) 572 child := mut(&n.children[i]) 573 xLa, grandChild := left.popBack() 574 yLa := n.items[i-1] 575 child.pushFront(yLa, grandChild) 576 n.items[i-1] = xLa 577 child.subtreeCount++ 578 left.subtreeCount-- 579 if grandChild != nil { 580 child.subtreeCount += grandChild.subtreeCount 581 left.subtreeCount -= grandChild.subtreeCount 582 } 583 584 case i < int(n.count) && n.children[i+1].count > minItems: 585 // Rebalance from right sibling. 586 // 587 // +-----------+ 588 // | y | 589 // +----/-\----+ 590 // / \ 591 // v v 592 // +-----------+ +-----------+ 593 // | | | x | 594 // +-----------+ +/----------+ 595 // / 596 // v 597 // a 598 // 599 // After: 600 // 601 // +-----------+ 602 // | x | 603 // +----/-\----+ 604 // / \ 605 // v v 606 // +-----------+ +-----------+ 607 // | y | | | 608 // +----------\+ +-----------+ 609 // \ 610 // v 611 // a 612 // 613 right := mut(&n.children[i+1]) 614 child := mut(&n.children[i]) 615 xLa, grandChild := right.popFront() 616 yLa := n.items[i] 617 child.pushBack(yLa, grandChild) 618 child.subtreeCount++ 619 right.subtreeCount-- 620 if grandChild != nil { 621 child.subtreeCount += grandChild.subtreeCount 622 right.subtreeCount -= grandChild.subtreeCount 623 } 624 n.items[i] = xLa 625 626 default: 627 // Merge with either the left or right sibling. 628 // 629 // +-----------+ 630 // | u y v | 631 // +----/-\----+ 632 // / \ 633 // v v 634 // +-----------+ +-----------+ 635 // | x | | z | 636 // +-----------+ +-----------+ 637 // 638 // After: 639 // 640 // +-----------+ 641 // | u v | 642 // +-----|-----+ 643 // | 644 // v 645 // +-----------+ 646 // | x y z | 647 // +-----------+ 648 // 649 if i >= int(n.count) { 650 i = int(n.count - 1) 651 } 652 child := mut(&n.children[i]) 653 // Make mergeChild mutable, bumping the refcounts on its children if necessary. 654 _ = mut(&n.children[i+1]) 655 mergeLa, mergeChild := n.removeAt(i) 656 child.items[child.count] = mergeLa 657 copy(child.items[child.count+1:], mergeChild.items[:mergeChild.count]) 658 if !child.leaf { 659 copy(child.children[child.count+1:], mergeChild.children[:mergeChild.count+1]) 660 } 661 child.count += mergeChild.count + 1 662 child.subtreeCount += mergeChild.subtreeCount + 1 663 664 mergeChild.decRef(false /* contentsToo */, nil) 665 } 666 } 667 668 // InvalidateAnnotation removes any existing cached annotations for the provided 669 // annotator from this node's subtree. 670 func (n *node) InvalidateAnnotation(a Annotator) { 671 // Find this annotator's annotation on this node. 672 var annot *annotation 673 for i := range n.annot { 674 if n.annot[i].annotator == a { 675 annot = &n.annot[i] 676 } 677 } 678 679 if annot != nil && annot.valid { 680 annot.valid = false 681 annot.v = a.Zero(annot.v) 682 } 683 if !n.leaf { 684 for i := int16(0); i <= n.count; i++ { 685 n.children[i].InvalidateAnnotation(a) 686 } 687 } 688 } 689 690 // Annotation retrieves, computing if not already computed, the provided 691 // annotator's annotation of this node. The second return value indicates 692 // whether the future reads of this annotation may use the first return value 693 // as-is. If false, the annotation is not stable and may change on a subsequent 694 // computation. 695 func (n *node) Annotation(a Annotator) (interface{}, bool) { 696 // Find this annotator's annotation on this node. 697 var annot *annotation 698 for i := range n.annot { 699 if n.annot[i].annotator == a { 700 annot = &n.annot[i] 701 } 702 } 703 704 // If it exists and is marked as valid, we can return it without 705 // recomputing anything. 706 if annot != nil && annot.valid { 707 return annot.v, true 708 } 709 710 if annot == nil { 711 // This is n's first time being annotated by a. 712 // Create a new zeroed annotation. 713 n.annot = append(n.annot, annotation{ 714 annotator: a, 715 v: a.Zero(nil), 716 }) 717 annot = &n.annot[len(n.annot)-1] 718 } else { 719 // There's an existing annotation that must be recomputed. 720 // Zero its value. 721 annot.v = a.Zero(annot.v) 722 } 723 724 annot.valid = true 725 for i := int16(0); i <= n.count; i++ { 726 if !n.leaf { 727 v, ok := n.children[i].Annotation(a) 728 annot.v = a.Merge(v, annot.v) 729 annot.valid = annot.valid && ok 730 } 731 if i < n.count { 732 v, ok := a.Accumulate(n.items[i], annot.v) 733 annot.v = v 734 annot.valid = annot.valid && ok 735 } 736 } 737 return annot.v, annot.valid 738 } 739 740 func (n *node) verifyInvariants() { 741 recomputedSubtreeCount := int(n.count) 742 if !n.leaf { 743 for i := int16(0); i <= n.count; i++ { 744 n.children[i].verifyInvariants() 745 recomputedSubtreeCount += n.children[i].subtreeCount 746 } 747 } 748 if recomputedSubtreeCount != n.subtreeCount { 749 panic(fmt.Sprintf("recomputed subtree count (%d) ≠ n.subtreeCount (%d)", 750 recomputedSubtreeCount, n.subtreeCount)) 751 } 752 } 753 754 // btree is an implementation of a B-Tree. 755 // 756 // btree stores FileMetadata in an ordered structure, allowing easy insertion, 757 // removal, and iteration. The B-Tree stores items in order based on cmp. The 758 // first level of the LSM uses a cmp function that compares sequence numbers. 759 // All other levels compare using the FileMetadata.Smallest. 760 // 761 // Write operations are not safe for concurrent mutation by multiple 762 // goroutines, but Read operations are. 763 type btree struct { 764 root *node 765 cmp btreeCmp 766 } 767 768 // Release dereferences and clears the root node of the btree, removing all 769 // items from the btree. In doing so, it decrements contained file counts. 770 // It returns a slice of newly obsolete backing files, if any. 771 func (t *btree) Release() (obsolete []*FileBacking) { 772 if t.root != nil { 773 t.root.decRef(true /* contentsToo */, &obsolete) 774 t.root = nil 775 } 776 return obsolete 777 } 778 779 // Clone clones the btree, lazily. It does so in constant time. 780 func (t *btree) Clone() btree { 781 c := *t 782 if c.root != nil { 783 // Incrementing the reference count on the root node is sufficient to 784 // ensure that no node in the cloned tree can be mutated by an actor 785 // holding a reference to the original tree and vice versa. This 786 // property is upheld because the root node in the receiver btree and 787 // the returned btree will both necessarily have a reference count of at 788 // least 2 when this method returns. All tree mutations recursively 789 // acquire mutable node references (see mut) as they traverse down the 790 // tree. The act of acquiring a mutable node reference performs a clone 791 // if a node's reference count is greater than one. Cloning a node (see 792 // clone) increases the reference count on each of its children, 793 // ensuring that they have a reference count of at least 2. This, in 794 // turn, ensures that any of the child nodes that are modified will also 795 // be copied-on-write, recursively ensuring the immutability property 796 // over the entire tree. 797 c.root.incRef() 798 } 799 return c 800 } 801 802 // Delete removes the provided file from the tree. 803 // It returns true if the file now has a zero reference count. 804 func (t *btree) Delete(item *FileMetadata) (obsolete bool) { 805 if t.root == nil || t.root.count == 0 { 806 return false 807 } 808 if out := mut(&t.root).Remove(t.cmp, item); out != nil { 809 obsolete = out.Unref() == 0 810 } 811 if invariants.Enabled { 812 t.root.verifyInvariants() 813 } 814 if t.root.count == 0 { 815 old := t.root 816 if t.root.leaf { 817 t.root = nil 818 } else { 819 t.root = t.root.children[0] 820 } 821 old.decRef(false /* contentsToo */, nil) 822 } 823 return obsolete 824 } 825 826 // Insert adds the given item to the tree. If a item in the tree already 827 // equals the given one, Insert panics. 828 func (t *btree) Insert(item *FileMetadata) error { 829 if t.root == nil { 830 t.root = newLeafNode() 831 } else if t.root.count >= maxItems { 832 splitLa, splitNode := mut(&t.root).split(maxItems / 2) 833 newRoot := newNode() 834 newRoot.count = 1 835 newRoot.items[0] = splitLa 836 newRoot.children[0] = t.root 837 newRoot.children[1] = splitNode 838 newRoot.subtreeCount = t.root.subtreeCount + splitNode.subtreeCount + 1 839 t.root = newRoot 840 } 841 item.Ref() 842 err := mut(&t.root).Insert(t.cmp, item) 843 if invariants.Enabled { 844 t.root.verifyInvariants() 845 } 846 return err 847 } 848 849 // Iter returns a new iterator object. It is not safe to continue using an 850 // iterator after modifications are made to the tree. If modifications are made, 851 // create a new iterator. 852 func (t *btree) Iter() iterator { 853 return iterator{r: t.root, pos: -1, cmp: t.cmp} 854 } 855 856 // Count returns the number of files contained within the B-Tree. 857 func (t *btree) Count() int { 858 if t.root == nil { 859 return 0 860 } 861 return t.root.subtreeCount 862 } 863 864 // String returns a string description of the tree. The format is 865 // similar to the https://en.wikipedia.org/wiki/Newick_format. 866 func (t *btree) String() string { 867 if t.Count() == 0 { 868 return ";" 869 } 870 var b strings.Builder 871 t.root.writeString(&b) 872 return b.String() 873 } 874 875 func (n *node) writeString(b *strings.Builder) { 876 if n.leaf { 877 for i := int16(0); i < n.count; i++ { 878 if i != 0 { 879 b.WriteString(",") 880 } 881 b.WriteString(n.items[i].String()) 882 } 883 return 884 } 885 for i := int16(0); i <= n.count; i++ { 886 b.WriteString("(") 887 n.children[i].writeString(b) 888 b.WriteString(")") 889 if i < n.count { 890 b.WriteString(n.items[i].String()) 891 } 892 } 893 } 894 895 // iterStack represents a stack of (node, pos) tuples, which captures 896 // iteration state as an iterator descends a btree. 897 type iterStack struct { 898 // a contains aLen stack frames when an iterator stack is short enough. 899 // If the iterator stack overflows the capacity of iterStackArr, the stack 900 // is moved to s and aLen is set to -1. 901 a iterStackArr 902 aLen int16 // -1 when using s 903 s []iterFrame 904 } 905 906 // Used to avoid allocations for stacks below a certain size. 907 type iterStackArr [3]iterFrame 908 909 type iterFrame struct { 910 n *node 911 pos int16 912 } 913 914 func (is *iterStack) push(f iterFrame) { 915 if is.aLen == -1 { 916 is.s = append(is.s, f) 917 } else if int(is.aLen) == len(is.a) { 918 is.s = make([]iterFrame, int(is.aLen)+1, 2*int(is.aLen)) 919 copy(is.s, is.a[:]) 920 is.s[int(is.aLen)] = f 921 is.aLen = -1 922 } else { 923 is.a[is.aLen] = f 924 is.aLen++ 925 } 926 } 927 928 func (is *iterStack) pop() iterFrame { 929 if is.aLen == -1 { 930 f := is.s[len(is.s)-1] 931 is.s = is.s[:len(is.s)-1] 932 return f 933 } 934 is.aLen-- 935 return is.a[is.aLen] 936 } 937 938 func (is *iterStack) len() int { 939 if is.aLen == -1 { 940 return len(is.s) 941 } 942 return int(is.aLen) 943 } 944 945 func (is *iterStack) clone() iterStack { 946 // If the iterator is using the embedded iterStackArr, we only need to 947 // copy the struct itself. 948 if is.s == nil { 949 return *is 950 } 951 clone := *is 952 clone.s = make([]iterFrame, len(is.s)) 953 copy(clone.s, is.s) 954 return clone 955 } 956 957 func (is *iterStack) nth(n int) (f iterFrame, ok bool) { 958 if is.aLen == -1 { 959 if n >= len(is.s) { 960 return f, false 961 } 962 return is.s[n], true 963 } 964 if int16(n) >= is.aLen { 965 return f, false 966 } 967 return is.a[n], true 968 } 969 970 func (is *iterStack) reset() { 971 if is.aLen == -1 { 972 is.s = is.s[:0] 973 } else { 974 is.aLen = 0 975 } 976 } 977 978 // iterator is responsible for search and traversal within a btree. 979 type iterator struct { 980 // the root node of the B-Tree. 981 r *node 982 // n and pos make up the current position of the iterator. 983 // If valid, n.items[pos] is the current value of the iterator. 984 // 985 // n may be nil iff i.r is nil. 986 n *node 987 pos int16 988 // cmp dictates the ordering of the FileMetadata. 989 cmp func(*FileMetadata, *FileMetadata) int 990 // a stack of n's ancestors within the B-Tree, alongside the position 991 // taken to arrive at n. If non-empty, the bottommost frame of the stack 992 // will always contain the B-Tree root. 993 s iterStack 994 } 995 996 // countLeft returns the count of files that are to the left of the current 997 // iterator position. 998 func (i *iterator) countLeft() int { 999 if i.r == nil { 1000 return 0 1001 } 1002 1003 // Each iterator has a stack of frames marking the path from the root node 1004 // to the current iterator position. All files (n.items) and all subtrees 1005 // (n.children) with indexes less than [pos] are to the left of the current 1006 // iterator position. 1007 // 1008 // +------------------------+ - 1009 // | Root pos:5 | | 1010 // +------------------------+ | stack 1011 // | Root/5 pos:3 | | frames 1012 // +------------------------+ | [i.s] 1013 // | Root/5/3 pos:9 | | 1014 // +========================+ - 1015 // | | 1016 // | i.n: Root/5/3/9 i.pos:2| 1017 // +------------------------+ 1018 // 1019 var count int 1020 // Walk all the ancestors in the iterator stack [i.s], tallying up all the 1021 // files and subtrees to the left of the stack frame's position. 1022 f, ok := i.s.nth(0) 1023 for fi := 0; ok; fi++ { 1024 // There are [f.pos] files contained within [f.n.items] that sort to the 1025 // left of the subtree the iterator has descended. 1026 count += int(f.pos) 1027 // Any subtrees that fall before the stack frame's position are entirely 1028 // to the left of the iterator's current position. 1029 for j := int16(0); j < f.pos; j++ { 1030 count += f.n.children[j].subtreeCount 1031 } 1032 f, ok = i.s.nth(fi + 1) 1033 } 1034 1035 // The bottommost stack frame is inlined within the iterator struct. Again, 1036 // [i.pos] files fall to the left of the current iterator position. 1037 count += int(i.pos) 1038 if !i.n.leaf { 1039 // NB: Unlike above, we use a `<= i.pos` comparison. The iterator is 1040 // positioned at item `i.n.items[i.pos]`, which sorts after everything 1041 // in the subtree at `i.n.children[i.pos]`. 1042 for j := int16(0); j <= i.pos; j++ { 1043 count += i.n.children[j].subtreeCount 1044 } 1045 } 1046 return count 1047 } 1048 1049 func (i *iterator) clone() iterator { 1050 c := *i 1051 c.s = i.s.clone() 1052 return c 1053 } 1054 1055 func (i *iterator) reset() { 1056 i.n = i.r 1057 i.pos = -1 1058 i.s.reset() 1059 } 1060 1061 func (i iterator) String() string { 1062 var buf bytes.Buffer 1063 for n := 0; ; n++ { 1064 f, ok := i.s.nth(n) 1065 if !ok { 1066 break 1067 } 1068 fmt.Fprintf(&buf, "%p: %02d/%02d\n", f.n, f.pos, f.n.count) 1069 } 1070 if i.r == nil { 1071 fmt.Fprintf(&buf, "<nil>: %02d", i.pos) 1072 } else { 1073 fmt.Fprintf(&buf, "%p: %02d/%02d", i.n, i.pos, i.n.count) 1074 } 1075 return buf.String() 1076 } 1077 1078 func cmpIter(a, b iterator) int { 1079 if a.r != b.r { 1080 panic("compared iterators from different btrees") 1081 } 1082 1083 // Each iterator has a stack of frames marking the path from the root node 1084 // to the current iterator position. We walk both paths formed by the 1085 // iterators' stacks simultaneously, descending from the shared root node, 1086 // always comparing nodes at the same level in the tree. 1087 // 1088 // If the iterators' paths ever diverge and point to different nodes, the 1089 // iterators are not equal and we use the node positions to evaluate the 1090 // comparison. 1091 // 1092 // If an iterator's stack ends, we stop descending and use its current 1093 // node and position for the final comparison. One iterator's stack may 1094 // end before another's if one iterator is positioned deeper in the tree. 1095 // 1096 // a b 1097 // +------------------------+ +--------------------------+ - 1098 // | Root pos:5 | = | Root pos:5 | | 1099 // +------------------------+ +--------------------------+ | stack 1100 // | Root/5 pos:3 | = | Root/5 pos:3 | | frames 1101 // +------------------------+ +--------------------------+ | 1102 // | Root/5/3 pos:9 | > | Root/5/3 pos:1 | | 1103 // +========================+ +==========================+ - 1104 // | | | | 1105 // | a.n: Root/5/3/9 a.pos:2| | b.n: Root/5/3/1, b.pos:5 | 1106 // +------------------------+ +--------------------------+ 1107 1108 // Initialize with the iterator's current node and position. These are 1109 // conceptually the most-recent/current frame of the iterator stack. 1110 an, apos := a.n, a.pos 1111 bn, bpos := b.n, b.pos 1112 1113 // aok, bok are set while traversing the iterator's path down the B-Tree. 1114 // They're declared in the outer scope because they help distinguish the 1115 // sentinel case when both iterators' first frame points to the last child 1116 // of the root. If an iterator has no other frames in its stack, it's the 1117 // end sentinel state which sorts after everything else. 1118 var aok, bok bool 1119 for i := 0; ; i++ { 1120 var af, bf iterFrame 1121 af, aok = a.s.nth(i) 1122 bf, bok = b.s.nth(i) 1123 if !aok || !bok { 1124 if aok { 1125 // Iterator a, unlike iterator b, still has a frame. Set an, 1126 // apos so we compare using the frame from the stack. 1127 an, apos = af.n, af.pos 1128 } 1129 if bok { 1130 // Iterator b, unlike iterator a, still has a frame. Set bn, 1131 // bpos so we compare using the frame from the stack. 1132 bn, bpos = bf.n, bf.pos 1133 } 1134 break 1135 } 1136 1137 // aok && bok 1138 if af.n != bf.n { 1139 panic("nonmatching nodes during btree iterator comparison") 1140 } 1141 switch { 1142 case af.pos < bf.pos: 1143 return -1 1144 case af.pos > bf.pos: 1145 return +1 1146 default: 1147 // Continue up both iterators' stacks (equivalently, down the 1148 // B-Tree away from the root). 1149 } 1150 } 1151 1152 if aok && bok { 1153 panic("expected one or more stacks to have been exhausted") 1154 } 1155 if an != bn { 1156 panic("nonmatching nodes during btree iterator comparison") 1157 } 1158 switch { 1159 case apos < bpos: 1160 return -1 1161 case apos > bpos: 1162 return +1 1163 default: 1164 switch { 1165 case aok: 1166 // a is positioned at a leaf child at this position and b is at an 1167 // end sentinel state. 1168 return -1 1169 case bok: 1170 // b is positioned at a leaf child at this position and a is at an 1171 // end sentinel state. 1172 return +1 1173 default: 1174 return 0 1175 } 1176 } 1177 } 1178 1179 func (i *iterator) descend(n *node, pos int16) { 1180 i.s.push(iterFrame{n: n, pos: pos}) 1181 i.n = n.children[pos] 1182 i.pos = 0 1183 } 1184 1185 // ascend ascends up to the current node's parent and resets the position 1186 // to the one previously set for this parent node. 1187 func (i *iterator) ascend() { 1188 f := i.s.pop() 1189 i.n = f.n 1190 i.pos = f.pos 1191 } 1192 1193 // seek repositions the iterator over the first file for which fn returns 1194 // true, mirroring the semantics of the standard library's sort.Search 1195 // function. Like sort.Search, seek requires the iterator's B-Tree to be 1196 // ordered such that fn returns false for some (possibly empty) prefix of the 1197 // tree's files, and then true for the (possibly empty) remainder. 1198 func (i *iterator) seek(fn func(*FileMetadata) bool) { 1199 i.reset() 1200 if i.r == nil { 1201 return 1202 } 1203 1204 for { 1205 // Logic copied from sort.Search. 1206 j, k := 0, int(i.n.count) 1207 for j < k { 1208 h := int(uint(j+k) >> 1) // avoid overflow when computing h 1209 1210 // j ≤ h < k 1211 if !fn(i.n.items[h]) { 1212 j = h + 1 // preserves f(j-1) == false 1213 } else { 1214 k = h // preserves f(k) == true 1215 } 1216 } 1217 1218 i.pos = int16(j) 1219 if i.n.leaf { 1220 if i.pos == i.n.count { 1221 i.next() 1222 } 1223 return 1224 } 1225 i.descend(i.n, i.pos) 1226 } 1227 } 1228 1229 // first seeks to the first item in the btree. 1230 func (i *iterator) first() { 1231 i.reset() 1232 if i.r == nil { 1233 return 1234 } 1235 for !i.n.leaf { 1236 i.descend(i.n, 0) 1237 } 1238 i.pos = 0 1239 } 1240 1241 // last seeks to the last item in the btree. 1242 func (i *iterator) last() { 1243 i.reset() 1244 if i.r == nil { 1245 return 1246 } 1247 for !i.n.leaf { 1248 i.descend(i.n, i.n.count) 1249 } 1250 i.pos = i.n.count - 1 1251 } 1252 1253 // next positions the iterator to the item immediately following 1254 // its current position. 1255 func (i *iterator) next() { 1256 if i.r == nil { 1257 return 1258 } 1259 1260 if i.n.leaf { 1261 if i.pos < i.n.count { 1262 i.pos++ 1263 } 1264 if i.pos < i.n.count { 1265 return 1266 } 1267 for i.s.len() > 0 && i.pos >= i.n.count { 1268 i.ascend() 1269 } 1270 return 1271 } 1272 1273 i.descend(i.n, i.pos+1) 1274 for !i.n.leaf { 1275 i.descend(i.n, 0) 1276 } 1277 i.pos = 0 1278 } 1279 1280 // prev positions the iterator to the item immediately preceding 1281 // its current position. 1282 func (i *iterator) prev() { 1283 if i.r == nil { 1284 return 1285 } 1286 1287 if i.n.leaf { 1288 i.pos-- 1289 if i.pos >= 0 { 1290 return 1291 } 1292 for i.s.len() > 0 && i.pos < 0 { 1293 i.ascend() 1294 i.pos-- 1295 } 1296 return 1297 } 1298 1299 i.descend(i.n, i.pos) 1300 for !i.n.leaf { 1301 i.descend(i.n, i.n.count) 1302 } 1303 i.pos = i.n.count - 1 1304 } 1305 1306 // valid returns whether the iterator is positioned at a valid position. 1307 func (i *iterator) valid() bool { 1308 return i.r != nil && i.pos >= 0 && i.pos < i.n.count 1309 } 1310 1311 // cur returns the item at the iterator's current position. It is illegal 1312 // to call cur if the iterator is not valid. 1313 func (i *iterator) cur() *FileMetadata { 1314 if invariants.Enabled && !i.valid() { 1315 panic("btree iterator.cur invoked on invalid iterator") 1316 } 1317 return i.n.items[i.pos] 1318 }