cubefs

1
// Copyright 2018 The CubeFS Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12
// implied. See the License for the specific language governing
13
// permissions and limitations under the License.
14

15
package metanode
16

17
import (
18
	"sync"
19

20
	"github.com/cubefs/cubefs/util/btree"
21
)
22

23
const defaultBTreeDegree = 32
24

25
type (
26
	// BtreeItem type alias google btree Item
27
	BtreeItem = btree.Item
28
)
29

30
// BTree is the wrapper of Google's btree.
31
type BTree struct {
32
	sync.RWMutex
33
	tree *btree.BTree
34
}
35

36
// NewBtree creates a new btree.
37
func NewBtree() *BTree {
38
	return &BTree{
39
		tree: btree.New(defaultBTreeDegree),
40
	}
41
}
42

43
// Get returns the object of the given key in the btree.
44
func (b *BTree) Get(key BtreeItem) (item BtreeItem) {
45
	b.RLock()
46
	item = b.tree.Get(key)
47
	b.RUnlock()
48
	return
49
}
50

51
func (b *BTree) CopyGet(key BtreeItem) (item BtreeItem) {
52
	b.Lock()
53
	item = b.tree.CopyGet(key)
54
	b.Unlock()
55
	return
56
}
57

58
// Find searches for the given key in the btree.
59
func (b *BTree) Find(key BtreeItem, fn func(i BtreeItem)) {
60
	b.RLock()
61
	item := b.tree.Get(key)
62
	b.RUnlock()
63
	if item == nil {
64
		return
65
	}
66
	fn(item)
67
}
68

69
func (b *BTree) CopyFind(key BtreeItem, fn func(i BtreeItem)) {
70
	b.Lock()
71
	item := b.tree.CopyGet(key)
72
	fn(item)
73
	b.Unlock()
74
}
75

76
// Has checks if the key exists in the btree.
77
func (b *BTree) Has(key BtreeItem) (ok bool) {
78
	b.RLock()
79
	ok = b.tree.Has(key)
80
	b.RUnlock()
81
	return
82
}
83

84
// Delete deletes the object by the given key.
85
func (b *BTree) Delete(key BtreeItem) (item BtreeItem) {
86
	b.Lock()
87
	item = b.tree.Delete(key)
88
	b.Unlock()
89
	return
90
}
91

92
func (b *BTree) Execute(fn func(tree *btree.BTree) interface{}) interface{} {
93
	b.Lock()
94
	defer b.Unlock()
95
	return fn(b.tree)
96
}
97

98
// ReplaceOrInsert is the wrapper of google's btree ReplaceOrInsert.
99
func (b *BTree) ReplaceOrInsert(key BtreeItem, replace bool) (item BtreeItem, ok bool) {
100
	b.Lock()
101
	if replace {
102
		item = b.tree.ReplaceOrInsert(key)
103
		b.Unlock()
104
		ok = true
105
		return
106
	}
107

108
	item = b.tree.Get(key)
109
	if item == nil {
110
		item = b.tree.ReplaceOrInsert(key)
111
		b.Unlock()
112
		ok = true
113
		return
114
	}
115
	ok = false
116
	b.Unlock()
117
	return
118
}
119

120
// Ascend is the wrapper of the google's btree Ascend.
121
// This function scans the entire btree. When the data is huge, it is not recommended to use this function online.
122
// Instead, it is recommended to call GetTree to obtain the snapshot of the current btree, and then do the scan on the snapshot.
123
func (b *BTree) Ascend(fn func(i BtreeItem) bool) {
124
	b.RLock()
125
	b.tree.Ascend(fn)
126
	b.RUnlock()
127
}
128

129
// AscendRange is the wrapper of the google's btree AscendRange.
130
func (b *BTree) AscendRange(greaterOrEqual, lessThan BtreeItem, iterator func(i BtreeItem) bool) {
131
	b.RLock()
132
	b.tree.AscendRange(greaterOrEqual, lessThan, iterator)
133
	b.RUnlock()
134
}
135

136
// AscendGreaterOrEqual is the wrapper of the google's btree AscendGreaterOrEqual
137
func (b *BTree) AscendGreaterOrEqual(pivot BtreeItem, iterator func(i BtreeItem) bool) {
138
	b.RLock()
139
	b.tree.AscendGreaterOrEqual(pivot, iterator)
140
	b.RUnlock()
141
}
142

143
// GetTree returns the snapshot of a btree.
144
func (b *BTree) GetTree() *BTree {
145
	b.Lock()
146
	t := b.tree.Clone()
147
	b.Unlock()
148
	nb := NewBtree()
149
	nb.tree = t
150
	return nb
151
}
152

153
// Reset resets the current btree.
154
func (b *BTree) Reset() {
155
	b.Lock()
156
	b.tree.Clear(true)
157
	b.Unlock()
158
}
159

160
// Len returns the total number of items in the btree.
161
func (b *BTree) Len() (size int) {
162
	b.RLock()
163
	size = b.tree.Len()
164
	b.RUnlock()
165
	return
166
}
167

168
// MaxItem returns the largest item in the btree.
169
func (b *BTree) MaxItem() BtreeItem {
170
	b.RLock()
171
	item := b.tree.Max()
172
	b.RUnlock()
173
	return item
174
}
175