Skip to content

Commit e9d55b6

Browse files
lhy1024lhy1024
committed
update, ref tikv#4399
Signed-off-by: lhy1024 <[email protected]>
1 parent 1b8f823 commit e9d55b6

File tree

5 files changed

+228
-122
lines changed

5 files changed

+228
-122
lines changed

‎pkg/heap/indexed_heap.go

+136
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
// Copyright 2022 TiKV Project 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 implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package heap
16+
17+
import "container/heap"
18+
19+
// TopNItem represents a single object in TopN.
20+
type TopNItem interface {
21+
// ID is used to check identity.
22+
ID() uint64
23+
// Less tests whether the current item is less than the given argument in the `k`th dimension.
24+
Less(k int, than TopNItem) bool
25+
}
26+
27+
// IndexedHeap is a heap with index.
28+
type IndexedHeap struct {
29+
k int
30+
rev bool
31+
items []TopNItem
32+
index map[uint64]int
33+
}
34+
35+
// NewMinHeap returns a min heap with the given hint.
36+
func NewMinHeap(k, hint int) *IndexedHeap {
37+
return &IndexedHeap{
38+
k: k,
39+
rev: false,
40+
items: make([]TopNItem, 0, hint),
41+
index: map[uint64]int{},
42+
}
43+
}
44+
45+
// NewMaxHeap returns a max heap with the given hint.
46+
func NewMaxHeap(k, hint int) *IndexedHeap {
47+
return &IndexedHeap{
48+
k: k,
49+
rev: true,
50+
items: make([]TopNItem, 0, hint),
51+
index: map[uint64]int{},
52+
}
53+
}
54+
55+
// Implementing heap.Interface.
56+
func (hp *IndexedHeap) Len() int {
57+
return len(hp.items)
58+
}
59+
60+
// Implementing heap.Interface.
61+
func (hp *IndexedHeap) Less(i, j int) bool {
62+
if !hp.rev {
63+
return hp.items[i].Less(hp.k, hp.items[j])
64+
}
65+
return hp.items[j].Less(hp.k, hp.items[i])
66+
}
67+
68+
// Implementing heap.Interface.
69+
func (hp *IndexedHeap) Swap(i, j int) {
70+
lid := hp.items[i].ID()
71+
rid := hp.items[j].ID()
72+
hp.items[i], hp.items[j] = hp.items[j], hp.items[i]
73+
hp.index[lid] = j
74+
hp.index[rid] = i
75+
}
76+
77+
// Implementing heap.Interface.
78+
func (hp *IndexedHeap) Push(x interface{}) {
79+
item := x.(TopNItem)
80+
hp.index[item.ID()] = hp.Len()
81+
hp.items = append(hp.items, item)
82+
}
83+
84+
// Implementing heap.Interface.
85+
func (hp *IndexedHeap) Pop() interface{} {
86+
l := hp.Len()
87+
item := hp.items[l-1]
88+
hp.items = hp.items[:l-1]
89+
delete(hp.index, item.ID())
90+
return item
91+
}
92+
93+
// Top returns the top item.
94+
func (hp *IndexedHeap) Top() TopNItem {
95+
if hp.Len() <= 0 {
96+
return nil
97+
}
98+
return hp.items[0]
99+
}
100+
101+
// Get returns item with the given ID.
102+
func (hp *IndexedHeap) Get(id uint64) TopNItem {
103+
idx, ok := hp.index[id]
104+
if !ok {
105+
return nil
106+
}
107+
item := hp.items[idx]
108+
return item
109+
}
110+
111+
// GetAll returns all the items.
112+
func (hp *IndexedHeap) GetAll() []TopNItem {
113+
all := make([]TopNItem, len(hp.items))
114+
copy(all, hp.items)
115+
return all
116+
}
117+
118+
// Put inserts item or updates the old item if it exists.
119+
func (hp *IndexedHeap) Put(item TopNItem) (isUpdate bool) {
120+
if idx, ok := hp.index[item.ID()]; ok {
121+
hp.items[idx] = item
122+
heap.Fix(hp, idx)
123+
return true
124+
}
125+
heap.Push(hp, item)
126+
return false
127+
}
128+
129+
// Remove deletes item by ID and returns it.
130+
func (hp *IndexedHeap) Remove(id uint64) TopNItem {
131+
if idx, ok := hp.index[id]; ok {
132+
item := heap.Remove(hp, idx)
133+
return item.(TopNItem)
134+
}
135+
return nil
136+
}

