Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use z.Buffer backing for B+ tree #268

Merged
merged 9 commits into from
Apr 28, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 36 additions & 15 deletions z/btree.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
29 changes: 20 additions & 9 deletions z/btree_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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++ {
Expand All @@ -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)
Expand All @@ -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++ {
Expand Down Expand Up @@ -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++ {
Expand All @@ -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)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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))
}
Expand All @@ -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()
Expand Down Expand Up @@ -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)
Expand Down
Loading