1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
|
package core
import (
"container/heap"
"context"
)
// If using rockDB no need priority queue anymore
// implement priority queue from golang
// ref: https://pkg.go.dev/container/heap#example-package-PriorityQueue
type timerServiceKey struct{}
// InjectTimerService inject the timer into Context
func InjectTimerService(ctx context.Context, ts TimerService) context.Context {
return context.WithValue(ctx, timerServiceKey{}, ts)
}
// ExtractTimerService get the timer for user from context
func ExtractTimerService(ctx context.Context) (TimerService, bool) {
ts, ok := ctx.Value(timerServiceKey{}).(TimerService)
return ts, ok
}
type Timer struct {
// Timestamp is the time when Timer is triggered
Timestamp int64
// Key is for who User involve
Key string
}
// Priority Queue (Min Heap). Smallest timestamp will be on top
type timerHeap []Timer
func (h *timerHeap) Len() int { return len(*h) }
func (h *timerHeap) Less(i, j int) bool { return (*h)[i].Timestamp < (*h)[j].Timestamp }
func (h *timerHeap) Swap(i, j int) {
(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}
func (h *timerHeap) Push(x any) {
*h = append(*h, x.(Timer))
}
func (h *timerHeap) Pop() any {
old := *h
n := len(old)
last := old[n-1]
*h = old[:n-1]
return last
}
type TimerService interface {
Register(time int64, key string)
CurrentWatermark() int64
}
type InternalTimerService interface {
TimerService
AdvanceWatermark(watermark int64) []Timer
}
type timerService struct {
queue timerHeap
currentWatermark int64
}
func NewTimerService() InternalTimerService {
return &timerService{
queue: make(timerHeap, 0),
currentWatermark: -1,
}
}
func (t *timerService) Register(time int64, key string) {
heap.Push(&t.queue, Timer{
Timestamp: time,
Key: key,
})
}
func (t *timerService) CurrentWatermark() int64 {
return t.currentWatermark
}
func (t *timerService) AdvanceWatermark(watermark int64) []Timer {
if watermark < t.currentWatermark {
// no need to advance when get smaller
return nil
}
t.currentWatermark = watermark
var timers []Timer
for t.queue.Len() > 0 {
earliestTimer := t.queue[0]
if earliestTimer.Timestamp > watermark {
break
}
pop := heap.Pop(&t.queue).(Timer)
timers = append(timers, pop)
}
return timers
}
|