From 797235739fd72c013f22ffa791decd3f0a816adb Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Mon, 29 Oct 2018 23:57:43 -0400
Subject: [PATCH] util/interval: Add a Clear method that clears all nodes into
 a freelist

google/btree#21

Release note: None
---
 pkg/util/interval/btree_based_interval.go | 84 ++++++++++++++++++++---
 1 file changed, 74 insertions(+), 10 deletions(-)

diff --git a/pkg/util/interval/btree_based_interval.go b/pkg/util/interval/btree_based_interval.go
index 84a33d4cc124..e642f5f1b115 100644
--- a/pkg/util/interval/btree_based_interval.go
+++ b/pkg/util/interval/btree_based_interval.go
@@ -12,7 +12,7 @@
 // implied. See the License for the specific language governing
 // permissions and limitations under the License.
 //
-// This code is based on: https://github.com/google/btree
+// This code is based on: https://github.com/google/btree.
 
 package interval
 
@@ -66,12 +66,16 @@ func (f *FreeList) newNode() (n *node) {
 	return
 }
 
-func (f *FreeList) freeNode(n *node) {
+// freeNode adds the given node to the list, returning true if it was added
+// and false if it was discarded.
+func (f *FreeList) freeNode(n *node) (out bool) {
 	f.mu.Lock()
 	if len(f.freelist) < cap(f.freelist) {
 		f.freelist = append(f.freelist, n)
+		out = true
 	}
 	f.mu.Unlock()
+	return
 }
 
 // newBTree creates a new interval tree with the given overlapper function and
@@ -930,12 +934,16 @@ func (t *btree) AdjustRanges() {
 	if t.isEmpty() {
 		return
 	}
-	t.root.adjustRanges()
+	t.root.adjustRanges(t.root.cow)
 }
 
-func (n *node) adjustRanges() {
-	for _, c := range n.children {
-		c.adjustRanges()
+func (n *node) adjustRanges(c *copyOnWriteContext) {
+	if n.cow != c {
+		// Could not have been modified.
+		return
+	}
+	for _, child := range n.children {
+		child.adjustRanges(c)
 	}
 	n.adjustRange()
 }
@@ -957,14 +965,29 @@ func (c *copyOnWriteContext) newNode() (n *node) {
 	return
 }
 
-func (c *copyOnWriteContext) freeNode(n *node) {
+type freeType int
+
+const (
+	ftFreelistFull freeType = iota // node was freed (available for GC, not stored in freelist)
+	ftStored                       // node was stored in the freelist for later use
+	ftNotOwned                     // node was ignored by COW, since it's owned by another one
+)
+
+// freeNode frees a node within a given COW context, if it's owned by that
+// context.  It returns what happened to the node (see freeType const
+// documentation).
+func (c *copyOnWriteContext) freeNode(n *node) freeType {
 	if n.cow == c {
 		// clear to allow GC
 		n.items.truncate(0)
 		n.children.truncate(0)
 		n.cow = nil // clear to allow GC
-		c.freelist.freeNode(n)
+		if c.freelist.freeNode(n) {
+			return ftStored
+		}
+		return ftFreelistFull
 	}
+	return ftNotOwned
 }
 
 func (t *btree) Insert(e Interface, fast bool) (err error) {
@@ -1075,7 +1098,48 @@ func (t *btree) Iterator() TreeIterator {
 	return &ti
 }
 
+// Clear removes all items from the btree.  If addNodesToFreelist is true,
+// t's nodes are added to its freelist as part of this call, until the freelist
+// is full.  Otherwise, the root node is simply dereferenced and the subtree
+// left to Go's normal GC processes.
+//
+// This can be much faster
+// than calling Delete on all elements, because that requires finding/removing
+// each element in the tree and updating the tree accordingly.  It also is
+// somewhat faster than creating a new tree to replace the old one, because
+// nodes from the old tree are reclaimed into the freelist for use by the new
+// one, instead of being lost to the garbage collector.
+//
+// This call takes:
+//   O(1): when addNodesToFreelist is false, this is a single operation.
+//   O(1): when the freelist is already full, it breaks out immediately
+//   O(freelist size):  when the freelist is empty and the nodes are all owned
+//       by this tree, nodes are added to the freelist until full.
+//   O(tree size):  when all nodes are owned by another tree, all nodes are
+//       iterated over looking for nodes to add to the freelist, and due to
+//       ownership, none are.
+func (t *btree) ClearWithOpt(addNodesToFreelist bool) {
+	if t.root != nil && addNodesToFreelist {
+		t.root.reset(t.cow)
+	}
+	t.root, t.length = nil, 0
+}
+
 func (t *btree) Clear() {
-	t.root = nil
-	t.length = 0
+	t.ClearWithOpt(true)
+}
+
+// reset returns a subtree to the freelist.  It breaks out immediately if the
+// freelist is full, since the only benefit of iterating is to fill that
+// freelist up.  Returns true if parent reset call should continue.
+func (n *node) reset(c *copyOnWriteContext) bool {
+	if n.cow != c {
+		return false
+	}
+	for _, child := range n.children {
+		if !child.reset(c) {
+			return false
+		}
+	}
+	return c.freeNode(n) != ftFreelistFull
 }