Min heap (#10)

* min heap

* 更新

* 更新action

* 更新代码

* 更新run函数

* fix

* 完成最小堆

Author: guonaihong

.github/workflows/go.yml

@@ -10,12 +10,12 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        go: [ '1.13', '1.14']
+        go: ['1.19']
     name: Go ${{ matrix.go }} sample
 
     steps:
 
-    - name: Set up Go 1.13
+    - name: Set up Go 1.19
       uses: actions/setup-go@v1
       with:
         go-version: ${{ matrix.go }}

.gitignore

@@ -13,6 +13,7 @@ _long-time-test/long-time-test
 
 # Output of the go coverage tool, specifically when used with LiteIDE
 *.out
+cover.cov
 
 # Dependency directories (remove the comment below to include it)
 # vendor/

README.md

@@ -6,6 +6,7 @@ timer是高性能定时器库
 ## feature
 * 支持一次性定时器
 * 支持周期性定时器
+* 支持多种数据结构后端，最小堆，5级时间轮
 
 ## 一次性定时器
 ```go
@@ -73,6 +74,17 @@ func main() {
 	tm.Run()
 }
 ```
+## 选择不同的的数据结构
+```go
+import (
+    "github.com/antlabs/timer"
+    "log"
+)
+
+func main() {
+        tm := timer.NewTimer(timer.WithMinHeap())// 选择最小堆
+}
+```
 ## benchmark
 
 github.com/antlabs/timer 性能最高

go.mod

@@ -1,9 +1,15 @@
 module github.com/antlabs/timer
 
-go 1.13
+go 1.19
 
 require (
 	github.com/antlabs/stl v0.0.1
 	github.com/stretchr/testify v1.6.1
 	gopkg.in/go-playground/assert.v1 v1.2.1
 )
+
+require (
+	github.com/davecgh/go-spew v1.1.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
+)

go.sum

@@ -7,6 +7,7 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/go-playground/assert.v1 v1.2.1 h1:xoYuJVE7KT85PYWrN730RguIQO0ePzVRfFMXadIrXTM=
 gopkg.in/go-playground/assert.v1 v1.2.1/go.mod h1:9RXL0bg/zibRAgZUYszZSwO/z8Y/a8bDuhia5mkpMnE=

min_heap.go

@@ -0,0 +1,136 @@
+package timer
+
+import (
+	"container/heap"
+	"context"
+	"math"
+	"sync"
+	"time"
+)
+
+var _ Timer = (*minHeap)(nil)
+
+type minHeap struct {
+	mu sync.Mutex
+	minHeaps
+	chAdd  chan struct{}
+	ctx    context.Context
+	cancel context.CancelFunc
+	wait   sync.WaitGroup
+}
+
+// 一次性定时器
+func (m *minHeap) AfterFunc(expire time.Duration, callback func()) TimeNoder {
+	return m.addCallback(expire, callback, false)
+}
+
+// 加任务
+func (m *minHeap) addCallback(expire time.Duration, callback func(), isSchedule bool) TimeNoder {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	node := minHeapNode{
+		callback:   callback,
+		userExpire: expire,
+		absExpire:  time.Now().Add(expire),
+		isSchedule: isSchedule,
+		root:       m,
+	}
+
+	heap.Push(&m.minHeaps, node)
+	select {
+	case m.chAdd <- struct{}{}:
+	default:
+	}
+
+	return &node
+
+}
+
+func (m *minHeap) removeTimeNode(node *minHeapNode) {
+	m.mu.Lock()
+	if node.index < 0 || node.index > len(m.minHeaps) || len(m.minHeaps) == 0 {
+		m.mu.Unlock()
+		return
+	}
+
+	heap.Remove(&m.minHeaps, node.index)
+	m.mu.Unlock()
+}
+
+// 周期性定时器
+func (m *minHeap) ScheduleFunc(expire time.Duration, callback func()) TimeNoder {
+	return m.addCallback(expire, callback, true)
+}
+
+// 运行
+// 为了避免空转cpu, 会等待一个chan, 只要AfterFunc或者ScheduleFunc被调用就会往这个chan里面写值
+func (m *minHeap) Run() {
+
+	tm := time.NewTimer(time.Hour)
+	for {
+		select {
+		case <-tm.C:
+			for {
+				m.mu.Lock()
+				now := time.Now()
+				if m.minHeaps.Len() == 0 {
+					m.mu.Unlock()
+					goto next
+				}
+
+				first := &m.minHeaps[0]
+
+				if now.After(first.absExpire) {
+					callback := first.callback
+					if first.isSchedule {
+						first.absExpire = now.Add(first.userExpire)
+						heap.Fix(&m.minHeaps, first.index)
+					} else {
+						m.minHeaps.Pop()
+					}
+					go callback()
+				}
+
+				if m.minHeaps.Len() == 0 {
+					m.mu.Unlock()
+					goto next
+				}
+				first = &m.minHeaps[0]
+				if time.Now().Before(first.absExpire) {
+					to := time.Duration(math.Abs(float64(time.Since(m.minHeaps[0].absExpire))))
+					tm.Reset(to)
+					m.mu.Unlock()
+					goto next
+				}
+				m.mu.Unlock()
+			}
+		case <-m.chAdd:
+			m.mu.Lock()
+			// 极端情况，加完任务立即给删除了, 判断下当前堆中是否有元素
+			if m.minHeaps.Len() > 0 {
+				tm.Reset(time.Since(m.minHeaps[0].absExpire))
+			}
+			m.mu.Unlock()
+			// 进入事件循环，如果为空就会从事件循环里面退出
+		case <-m.ctx.Done():
+			// 等待所有任务结束
+			m.wait.Wait()
+			return
+		}
+	next:
+	}
+}
+
+// 停止所有定时器
+func (m *minHeap) Stop() {
+	m.cancel()
+}
+
+func newMinHeap() (mh *minHeap) {
+	mh = &minHeap{}
+	heap.Init(&mh.minHeaps)
+	mh.chAdd = make(chan struct{}, 1)
+	mh.ctx, mh.cancel = context.WithCancel(context.TODO())
+	return
+}

min_heap_node.go

@@ -0,0 +1,40 @@
+package timer
+
+import (
+	"time"
+)
+
+type minHeapNode struct {
+	callback   func()        //用户的callback
+	absExpire  time.Time     //绝对时间
+	userExpire time.Duration //过期时间段
+	isSchedule bool          //是否是周期性任务
+	index      int           //在min heap中的索引，方便删除或者重新推入堆中
+	root       *minHeap
+}
+
+func (m *minHeapNode) Stop() {
+	m.root.removeTimeNode(m)
+}
+
+type minHeaps []minHeapNode
+
+func (m minHeaps) Len() int           { return len(m) }
+func (m minHeaps) Less(i, j int) bool { return m[i].absExpire.Before(m[j].absExpire) }
+func (m minHeaps) Swap(i, j int)      { m[i], m[j] = m[j], m[i] }
+
+func (m *minHeaps) Push(x any) {
+	// Push and Pop use pointer receivers because they modify the slice's length,
+	// not just its contents.
+	*m = append(*m, x.(minHeapNode))
+	lastIndex := len(*m) - 1
+	(*m)[lastIndex].index = lastIndex
+}
+
+func (m *minHeaps) Pop() any {
+	old := *m
+	n := len(old)
+	x := old[n-1]
+	*m = old[0 : n-1]
+	return x
+}