Initialize persistent tree correctly (#271)

z/btree.go

@@ -28,9 +28,12 @@ import (
 )
 
 var (
-	pageSize    = os.Getpagesize()
-	maxKeys     = (pageSize / 16) - 1
-	oneThird    = int(float64(maxKeys) / 3)
+	pageSize = os.Getpagesize()
+	maxKeys  = (pageSize / 16) - 1
+	oneThird = int(float64(maxKeys) / 3)
+)
+
+const (
 	absoluteMax = uint64(math.MaxUint64 - 1)
 	minSize     = 1 << 20
 )
@@ -65,24 +68,92 @@ func NewTree(tag string) *Tree {
 
 // NewTree returns a persistent on-disk B+ tree.
 func NewTreePersistent(path string) (*Tree, error) {
-	buffer, err := NewBufferPersistent(path)
+	t := &Tree{}
+	var err error
+
+	// Open the buffer from disk and set it to the maximum allocated size.
+	t.buffer, err = NewBufferPersistent(path, minSize)
 	if err != nil {
 		return nil, err
 	}
-	tree := &Tree{buffer: buffer}
-	tree.Reset()
-	return tree, nil
+	t.buffer.offset = uint64(len(t.buffer.buf))
+	t.data = t.buffer.Bytes()
+
+	// pageID can never be 0 if the tree has been initialized.
+	root := t.node(1)
+	isInitialized := root.pageID() != 0
+
+	if !isInitialized {
+		t.nextPage = 1
+		t.freePage = 0
+		t.initRootNode()
+	} else {
+		t.reinit()
+	}
+
+	return t, nil
+}
+
+// reinit sets the internal variables of a Tree, which are normally stored in
+// memory, but are lost when loading from disk.
+func (t *Tree) reinit() {
+	// Calculate t.nextPage by finding the highest pageId among all the nodes.
+	maxPageId := uint64(0)
+	t.Iterate(func(n node) {
+		if pageId := n.pageID(); pageId > t.nextPage {
+			maxPageId = pageId
+		}
+		// If this is a leaf node, increment the stats.
+		if n.isLeaf() {
+			t.stats.NumLeafKeys += n.numKeys()
+		}
+	})
+	t.nextPage = maxPageId + 1
+
+	// Calculate t.freePage by finding the page to which no other page points.
+	// This would be the root of the page tree.
+	// childPages[i] is true if pageId i+1 is a child page.
+	childPages := make([]bool, maxPageId)
+	// Mark all pages containing nodes as child pages.
+	t.Iterate(func(n node) {
+		i := n.pageID() - 1
+		childPages[i] = true
+	})
+	// pointedPages is a list of page IDs that the child pages point to.
+	pointedPages := make([]uint64, 0)
+	for i, isChild := range childPages {
+		if !isChild {
+			pageId := uint64(i) + 1
+			pointedPages = append(pointedPages, t.node(pageId).uint64(0))
+			t.stats.NumPagesFree++
+		}
+	}
+	// Mark all pages being pointed to as child pages.
+	for _, pageId := range pointedPages {
+		i := pageId - 1
+		childPages[i] = true
+	}
+	// There should only be one root page left.
+	for i, isChild := range childPages {
+		if !isChild {
+			pageId := uint64(i) + 1
+			t.freePage = pageId
+			break
+		}
+	}
 }
 
 // Reset resets the tree and truncates it to maxSz.
 func (t *Tree) Reset() {
-	t.nextPage = 1
-	t.freePage = 0
+	// Tree relies on uninitialized data being zeroed out, so we need to Memclr
+	// the data before using it again.
 	Memclr(t.buffer.buf)
 	t.buffer.Reset()
 	t.buffer.AllocateOffset(minSize)
 	t.data = t.buffer.Bytes()
 	t.stats = TreeStats{}
+	t.nextPage = 1
+	t.freePage = 0
 	t.initRootNode()
 }
 
@@ -119,7 +190,7 @@ func (t *Tree) Stats() TreeStats {
 	return out
 }
 
-// BytesToU32Slice converts the given byte slice to uint32 slice
+// BytesToUint64Slice converts a byte slice to a uint64 slice.
 func BytesToUint64Slice(b []byte) []uint64 {
 	if len(b) == 0 {
 		return nil

z/btree_test.go

@@ -22,6 +22,7 @@ import (
 	"math"
 	"math/rand"
 	"os"
+	"path/filepath"
 	"sort"
 	"testing"
 	"time"
@@ -59,6 +60,44 @@ func TestTree(t *testing.T) {
 	}
 }
 
+func TestTreePersistent(t *testing.T) {
+	dir, err := ioutil.TempDir("", "")
+	require.NoError(t, err)
+	defer os.RemoveAll(dir)
+	path := filepath.Join(dir, "tree.buf")
+
+	// Create a tree and validate the data.
+	bt, err := NewTreePersistent(path)
+	require.NoError(t, err)
+	N := uint64(64 << 10)
+	for i := uint64(1); i < N; i++ {
+		bt.Set(i, i*2)
+	}
+	for i := uint64(1); i < N; i++ {
+		require.Equal(t, i*2, bt.Get(i))
+	}
+	require.NoError(t, bt.Close())
+	freePage := bt.freePage
+	nextPage := bt.nextPage
+	stats := bt.Stats()
+
+	// Reopen tree and validate the data.
+	bt, err = NewTreePersistent(path)
+	require.NoError(t, err)
+	require.Equal(t, freePage, bt.freePage)
+	require.Equal(t, nextPage, bt.nextPage)
+	statsNew := bt.Stats()
+	// When reopening a tree, the allocated size becomes the file size.
+	// We don't need to compare this, it doesn't change anything in the tree.
+	statsNew.Allocated = stats.Allocated
+	require.Equal(t, stats, statsNew)
+	for i := uint64(1); i < N; i++ {
+		require.Equal(t, i*2, bt.Get(i))
+	}
+	require.NoError(t, err)
+	require.NoError(t, bt.Close())
+}
+
 func TestTreeBasic(t *testing.T) {
 	setAndGet := func() {
 		bt := NewTree("TestTreeBasic")

z/buffer.go

@@ -73,12 +73,14 @@ func NewBuffer(capacity int, tag string) *Buffer {
 	}
 }
 
-func NewBufferPersistent(path string) (*Buffer, error) {
+// It is the caller's responsibility to set offset after this, because Buffer
+// doesn't remember what it was.
+func NewBufferPersistent(path string, capacity int) (*Buffer, error) {
 	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0666)
 	if err != nil {
 		return nil, err
 	}
-	buffer, err := newBufferFile(file, 0)
+	buffer, err := newBufferFile(file, capacity)
 	if err != nil {
 		return nil, err
 	}
@@ -87,9 +89,6 @@ func NewBufferPersistent(path string) (*Buffer, error) {
 }
 
 func NewBufferTmp(dir string, capacity int) (*Buffer, error) {
-	if capacity == 0 {
-		capacity = defaultCapacity
-	}
 	if dir == "" {
 		dir = tmpDir
 	}
@@ -101,14 +100,17 @@ func NewBufferTmp(dir string, capacity int) (*Buffer, error) {
 }
 
 func newBufferFile(file *os.File, capacity int) (*Buffer, error) {
+	if capacity == 0 {
+		capacity = defaultCapacity
+	}
 	mmapFile, err := OpenMmapFileUsing(file, capacity, true)
 	if err != nil && err != NewFile {
 		return nil, err
 	}
 	buf := &Buffer{
 		buf:      mmapFile.Data,
 		bufType:  UseMmap,
-		curSz:    capacity,
+		curSz:    len(mmapFile.Data),
 		mmapFile: mmapFile,
 		offset:   8,
 		padding:  8,