delay queue tool

cmd/main.go

@@ -20,15 +20,33 @@ func startServer() {
 }
 
 func main() {
-	//startServer()
+	startServer()
-	ms := time.Now().UnixMilli()
-	now := time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)
 
-	fmt.Println(ms)
+	//m := treemap.NewWithIntComparator()
-	fmt.Println(ms << 22)
+	//k, v := m.Min()
-	fmt.Println(ms - now.UnixMilli())
+	//fmt.Println(k, v)
-	key := (ms-now.UnixMilli())<<22 + 138
-	fmt.Println(key >> 22)
-	fmt.Println(key << 42 >> 42)
 
+	//wait := &sync.WaitGroup{}
+	//wait.Add(1)
+	//
+	//times := 100
+	//
+	//q := collect.NewDelayQueue(func(data int, now time.Time) {
+	//	fmt.Println(data, now)
+	//	if data == times-1 {
+	//		wait.Done()
+	//	}
+	//})
+	//for i := 0; i < times; i++ {
+	//	key := q.Add(time.Second*time.Duration(i), i)
+	//	if i%5 == 0 {
+	//		go func(k int64, idx int) {
+	//			<-time.After(time.Second * 3)
+	//			q.Cancel(k)
+	//			fmt.Println("cancel", idx, k)
+	//		}(key, i)
+	//	}
+	//}
+	//
+	//wait.Wait()
 }

tool/collect/delay_queue.go

@@ -1,38 +1,130 @@
 package collect
 
 import (
-	"fmt"
 	"github.com/emirpasic/gods/maps/treemap"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
 // DelayQueue key: int64(64bit) = 1(0固定值)+41(ms)+22(index)
 type DelayQueue[T any] struct {
-	taskQueueMap *treemap.Map
+	taskQueueMap   *treemap.Map
-	lock         *sync.RWMutex
+	lock           *sync.RWMutex
-	offsetMs     int64
+	offsetMs       int64
+	handler        func(data T, now time.Time)
+	supplying      *atomic.Bool
+	consumeChannel chan []T
 }
 
-func NewDelayQueue[T any](handler func(data T)) {
+func NewDelayQueue[T any](handler func(data T, now time.Time)) *DelayQueue[T] {
 	q := &DelayQueue[T]{}
 	q.taskQueueMap = treemap.NewWithIntComparator()
 	q.lock = &sync.RWMutex{}
 	q.offsetMs = time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC).UnixMilli()
+	q.handler = handler
+	q.supplying = &atomic.Bool{}
+	q.consumeChannel = make(chan []T, 32)
 
+	// 开启生产者轮训
+	q.supplier()
+	// 开启消费者
+	go func() {
+		for {
+			values := <-q.consumeChannel
+			if IsEmptySlice(values) {
+				continue
+			}
+			for _, data := range values {
+				q.handler(data, time.Now())
+			}
+		}
+	}()
+	return q
 }
 
+// supplier 定时开始检查
+func (q *DelayQueue[T]) supplier() {
+	// 正在轮训中...
+	if !q.supplying.CompareAndSwap(false, true) {
+		return
+	}
+	go func() {
+		for {
+			now := <-time.After(time.Millisecond * 100)
+			_, queue := q.next(now.UnixMilli())
+			if IsEmptySlice(queue) {
+				continue
+			}
+			q.consumeChannel <- queue
+			// 没元素了,停下
+			if q.taskQueueMap.Size() == 0 {
+				q.supplying.Store(false)
+				break
+			}
+		}
+	}()
+}
+
+// next 寻找下一个已到时间的任务数据
+func (q *DelayQueue[T]) next(nowMs int64) (int64, []T) {
+	k, v := q.taskQueueMap.Min()
+	if k == nil {
+		return 0, nil
+	}
+	ms := k.(int)
+	//fmt.Println("k", k, "now", nowMs)
+	if int64(ms) > nowMs {
+		// 还没到时间
+		return 0, nil
+	}
+	q.lock.Lock()
+	defer q.lock.Unlock()
+	q.taskQueueMap.Remove(int(ms))
+	return int64(ms), v.([]T)
+}
+
+// Add 添加任务
 func (q *DelayQueue[T]) Add(after time.Duration, data T) int64 {
 	ms := time.Now().UnixMilli() + after.Milliseconds()
 	q.lock.Lock()
 	defer q.lock.Unlock()
-	val, found := q.taskQueueMap.Get(ms)
+	val, found := q.taskQueueMap.Get(int(ms))
 	if !found {
-		q.taskQueueMap.Put(ms, []T{data})
+		q.taskQueueMap.Put(int(ms), []T{data})
-		return ms
+		return q.genKey(ms, 0)
+	}
+	values := val.([]T)
+	q.taskQueueMap.Put(int(ms), append(values, data))
+	// 通知生产者goroutine轮训数据
+	q.supplier()
+	return q.genKey(ms, len(values))
+}
+
+// Cancel 取消任务
+func (q *DelayQueue[T]) Cancel(key int64) {
+	ms, index := q.parseKey(key)
+	v, found := q.taskQueueMap.Get(int(ms))
+	if !found {
+		return
+	}
+	values := v.([]T)
+	if IsEmptySlice(values) || len(values) <= index {
+		return
+	}
+	var newValues []T
+	q.lock.Lock()
+	defer q.lock.Unlock()
+	// 只一个任务直接干掉
+	if len(values) == 1 {
+		q.taskQueueMap.Remove(int(ms))
+		return
 	}
-	q.taskQueueMap.Put(ms, append(val.([]T), data))
+	// 过滤掉要取消的任务
-	return 1
+	newValues = Filter(values, func(i int, _ T) bool {
+		return i != index
+	})
+	q.taskQueueMap.Put(int(ms), newValues)
 }
 
 func (q *DelayQueue[T]) genKey(ms int64, index int) int64 {
@@ -40,8 +132,6 @@ func (q *DelayQueue[T]) genKey(ms int64, index int) int64 {
 	return key
 }
 
-func parseKey(key int64) (int64, int) {
+func (q *DelayQueue[T]) parseKey(key int64) (int64, int) {
-	fmt.Println(key >> 22)
+	return key>>22 + q.offsetMs, int(key << 42 >> 42)
-	fmt.Println(key << 42 >> 42)
-	return 0, 0
 }