161 lines
3.2 KiB
Go
161 lines
3.2 KiB
Go
package timer
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/go-redis/redis/v8"
|
|
)
|
|
|
|
// 单次的任务队列
|
|
type worker struct {
|
|
ctx context.Context
|
|
zsetKey string
|
|
listKey string
|
|
redis *redis.Client
|
|
worker WorkerInterface
|
|
}
|
|
|
|
type WorkerCode int
|
|
|
|
const (
|
|
WorkerCodeSuccess WorkerCode = 0
|
|
WorkerCodeAgain WorkerCode = -1
|
|
)
|
|
|
|
// 需要考虑执行失败重新放入队列的情况
|
|
type WorkerInterface interface {
|
|
Worker(uniqueKey string, jobType string, data map[string]interface{}) WorkerCode
|
|
}
|
|
|
|
var wo *worker = nil
|
|
var once sync.Once
|
|
|
|
type extendData struct {
|
|
Delay time.Duration
|
|
Data map[string]interface{}
|
|
}
|
|
|
|
func InitWorker(ctx context.Context, re *redis.Client, w WorkerInterface) *worker {
|
|
|
|
once.Do(func() {
|
|
wo = &worker{
|
|
ctx: ctx,
|
|
zsetKey: "timer:job_zsetkey",
|
|
listKey: "timer:job_listkey",
|
|
redis: re,
|
|
worker: w,
|
|
}
|
|
go wo.getJob()
|
|
go wo.execJob()
|
|
})
|
|
|
|
return wo
|
|
}
|
|
|
|
// 添加任务
|
|
// 重复插入就代表覆盖
|
|
func (w *worker) AddJob(uniqueKey string, jobType string, delayTime time.Duration, data map[string]interface{}) error {
|
|
if delayTime.Abs() != delayTime {
|
|
return fmt.Errorf("时间间隔不能为负数")
|
|
}
|
|
if delayTime == 0 {
|
|
return fmt.Errorf("时间间隔不能为0")
|
|
}
|
|
|
|
redisKey := fmt.Sprintf("%s[:]%s", uniqueKey, jobType)
|
|
|
|
ed := extendData{
|
|
Delay: delayTime,
|
|
Data: data,
|
|
}
|
|
b, _ := json.Marshal(ed)
|
|
|
|
_, err := w.redis.SetEX(w.ctx, redisKey, b, delayTime+time.Second*5).Result()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
_, err = w.redis.ZAdd(w.ctx, w.zsetKey, &redis.Z{
|
|
Score: float64(time.Now().Add(delayTime).UnixMilli()),
|
|
Member: redisKey,
|
|
}).Result()
|
|
|
|
return err
|
|
}
|
|
|
|
// 删除任务
|
|
func (w *worker) DelJob(uniqueKey string, jobType string) error {
|
|
redisKey := fmt.Sprintf("%s[:]%s", uniqueKey, jobType)
|
|
|
|
w.redis.Del(w.ctx, redisKey).Result()
|
|
|
|
w.redis.ZRem(w.ctx, w.zsetKey, redisKey).Result()
|
|
|
|
return nil
|
|
}
|
|
|
|
// 获取任务
|
|
func (w *worker) getJob() {
|
|
timer := time.NewTicker(time.Millisecond * 100)
|
|
defer timer.Stop()
|
|
|
|
Loop:
|
|
for {
|
|
select {
|
|
case <-timer.C:
|
|
script := `
|
|
local token = redis.call('zrangebyscore',KEYS[1],ARGV[1],ARGV[2])
|
|
for i,v in ipairs(token) do
|
|
redis.call('zrem',KEYS[1],v)
|
|
redis.call('lpush',KEYS[2],v)
|
|
end
|
|
return "OK"
|
|
`
|
|
w.redis.Eval(w.ctx, script, []string{w.zsetKey, w.listKey}, 0, time.Now().UnixMilli()).Result()
|
|
// fmt.Println(i, err)
|
|
|
|
case <-w.ctx.Done():
|
|
break Loop
|
|
}
|
|
}
|
|
}
|
|
|
|
// 执行任务
|
|
func (w *worker) execJob() {
|
|
for {
|
|
keys, err := w.redis.BLPop(w.ctx, time.Second*10, w.listKey).Result()
|
|
if err != nil {
|
|
fmt.Println("watch err:", err)
|
|
continue
|
|
}
|
|
|
|
go func() {
|
|
s := strings.Split(keys[1], "[:]")
|
|
|
|
// 读取数据
|
|
str, err := w.redis.Get(w.ctx, keys[1]).Result()
|
|
if err != nil {
|
|
fmt.Println("execJob err:", err)
|
|
return
|
|
}
|
|
ed := extendData{}
|
|
json.Unmarshal([]byte(str), &ed)
|
|
|
|
fmt.Println("开始时间：", time.Now().Format("2006-01-02 15:04:05"))
|
|
code := w.worker.Worker(s[0], s[1], ed.Data)
|
|
|
|
if code == WorkerCodeAgain {
|
|
// 重新放入队列
|
|
fmt.Println("重入时间：", time.Now().Format("2006-01-02 15:04:05"))
|
|
w.AddJob(s[0], s[1], ed.Delay, ed.Data)
|
|
}
|
|
}()
|
|
|
|
}
|
|
}
|