blob: 17ddb358e195fa7b56967cb4ebceee27a6d8cc99 [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001// Copyright 2018 The etcd 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
15package lease
16
17import "container/heap"
18
19// LeaseWithTime contains lease object with a time.
20// For the lessor's lease heap, time identifies the lease expiration time.
21// For the lessor's lease checkpoint heap, the time identifies the next lease checkpoint time.
22type LeaseWithTime struct {
23 id LeaseID
24 // Unix nanos timestamp.
25 time int64
26 index int
27}
28
29type LeaseQueue []*LeaseWithTime
30
31func (pq LeaseQueue) Len() int { return len(pq) }
32
33func (pq LeaseQueue) Less(i, j int) bool {
34 return pq[i].time < pq[j].time
35}
36
37func (pq LeaseQueue) Swap(i, j int) {
38 pq[i], pq[j] = pq[j], pq[i]
39 pq[i].index = i
40 pq[j].index = j
41}
42
43func (pq *LeaseQueue) Push(x interface{}) {
44 n := len(*pq)
45 item := x.(*LeaseWithTime)
46 item.index = n
47 *pq = append(*pq, item)
48}
49
50func (pq *LeaseQueue) Pop() interface{} {
51 old := *pq
52 n := len(old)
53 item := old[n-1]
54 item.index = -1 // for safety
55 *pq = old[0 : n-1]
56 return item
57}
58
59// LeaseExpiredNotifier is a queue used to notify lessor to revoke expired lease.
60// Only save one item for a lease, `Register` will update time of the corresponding lease.
61type LeaseExpiredNotifier struct {
62 m map[LeaseID]*LeaseWithTime
63 queue LeaseQueue
64}
65
66func newLeaseExpiredNotifier() *LeaseExpiredNotifier {
67 return &LeaseExpiredNotifier{
68 m: make(map[LeaseID]*LeaseWithTime),
69 queue: make(LeaseQueue, 0),
70 }
71}
72
73func (mq *LeaseExpiredNotifier) Init() {
74 heap.Init(&mq.queue)
75 mq.m = make(map[LeaseID]*LeaseWithTime)
76 for _, item := range mq.queue {
77 mq.m[item.id] = item
78 }
79}
80
81func (mq *LeaseExpiredNotifier) RegisterOrUpdate(item *LeaseWithTime) {
82 if old, ok := mq.m[item.id]; ok {
83 old.time = item.time
84 heap.Fix(&mq.queue, old.index)
85 } else {
86 heap.Push(&mq.queue, item)
87 mq.m[item.id] = item
88 }
89}
90
91func (mq *LeaseExpiredNotifier) Unregister() *LeaseWithTime {
92 item := heap.Pop(&mq.queue).(*LeaseWithTime)
93 delete(mq.m, item.id)
94 return item
95}
96
97func (mq *LeaseExpiredNotifier) Poll() *LeaseWithTime {
98 if mq.Len() == 0 {
99 return nil
100 }
101 return mq.queue[0]
102}
103
104func (mq *LeaseExpiredNotifier) Len() int {
105 return len(mq.m)
106}