feat(btree): allow option to reset btree and mmaping it to specified file.

z/btree.go

@@ -48,16 +48,22 @@ func (t *Tree) Release() {
 	}
 }
 
-// NewTree returns a memory mapped B+ tree.
-func NewTree(maxSz int) *Tree {
-	// Tell kernel that we'd be reading pages in random order, so don't do read ahead.
-	fd, err := ioutil.TempFile("", "btree")
-	check(err)
+func createFile(maxSz int, fname string) (*MmapFile, error) {
+	if fname == "" {
+		fd, err := ioutil.TempFile("", "btree")
+		check(err)
+		return OpenMmapFileUsing(fd, maxSz, true)
+	}
+	return OpenMmapFile(fname, os.O_RDWR|os.O_CREATE, maxSz)
+}
 
-	mf, err := OpenMmapFileUsing(fd, maxSz, true)
+// NewTree returns a memory mapped B+ tree with given filename.
+func NewTree(maxSz int, fname string) *Tree {
+	mf, err := createFile(maxSz, fname)
 	if err != NewFile {
 		check(err)
 	}
+	// Tell kernel that we'd be reading pages in random order, so don't do read ahead.
 	check(Madvise(mf.Data, false))
 
 	t := &Tree{
@@ -72,6 +78,19 @@ func NewTree(maxSz int) *Tree {
 	return t
 }
 
+// Reset resets the tree and truncates it to maxSz.
+func (t *Tree) Reset(maxSz int) {
+	t.nextPage = 1
+	t.freePage = 0
+	if maxSz < pageSize {
+		maxSz = pageSize
+	}
+	check(t.mf.Truncate(int64(maxSz)))
+	// Set the root.
+	t.newNode(0)
+	t.Set(absoluteMax, 0)
+}
+
 type TreeStats struct {
 	NextPage     int
 	NumNodes     int

z/btree_test.go

@@ -36,8 +36,7 @@ func setPageSize(sz int) {
 }
 
 func TestTree(t *testing.T) {
-	bt := NewTree(1 << 20)
-	// bt.Print()
+	bt := NewTree(1<<20, "")
 
 	N := uint64(256 * 256)
 	for i := uint64(1); i < N; i++ {
@@ -54,12 +53,11 @@ func TestTree(t *testing.T) {
 	for i := uint64(100); i < N; i++ {
 		require.Equal(t, i, bt.Get(i))
 	}
-	// bt.Print()
 }
 
 func TestTreeBasic(t *testing.T) {
 	setAndGet := func() {
-		bt := NewTree(1 << 20)
+		bt := NewTree(1<<20, "")
 		defer bt.Release()
 
 		N := uint64(1 << 20)
@@ -82,8 +80,43 @@ func TestTreeBasic(t *testing.T) {
 	setAndGet()
 }
 
+func TestTreeReset(t *testing.T) {
+	bt := NewTree(1<<20, "")
+	defer bt.Release()
+	N := 1 << 10
+	val := rand.Uint64()
+	for i := 0; i < N; i++ {
+		bt.Set(rand.Uint64(), val)
+	}
+
+	// Truncate it to small size that is less than pageSize.
+	bt.Reset(1 << 10)
+
+	stats := bt.Stats()
+	// Verify the tree stats.
+	require.Equal(t, 3, stats.NextPage)
+	require.Equal(t, 2, stats.NumNodes)
+	require.Equal(t, 1, stats.NumLeafKeys)
+	require.Equal(t, 1, stats.NumLeafKeys)
+	require.Equal(t, 2*pageSize, stats.Bytes)
+	expectedOcc := float64(2) * 100 / float64(stats.NumNodes*maxKeys)
+	require.InDelta(t, expectedOcc, stats.Occupancy, 0.01)
+	require.Zero(t, stats.FreePages)
+	// Check if we can reinsert the data.
+	mp := make(map[uint64]uint64)
+	for i := 0; i < N; i++ {
+		k := rand.Uint64()
+		mp[k] = val
+		bt.Set(k, val)
+	}
+	for k, v := range mp {
+		require.Equal(t, v, bt.Get(k))
+	}
+}
+
 func TestTreeCycle(t *testing.T) {
-	bt := NewTree(1 << 20)
+	bt := NewTree(1<<20, "")
+	defer bt.Release()
 	val := uint64(0)
 	for i := 0; i < 16; i++ {
 		for j := 0; j < 1e6+i*1e4; j++ {
@@ -109,7 +142,7 @@ func TestOccupancyRatio(t *testing.T) {
 	defer setPageSize(os.Getpagesize())
 	require.Equal(t, 4, maxKeys)
 
-	bt := NewTree(1 << 20)
+	bt := NewTree(1<<20, "")
 	defer bt.Release()
 
 	expectedRatio := float64(1) * 100 / float64(maxKeys)
@@ -212,7 +245,7 @@ func BenchmarkWrite(b *testing.B) {
 		}
 	})
 	b.Run("btree", func(b *testing.B) {
-		bt := NewTree(1 << 30)
+		bt := NewTree(1<<30, "")
 		defer bt.Release()
 		b.ResetTimer()
 		for n := 0; n < b.N; n++ {
@@ -246,7 +279,7 @@ func BenchmarkRead(b *testing.B) {
 		}
 	})
 
-	bt := NewTree(1 << 30)
+	bt := NewTree(1<<30, "")
 	defer bt.Release()
 	for i := 0; i < N; i++ {
 		k := uint64(rand.Intn(2*N)) + 1