go-omni/queue/redis/consumer.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)