jtimer/manager.go at master · wwj31/jtimer · GitHub

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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
package jtimer

import (
	"container/heap"
	"sync"
	"time"

	"github.com/rs/xid"
)

var timerPool = sync.Pool{New: func() interface{} { return new(timer) }}

type CallbackFn func(dt time.Duration)

type Manager struct {
	timers map[Id]*timer
	heap   timerHeap
}

func New() *Manager {
	h := make(timerHeap, 0)
	heap.Init(&h)

	return &Manager{
		timers: make(map[Id]*timer),
		heap:   h,
	}
}

func (m *Manager) Add(now, endAt time.Time, callback CallbackFn, execCount int, id ...Id) Id {
	var timerId Id
	if len(id) > 0 {
		timerId = id[0]
	} else {
		timerId = xid.New().String()
	}

	if endAt.Before(now) {
		endAt = now
	}

	if oldTimer, exist := m.timers[timerId]; exist {
		oldTimer.startAt = now
		oldTimer.endAt = endAt
		oldTimer.callback = callback
		oldTimer.execCount = execCount
		heap.Fix(&m.heap, oldTimer.index)
		return oldTimer.id
	}

	newTimer := timerPool.Get().(*timer)
	newTimer.id = timerId
	newTimer.startAt = now
	newTimer.endAt = endAt
	newTimer.callback = callback
	newTimer.execCount = execCount

	heap.Push(&m.heap, newTimer)
	m.timers[timerId] = newTimer
	return newTimer.id
}

func (m *Manager) Remove(id Id, softRemove ...bool) {
	var softrm bool
	if len(softRemove) > 0 && softRemove[0] {
		softrm = true
	}
	m.remove(id, softrm)
}

func (m *Manager) Len() int {
	return len(m.heap)
}

func (m *Manager) NextUpdateAt() (at time.Time) {
	headTimer := m.heap.peek()
	if headTimer == nil {
		return time.Unix(0, 0)
	}

	return headTimer.endAt
}

func (m *Manager) Update(now time.Time) {
	for m.heap.peek() != nil {
		headTimer := m.heap.peek()
		if now.Before(headTimer.endAt) {
			return
		}

		if headTimer.remove {
			m.remove(headTimer.id, false)
			continue
		}

		timerDuration := headTimer.endAt.Sub(headTimer.startAt)
		totalElapsed := now.Sub(headTimer.startAt)

		if headTimer.spareCount() {
			count := totalElapsed / timerDuration
			elapsedDuration := count * timerDuration

			headTimer.consumeCount(int(count))
			headTimer.startAt = headTimer.startAt.Add(elapsedDuration)
			headTimer.endAt = headTimer.startAt.Add(timerDuration)

			if headTimer.callback != nil {
				headTimer.callback(elapsedDuration)
			}
		}

		if headTimer.spareCount() && !headTimer.remove && m.timers[headTimer.id] != nil {
			heap.Fix(&m.heap, headTimer.index)
		} else {
			m.remove(headTimer.id, false)
		}
	}
}

func (m *Manager) remove(id Id, softRemove bool) {
	_timer, found := m.timers[id]
	if !found {
		return
	}

	if softRemove && _timer.spareCount() {
		_timer.remove = true
		return
	}

	heap.Remove(&m.heap, _timer.index)
	delete(m.timers, id)

	timerPool.Put(_timer)
}