xframe/component/queue/asynq/asynq_define.go

package asynq

import (
	"context"
	"net/http"
	"sync"
	"time"

	"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/errhandler"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/handlers"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/loggerx"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/maker"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/rate"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/queue/asynq/internal/utils"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/storage/redis"
	"github.com/hibiken/asynq"
	"github.com/hibiken/asynqmon"
	"github.com/pkg/errors"
)

const (
	queueDefault    = "default"
	queueSlow       = "slow"
	DefaultQueueQps = 500
	SlowQueueQps    = 50
)

type XcAsynq struct {
	cfg    *Config
	server *asynq.Server
	client *asynq.Client

	rdb      *redis.Redis
	rateCtrl *rate.RateCtrl

	queuePrefix string

	initLock  sync.Mutex
	startLock sync.Mutex
	inited    bool
	started   bool
}

func (q *XcAsynq) Init(ctx context.Context, cfg *Config) (err error) {
	// 上锁，防止重复init
	q.initLock.Lock()
	defer q.initLock.Unlock()

	if q.inited {
		return
	}
	defer func() {
		if err == nil {
			q.inited = true
		}
	}()

	q.cfg = cfg

	if cfg.DefaultQueueQps <= 0 {
		cfg.DefaultQueueQps = DefaultQueueQps
	}
	if cfg.SlowQueueQps <= 0 {
		cfg.SlowQueueQps = SlowQueueQps
	}
	if cfg.DefaultQueueQps < cfg.SlowQueueQps {
		err = errors.Errorf("default queue qps (%d) must larget than slow queue qps (%d)", cfg.DefaultQueueQps,
			cfg.SlowQueueQps)
		return
	}

	// 连redis
	rdb, err := redis.NewRedis(cfg.RedisConf)
	if err != nil {
		err = errors.Wrap(err, "new redis：")
		return
	}

	// 检查集群版redis的lua兼容性
	// 阿里的集群版做了特殊限制，要求keys必须显式传入，导致asynq的lua脚本报错：
	//reclaimStateAggregationSetsCmd、listLeaseExpiredCmd里面的key是动态获取的
	testLuaCmd := redis.NewScript(`return redis.call("TYPE", "test_cmd_key1")`)
	err = testLuaCmd.Run(ctx, rdb, []string{"fake"}).Err()
	if err != nil {
		err = errors.Wrap(err, "lua incompatible：")
		return
	}

	q.rdb = rdb
	q.rateCtrl = rate.NewRateCtrl(rdb)
	q.queuePrefix = utils.GenDefaultQueuePrefix()
	q.server = asynq.NewServer(maker.NewMaker(rdb), asynq.Config{
		BaseContext: utils.BaseContext,
		Queues: map[string]int{
			q.queuePrefix + queueDefault: 5,
			q.queuePrefix + queueSlow:    1,
		},
		IsFailure:                rate.IsFailure,
		RetryDelayFunc:           rate.RetryDelayFunc,
		DelayedTaskCheckInterval: time.Second / 2, // 超过1秒会导致超出qps的部分被delay过久，实际qps低于预期：qps60时想限制为50，结果实际跑出来10qps
		ErrorHandler:             errhandler.ErrorHandler{},
		Logger:                   loggerx.Logger{},
	})
	q.client = asynq.NewClient(maker.NewMaker(rdb))
	return
}

func (q *XcAsynq) Start() (err error) {
	q.startLock.Lock()
	defer q.startLock.Unlock()
	if !q.inited {
		return errors.New("not inited")
	}
	if q.started {
		return
	}

	defer func() {
		if err == nil {
			q.started = true
		}
	}()

	slowPattern, err := utils.ExtendUnifiedPattern(queueSlow, "")
	if err != nil {
		return
	}
	defaultPattern, err := utils.ExtendUnifiedPattern(queueDefault, "")
	if err != nil {
		return
	}

	mux := asynq.NewServeMux()
	mux.HandleFunc(slowPattern, handlers.GenHandlerFunc(q.rateCtrl, "slow", q.cfg.SlowQueueQps))
	mux.HandleFunc(defaultPattern, handlers.GenHandlerFunc(q.rateCtrl, "default", q.cfg.DefaultQueueQps))

	err = q.server.Start(mux)
	if err != nil {
		return
	}

	q.startWebUI()

	return
}

func (q *XcAsynq) Stop() {
	q.server.Stop()
	q.server.Shutdown()
}

// assignQueue 指定使用哪个队列，并给队列名加特定的前缀
func (q *XcAsynq) assignQueue(queue string, opts []asynq.Option) (result []asynq.Option, err error) {
	// opts里不能有对queue的配置，要报错
	for _, v := range opts {
		if v.Type() == asynq.QueueOpt {
			err = errors.Errorf("can not assign queue(%v) manually", v.Value())
			return
		}
	}
	result = append(opts, asynq.Queue(q.queuePrefix+queue))
	return
}

// assignDefaultMaxRetry 如果没有配置重试次数，补上默认的
func (q *XcAsynq) assignDefaultMaxRetry(opts []asynq.Option) []asynq.Option {
	for _, v := range opts {
		if v.Type() == asynq.MaxRetryOpt {
			return opts
		}
	}
	return append(opts, asynq.MaxRetry(3))
}

func (q *XcAsynq) addTask(ctx context.Context, queue, typ string, payload []byte, opts ...asynq.Option) (
	info *asynq.TaskInfo, err error) {
	opts, err = q.assignQueue(queue, opts)
	if err != nil {
		err = errors.Wrap(err, "assignQueue")
		return
	}
	opts = q.assignDefaultMaxRetry(opts)

	typ, err = utils.ExtendUnifiedPattern(queue, typ)
	if err != nil {
		err = errors.Wrap(err, "ExtendUnifiedPattern")
		return
	}

	task := asynq.NewTask(typ, payload, opts...)
	info, err = q.client.EnqueueContext(ctx, task)
	if err != nil {
		err = errors.Wrap(err, "EnqueueContext")
		return
	}
	return
}

func (q *XcAsynq) startWebUI() {
	h := asynqmon.New(asynqmon.Options{
		RedisConnOpt: maker.NewMaker(q.rdb),
	})

	http.Handle(h.RootPath()+"/", h)

	go http.ListenAndServe(":8889", nil)
}