‎pkg/heap/indexed_heap_test.go

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
// Copyright 2022 TiKV Project 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 implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package heap
16+
17+
import (
18+
"testing"
19+
20+
. "github.com/pingcap/check"
21+
)
22+
23+
func Test(t *testing.T) {
24+
TestingT(t)
25+
}
26+
27+
var _ = Suite(&testIndexedHeapSuite{})
28+
29+
type testIndexedHeapSuite struct {
30+
}
31+
32+
type testItem struct {
33+
id uint64
34+
value float64
35+
}
36+
37+
func (it *testItem) ID() uint64 {
38+
return it.id
39+
}
40+
41+
func (it *testItem) Less(k int, than TopNItem) bool {
42+
return it.value < than.(*testItem).value
43+
}
44+
45+
func NewTestItem(id uint64, value float64) *testItem {
46+
return &testItem{
47+
id: id,
48+
value: value,
49+
}
50+
}
51+
52+
func (s *testIndexedHeapSuite) Test(c *C) {
53+
h := NewMinHeap(0, 10)
54+
v := h.Top()
55+
c.Assert(v, Equals, nil)
56+
57+
h.Put(NewTestItem(1, 1.0))
58+
h.Put(NewTestItem(2, 2.0))
59+
v = h.Top()
60+
c.Assert(v.(*testItem).value, Equals, 1.0)
61+
62+
h.Put(NewTestItem(2, 0.0))
63+
v = h.Top()
64+
c.Assert(v.(*testItem).value, Equals, 0.0)
65+
66+
h.Put(NewTestItem(1, -1.0))
67+
v = h.Top()
68+
c.Assert(v.(*testItem).value, Equals, -1.0)
69+
70+
h = NewMaxHeap(0, 10)
71+
h.Put(NewTestItem(1, 1.0))
72+
h.Put(NewTestItem(2, 2.0))
73+
v = h.Top()
74+
c.Assert(v.(*testItem).value, Equals, 2.0)
75+
76+
h.Put(NewTestItem(2, 3.0))
77+
v = h.Top()
78+
c.Assert(v.(*testItem).value, Equals, 3.0)
79+
80+
h.Put(NewTestItem(1, 4.0))
81+
v = h.Top()
82+
c.Assert(v.(*testItem).value, Equals, 4.0)
83+
}

‎server/statistics/hot_peer.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"time"
2020

2121
"github.com/pingcap/kvproto/pkg/metapb"
22+
"github.com/tikv/pd/pkg/heap"
2223
"github.com/tikv/pd/pkg/movingaverage"
2324
"github.com/tikv/pd/pkg/slice"
2425
"go.uber.org/zap"
@@ -120,7 +121,7 @@ func (stat *HotPeerStat) ID() uint64 {
120121
}
121122

122123
// Less compares two HotPeerStat.Implementing TopNItem.
123-
func (stat *HotPeerStat) Less(k int, than TopNItem) bool {
124+
func (stat *HotPeerStat) Less(k int, than heap.TopNItem) bool {
124125
return stat.GetLoad(RegionStatKind(k)) < than.(*HotPeerStat).GetLoad(RegionStatKind(k))
125126
}
126127

‎server/statistics/topn.go

+6-121
Original file line numberDiff line numberDiff line change
@@ -20,15 +20,9 @@ import (
2020
"fmt"
2121
"sync"
2222
"time"
23-
)
2423

25-
// TopNItem represents a single object in TopN.
26-
type TopNItem interface {
27-
// ID is used to check identity.
28-
ID() uint64
29-
// Less tests whether the current item is less than the given argument in the `k`th dimension.
30-
Less(k int, than TopNItem) bool
31-
}
24+
. "github.com/tikv/pd/pkg/heap"
25+
)
3226

