feat: add independent platform modules and multi-provider storage
This commit is contained in:
213
queue/redis/consumer.go
Normal file
213
queue/redis/consumer.go
Normal file
@@ -0,0 +1,213 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user