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