3327
// TopN maintains the N largest items of multiple dimensions.
3428
type TopN struct {
@@ -130,16 +124,16 @@ func (tn *TopN) maintain() {
130124
type singleTopN struct {
131125
k int
132126
n int
133-
topn *indexedHeap
134-
rest *indexedHeap
127+
topn *IndexedHeap
128+
rest *IndexedHeap
135129
}
136130

137131
func newSingleTopN(k, n int) *singleTopN {
138132
return &singleTopN{
139133
k: k,
140134
n: n,
141-
topn: newTopNHeap(k, n),
142-
rest: newRevTopNHeap(k, n),
135+
topn: NewMinHeap(k, n),
136+
rest: NewMaxHeap(k, n),
143137
}
144138
}
145139

@@ -211,115 +205,6 @@ func (stn *singleTopN) maintain() {
211205
}
212206
}
213207

214-
// indexedHeap is a heap with index.
215-
type indexedHeap struct {
216-
k int
217-
rev bool
218-
items []TopNItem
219-
index map[uint64]int
220-
}
221-
222-
func newTopNHeap(k, hint int) *indexedHeap {
223-
return &indexedHeap{
224-
k: k,
225-
rev: false,
226-
items: make([]TopNItem, 0, hint),
227-
index: map[uint64]int{},
228-
}
229-
}
230-
231-
func newRevTopNHeap(k, hint int) *indexedHeap {
232-
return &indexedHeap{
233-
k: k,
234-
rev: true,
235-
items: make([]TopNItem, 0, hint),
236-
index: map[uint64]int{},
237-
}
238-
}
239-
240-
// Implementing heap.Interface.
241-
func (hp *indexedHeap) Len() int {
242-
return len(hp.items)
243-
}
244-
245-
// Implementing heap.Interface.
246-
func (hp *indexedHeap) Less(i, j int) bool {
247-
if !hp.rev {
248-
return hp.items[i].Less(hp.k, hp.items[j])
249-
}
250-
return hp.items[j].Less(hp.k, hp.items[i])
251-
}
252-
253-
// Implementing heap.Interface.
254-
func (hp *indexedHeap) Swap(i, j int) {
255-
lid := hp.items[i].ID()
256-
rid := hp.items[j].ID()
257-
hp.items[i], hp.items[j] = hp.items[j], hp.items[i]
258-
hp.index[lid] = j
259-
hp.index[rid] = i
260-
}
261-
262-
// Implementing heap.Interface.
263-
func (hp *indexedHeap) Push(x interface{}) {
264-
item := x.(TopNItem)
265-
hp.index[item.ID()] = hp.Len()
266-
hp.items = append(hp.items, item)
267-
}
268-
269-
// Implementing heap.Interface.
270-
func (hp *indexedHeap) Pop() interface{} {
271-
l := hp.Len()
272-
item := hp.items[l-1]
273-
hp.items = hp.items[:l-1]
274-
delete(hp.index, item.ID())
275-
return item
276-
}
277-
278-
// Top returns the top item.
279-
func (hp *indexedHeap) Top() TopNItem {
280-
if hp.Len() <= 0 {
281-
return nil
282-
}
283-
return hp.items[0]
284-
}
285-
286-
// Get returns item with the given ID.
287-
func (hp *indexedHeap) Get(id uint64) TopNItem {
288-
idx, ok := hp.index[id]
289-
if !ok {
290-
return nil
291-
}
292-
item := hp.items[idx]
293-
return item
294-
}
295-
296-
// GetAll returns all the items.
297-
func (hp *indexedHeap) GetAll() []TopNItem {
298-
all := make([]TopNItem, len(hp.items))
299-
copy(all, hp.items)
300-
return all
301-
}
302-
303-
// Put inserts item or updates the old item if it exists.
304-
func (hp *indexedHeap) Put(item TopNItem) (isUpdate bool) {
305-
if idx, ok := hp.index[item.ID()]; ok {
306-
hp.items[idx] = item
307-
heap.Fix(hp, idx)
308-
return true
309-
}
310-
heap.Push(hp, item)
311-
return false
312-
}
313-
314-
// Remove deletes item by ID and returns it.
315-
func (hp *indexedHeap) Remove(id uint64) TopNItem {
316-
if idx, ok := hp.index[id]; ok {
317-
item := heap.Remove(hp, idx)
318-
return item.(TopNItem)
319-
}
320-
return nil
321-
}
322-
323208
type ttlItem struct {
324209
id uint64
325210
expire time.Time

0 commit comments

Comments
 (0)