diff --git a/z/btree.go b/z/btree.go index 28a1c136..b9300cd0 100644 --- a/z/btree.go +++ b/z/btree.go @@ -48,30 +48,50 @@ 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) +} + +func (t *Tree) initRootNode() { + // This is the root node. + t.newNode(0) + // This acts as the rightmost pointer (all the keys are <= this key). + t.Set(absoluteMax, 0) +} - 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{ mf: mf, nextPage: 1, } - // This is the root node. - t.newNode(0) - - // This acts as the rightmost pointer (all the keys are <= this key). - t.Set(absoluteMax, 0) + t.initRootNode() 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))) + t.initRootNode() +} + type TreeStats struct { NextPage int NumNodes int diff --git a/z/btree_test.go b/z/btree_test.go index 3467459c..1e4c055b 100644 --- a/z/btree_test.go +++ b/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