214 lines
6.6 KiB
Go
214 lines
6.6 KiB
Go
package redis
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"git.yuntuo.hk/k2safe/go-omni/queue"
|
|
goredis "github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
type ConsumerConfig struct {
|
|
Workers int `json:"workers"`
|
|
MaxAttempts int64 `json:"maxAttempts"`
|
|
TaskTimeout time.Duration `json:"taskTimeout"`
|
|
VisibilityTimeout time.Duration `json:"visibilityTimeout"`
|
|
ShutdownTimeout time.Duration `json:"shutdownTimeout"`
|
|
}
|
|
type Consumer struct {
|
|
store *Store
|
|
cfg ConsumerConfig
|
|
used atomic.Bool
|
|
}
|
|
|
|
func NewConsumer(store *Store, cfg ConsumerConfig) (*Consumer, error) {
|
|
if store == nil || cfg.Workers < 0 || cfg.MaxAttempts < 0 || cfg.TaskTimeout < 0 || cfg.VisibilityTimeout < 0 || cfg.ShutdownTimeout < 0 {
|
|
return nil, errors.New("queue/redis: invalid consumer configuration")
|
|
}
|
|
if cfg.Workers == 0 {
|
|
cfg.Workers = 4
|
|
}
|
|
if cfg.MaxAttempts == 0 {
|
|
cfg.MaxAttempts = 3
|
|
}
|
|
if cfg.TaskTimeout == 0 {
|
|
cfg.TaskTimeout = 30 * time.Second
|
|
}
|
|
if cfg.VisibilityTimeout == 0 {
|
|
cfg.VisibilityTimeout = 2 * cfg.TaskTimeout
|
|
}
|
|
if cfg.ShutdownTimeout == 0 {
|
|
cfg.ShutdownTimeout = 5 * time.Second
|
|
}
|
|
if cfg.Workers > 128 || cfg.MaxAttempts > 100 || cfg.TaskTimeout < 10*time.Millisecond || cfg.TaskTimeout > time.Hour || cfg.VisibilityTimeout < 2*cfg.TaskTimeout || cfg.VisibilityTimeout > 24*time.Hour || cfg.ShutdownTimeout > time.Minute {
|
|
return nil, errors.New("queue/redis: consumer limits exceeded")
|
|
}
|
|
return &Consumer{store: store, cfg: cfg}, nil
|
|
}
|
|
|
|
// Run is called once. Cancellation stops fetching and cancels handlers; ACK is
|
|
// deferred to a later delivery if cancellation makes the outcome uncertain.
|
|
// A handler must honor context cancellation. Go cannot terminate an arbitrary
|
|
// function; an ignoring handler can outlive ShutdownTimeout (at most Workers).
|
|
func (c *Consumer) Run(ctx context.Context, handler queue.Handler) error {
|
|
if ctx == nil || handler == nil {
|
|
return errors.New("queue/redis: context and handler required")
|
|
}
|
|
if !c.used.CompareAndSwap(false, true) {
|
|
return errors.New("queue/redis: consumer already run")
|
|
}
|
|
signature := fmt.Sprintf("%d:%d:%d", c.cfg.MaxAttempts, c.cfg.TaskTimeout.Milliseconds(), c.cfg.VisibilityTimeout.Milliseconds())
|
|
if err := goredis.NewScript(`local v=redis.call('GET',KEYS[1]);if v and v~=ARGV[1] then return redis.error_reply('consumer configuration mismatch') end;redis.call('SET',KEYS[1],ARGV[1]);return 1`).Run(ctx, c.store.client, []string{c.store.keys[5]}, signature).Err(); err != nil {
|
|
return err
|
|
}
|
|
if err := c.store.client.XGroupCreateMkStream(ctx, c.store.keys[0], "workers", "0").Err(); err != nil && !strings.HasPrefix(err.Error(), "BUSYGROUP ") {
|
|
return err
|
|
}
|
|
var token [16]byte
|
|
if _, err := rand.Read(token[:]); err != nil {
|
|
return err
|
|
}
|
|
runCtx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
failures := make(chan error, c.cfg.Workers)
|
|
var wg sync.WaitGroup
|
|
for n := 0; n < c.cfg.Workers; n++ {
|
|
wg.Add(1)
|
|
go func(n int) {
|
|
defer wg.Done()
|
|
if err := c.loop(runCtx, hex.EncodeToString(token[:])+fmt.Sprint(n), handler); err != nil {
|
|
failures <- err
|
|
cancel()
|
|
}
|
|
}(n)
|
|
}
|
|
done := make(chan struct{})
|
|
go func() { wg.Wait(); close(done) }()
|
|
select {
|
|
case <-done:
|
|
case <-runCtx.Done():
|
|
select {
|
|
case <-done:
|
|
case <-time.After(c.cfg.ShutdownTimeout):
|
|
return errors.New("queue/redis: handler shutdown timeout")
|
|
}
|
|
}
|
|
select {
|
|
case err := <-failures:
|
|
return err
|
|
default:
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
func (c *Consumer) loop(ctx context.Context, owner string, handler queue.Handler) error {
|
|
s := c.store
|
|
defer func() {
|
|
cleanup, cancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancel()
|
|
_ = goredis.NewScript(`local p=redis.call('XPENDING',KEYS[1],'workers','-','+',1,ARGV[1]);if #p==0 then return redis.call('XGROUP','DELCONSUMER',KEYS[1],'workers',ARGV[1]) end;return 0`).Run(cleanup, s.client, []string{s.keys[0]}, owner).Err()
|
|
}()
|
|
cursor := "0-0"
|
|
for ctx.Err() == nil {
|
|
messages, next, err := s.client.XAutoClaim(ctx, &goredis.XAutoClaimArgs{Stream: s.keys[0], Group: "workers", Consumer: owner, MinIdle: c.cfg.VisibilityTimeout, Start: cursor, Count: 1}).Result()
|
|
if err != nil && err != goredis.Nil {
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
cursor = next
|
|
if cursor == "" {
|
|
cursor = "0-0"
|
|
}
|
|
if len(messages) == 0 {
|
|
streams, err := s.client.XReadGroup(ctx, &goredis.XReadGroupArgs{Group: "workers", Consumer: owner, Streams: []string{s.keys[0], ">"}, Count: 1, Block: -1}).Result()
|
|
if err != nil && err != goredis.Nil {
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
if len(streams) > 0 {
|
|
messages = streams[0].Messages
|
|
}
|
|
}
|
|
if len(messages) == 0 {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil
|
|
case <-time.After(25 * time.Millisecond):
|
|
}
|
|
continue
|
|
}
|
|
message := messages[0]
|
|
raw, ok := message.Values["message"].(string)
|
|
var m queue.Message
|
|
if !ok || len(raw) > s.cfg.MaxPayloadBytes+1024 || json.Unmarshal([]byte(raw), &m) != nil {
|
|
return errors.New("queue/redis: corrupt message; operator review required")
|
|
}
|
|
pending, err := s.client.XPendingExt(ctx, &goredis.XPendingExtArgs{Stream: s.keys[0], Group: "workers", Start: message.ID, End: message.ID, Count: 1}).Result()
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
if len(pending) != 1 || pending[0].Consumer != owner {
|
|
continue
|
|
}
|
|
delivery := queue.Delivery{Message: m, Attempt: pending[0].RetryCount}
|
|
reason := ""
|
|
if delivery.Attempt > c.cfg.MaxAttempts {
|
|
reason = "attemptsExhausted"
|
|
} else {
|
|
workCtx, cancel := context.WithTimeout(ctx, c.cfg.TaskTimeout)
|
|
err = invoke(workCtx, handler, delivery)
|
|
deadlineErr := workCtx.Err()
|
|
cancel()
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
var rejection *queue.Rejection
|
|
if errors.As(err, &rejection) {
|
|
reason = "rejected"
|
|
} else if err != nil || deadlineErr != nil {
|
|
if delivery.Attempt < c.cfg.MaxAttempts {
|
|
continue
|
|
}
|
|
reason = "attemptsExhausted"
|
|
}
|
|
}
|
|
ackCtx, ackCancel := context.WithTimeout(ctx, 3*time.Second)
|
|
err = s.finish(ackCtx, message.ID, owner, delivery, reason)
|
|
ackCancel()
|
|
if err != nil && !errors.Is(err, queue.ErrLeaseLost) {
|
|
if ctx.Err() != nil {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
func invoke(ctx context.Context, handler queue.Handler, delivery queue.Delivery) (err error) {
|
|
defer func() {
|
|
if recover() != nil {
|
|
err = errors.New("queue: handler panic")
|
|
}
|
|
}()
|
|
if err = ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
return handler(ctx, delivery)
|
|
}
|
|
|
|
var _ queue.Consumer = (*Consumer)(nil)
|