Use z.Buffer backing for B+ tree

z/btree.go

@@ -38,6 +38,7 @@ var (
 // Tree represents the structure for custom mmaped B+ tree.
 // It supports keys in range [1, math.MaxUint64-1] and values [1, math.Uint64].
 type Tree struct {
+	buffer   *Buffer
 	data     []byte
 	nextPage uint64
 	freePage uint64
@@ -51,22 +52,48 @@ func (t *Tree) initRootNode() {
 	t.Set(absoluteMax, 0)
 }
 
-// NewTree returns a memory mapped B+ tree with given filename.
-func NewTree() *Tree {
-	t := &Tree{}
+// NewTree returns an in-memory B+ tree.
+func NewTree(tag string) *Tree {
+	const defaultTag = "tree"
+	if tag == "" {
+		tag = defaultTag
+	}
+	t := &Tree{buffer: NewBuffer(minSize, tag)}
 	t.Reset()
 	return t
 }
 
+// NewTree returns a persistent on-disk B+ tree.
+func NewTreePersistent(path string) (*Tree, error) {
+	buffer, err := NewBufferPersistent(path)
+	if err != nil {
+		return nil, err
+	}
+	tree := &Tree{buffer: buffer}
+	tree.Reset()
+	return tree, nil
+}
+
 // Reset resets the tree and truncates it to maxSz.
 func (t *Tree) Reset() {
 	t.nextPage = 1
 	t.freePage = 0
-	t.data = make([]byte, minSize)
+	Memclr(t.buffer.buf)
+	t.buffer.Reset()
+	t.buffer.AllocateOffset(minSize)
+	t.data = t.buffer.Bytes()
 	t.stats = TreeStats{}
 	t.initRootNode()
 }
 
+// Close releases the memory used by the tree.
+func (t *Tree) Close() error {
+	if t == nil {
+		return nil
+	}
+	return t.buffer.Release()
+}
+
 type TreeStats struct {
 	Allocated    int     // Derived.
 	Bytes        int     // Derived.
@@ -82,7 +109,7 @@ func (t *Tree) Stats() TreeStats {
 	numPages := int(t.nextPage - 1)
 	out := TreeStats{
 		Bytes:        numPages * pageSize,
-		Allocated:    cap(t.data),
+		Allocated:    len(t.data),
 		NumLeafKeys:  t.stats.NumLeafKeys,
 		NumPages:     numPages,
 		NumPagesFree: t.stats.NumPagesFree,
@@ -114,16 +141,10 @@ func (t *Tree) newNode(bit uint64) node {
 		pageId = t.nextPage
 		t.nextPage++
 		offset := int(pageId) * pageSize
-		// Double the size with an upper cap of 1GB, if current buffer is insufficient.
-		if offset+pageSize > len(t.data) {
-			const oneGB = 1 << 30
-			newSz := 2 * len(t.data)
-			if newSz > len(t.data)+oneGB {
-				newSz = len(t.data) + oneGB
-			}
-			out := make([]byte, newSz)
-			copy(out, t.data)
-			t.data = out
+		reqSize := offset + pageSize
+		if reqSize > len(t.data) {
+			t.buffer.AllocateOffset(reqSize - len(t.data))
+			t.data = t.buffer.Bytes()
 		}
 	}
 	n := t.node(pageId)

z/btree_test.go

@@ -39,7 +39,8 @@ func setPageSize(sz int) {
 }
 
 func TestTree(t *testing.T) {
-	bt := NewTree()
+	bt := NewTree("TestTree")
+	defer func() { require.NoError(t, bt.Close()) }()
 
 	N := uint64(256 * 256)
 	for i := uint64(1); i < N; i++ {
@@ -60,7 +61,8 @@ func TestTree(t *testing.T) {
 
 func TestTreeBasic(t *testing.T) {
 	setAndGet := func() {
-		bt := NewTree()
+		bt := NewTree("TestTreeBasic")
+		defer func() { require.NoError(t, bt.Close()) }()
 
 		N := uint64(1 << 20)
 		mp := make(map[uint64]uint64)
@@ -83,7 +85,9 @@ func TestTreeBasic(t *testing.T) {
 }
 
 func TestTreeReset(t *testing.T) {
-	bt := NewTree()
+	bt := NewTree("TestTreeReset")
+	defer func() { require.NoError(t, bt.Close()) }()
+
 	N := 1 << 10
 	val := rand.Uint64()
 	for i := 0; i < N; i++ {
@@ -114,7 +118,9 @@ func TestTreeReset(t *testing.T) {
 }
 
 func TestTreeCycle(t *testing.T) {
-	bt := NewTree()
+	bt := NewTree("TestTreeCycle")
+	defer func() { require.NoError(t, bt.Close()) }()
+
 	val := uint64(0)
 	for i := 0; i < 16; i++ {
 		for j := 0; j < 1e6+i*1e4; j++ {
@@ -135,7 +141,8 @@ func TestTreeCycle(t *testing.T) {
 }
 
 func TestTreeIterateKV(t *testing.T) {
-	bt := NewTree()
+	bt := NewTree("TestTreeIterateKV")
+	defer func() { require.NoError(t, bt.Close()) }()
 
 	// Set entries: (i, i*10)
 	const n = uint64(1 << 20)
@@ -169,7 +176,8 @@ func TestOccupancyRatio(t *testing.T) {
 	defer setPageSize(os.Getpagesize())
 	require.Equal(t, 4, maxKeys)
 
-	bt := NewTree()
+	bt := NewTree("TestOccupancyRatio")
+	defer func() { require.NoError(t, bt.Close()) }()
 
 	expectedRatio := float64(1) * 100 / float64(2*maxKeys) // 2 because we'll have 2 pages.
 	stats := bt.Stats()
@@ -276,7 +284,8 @@ func BenchmarkPurge(b *testing.B) {
 
 	b.Run("btree", func(b *testing.B) {
 		start := time.Now()
-		bt := NewTree()
+		bt := NewTree("BenchmarkPurge")
+		defer func() { require.NoError(b, bt.Close()) }()
 		for i := 0; i < N; i++ {
 			bt.Set(rand.Uint64(), uint64(i))
 		}
@@ -299,7 +308,8 @@ func BenchmarkWrite(b *testing.B) {
 		}
 	})
 	b.Run("btree", func(b *testing.B) {
-		bt := NewTree()
+		bt := NewTree("BenchmarkWrite")
+		defer func() { require.NoError(b, bt.Close()) }()
 		b.ResetTimer()
 		for n := 0; n < b.N; n++ {
 			k := rand.Uint64()
@@ -332,7 +342,8 @@ func BenchmarkRead(b *testing.B) {
 		}
 	})
 
-	bt := NewTree()
+	bt := NewTree("BenchmarkRead")
+	defer func() { require.NoError(b, bt.Close()) }()
 	for i := 0; i < N; i++ {
 		k := uint64(rand.Intn(2*N)) + 1
 		bt.Set(k, k)