feat: add on timer & watermarkmain

1 files changed, 108 insertions, 0 deletions

diff --git a/core/timer.go b/core/timer.go
new file mode 100644
index 0000000..f744382
--- /dev/null
+++ b/core/timer.go
@@ -0,0 +1,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
+}