118 lines
3.0 KiB
Go
Executable File
118 lines
3.0 KiB
Go
Executable File
package asynq
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"net/http"
|
||
"reflect"
|
||
|
||
"git.wishpal.cn/wishpal_ironfan/xframe/base/proc"
|
||
"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/handlers"
|
||
"git.wishpal.cn/wishpal_ironfan/xframe/component/storage/redis"
|
||
"github.com/hibiken/asynq"
|
||
"github.com/pkg/errors"
|
||
)
|
||
|
||
var (
|
||
xcAsynq = new(XcAsynq)
|
||
)
|
||
|
||
type Config struct {
|
||
redis.RedisConf
|
||
|
||
// 可选
|
||
DefaultQueueQps int `json:",default=500"` // 常规队列的qps限制，默认500
|
||
SlowQueueQps int `json:",default=50"` // 慢速队列的qps限制，默认50
|
||
}
|
||
|
||
// Init 初始化相关资源，但不开始执行异步任务
|
||
func Init(ctx context.Context, cfg *Config) error {
|
||
return xcAsynq.Init(ctx, cfg)
|
||
}
|
||
|
||
// Start 真的开始执行异步任务。和Init拆开，以防止handler依赖的模块还没初始化好（比如mongo连接）
|
||
func Start() error {
|
||
if err := xcAsynq.Start(); err != nil {
|
||
return err
|
||
}
|
||
|
||
proc.AddShutdownListener(func() {
|
||
Stop()
|
||
})
|
||
return nil
|
||
}
|
||
|
||
// Stop 优雅关闭异步任务
|
||
func Stop() {
|
||
if xcAsynq.server == nil {
|
||
return
|
||
}
|
||
xcAsynq.server.Stop()
|
||
xcAsynq.server.Shutdown()
|
||
}
|
||
|
||
// MustRegistHandler 注册handler，失败则panic。
|
||
// 支持的T类型：map、slice、array、struct、string/int等，及其指针。
|
||
// 不支持：interface{}。
|
||
func MustRegistHandler[T any](typ string, f func(context.Context, T) error) {
|
||
handlers.MustRegistHandler(typ, f)
|
||
}
|
||
|
||
// MustRegistHttpHandler
|
||
// 业务不要直接用，框架用于处理异步http请求
|
||
func MustRegistHttpHandler(typ string, f http.HandlerFunc) {
|
||
handlers.MustRegistHttpHandler(typ, f)
|
||
}
|
||
|
||
func GenHttpTaskPayload(r *http.Request) *JsonHelper {
|
||
return Json(handlers.RequestToPayload(r))
|
||
}
|
||
|
||
// JsonHelper 封装data来获得MarshalBinary方法
|
||
type JsonHelper struct {
|
||
data interface{}
|
||
}
|
||
|
||
func (j JsonHelper) MarshalBinary() ([]byte, error) {
|
||
return json.Marshal(j.data)
|
||
}
|
||
|
||
// TaskData 兼容的task数据类型
|
||
type TaskData interface {
|
||
[]byte | *JsonHelper
|
||
}
|
||
|
||
// Json 封装data以支持自动json序列化
|
||
func Json(data interface{}) *JsonHelper {
|
||
return &JsonHelper{data: data}
|
||
}
|
||
|
||
// AddTask 添加一个常规速度的任务
|
||
func AddTask[T TaskData](ctx context.Context, typ string, payload T, opts ...asynq.Option) (info *asynq.TaskInfo, err error) {
|
||
return addTask[T](ctx, queueDefault, typ, payload, opts...)
|
||
}
|
||
|
||
// AddSlowTask 添加一个低速的任务
|
||
func AddSlowTask[T TaskData](ctx context.Context, typ string, payload T, opts ...asynq.Option) (info *asynq.TaskInfo, err error) {
|
||
return addTask[T](ctx, queueSlow, typ, payload, opts...)
|
||
}
|
||
|
||
func addTask[T TaskData](ctx context.Context, queue, typ string, data T, opts ...asynq.Option) (info *asynq.TaskInfo, err error) {
|
||
var payload []byte
|
||
switch tmp := interface{}(data).(type) {
|
||
case []byte:
|
||
payload = tmp
|
||
case *JsonHelper:
|
||
payload, err = tmp.MarshalBinary()
|
||
if err != nil {
|
||
err = errors.Wrap(err, "JsonHelper.MarshalBinary")
|
||
return
|
||
}
|
||
default:
|
||
err = errors.Errorf("unexpected data type: %v", reflect.TypeOf(data))
|
||
return
|
||
}
|
||
|
||
return xcAsynq.addTask(ctx, queue, typ, payload, opts...)
|
||
}
|