497 lines
15 KiB
Go
497 lines
15 KiB
Go
package redisqueue
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"os"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"github.com/centrifugal/centrifugo/v6/internal/redisshard"
|
||
|
||
"github.com/redis/rueidis"
|
||
)
|
||
|
||
// ConsumerFunc is a type alias for the functions that will be used to handle
|
||
// and process Messages.
|
||
type ConsumerFunc func(*Message) error
|
||
|
||
// ConsumerOptions provide options to configure the Consumer.
|
||
type ConsumerOptions struct {
|
||
// Stream is stream to consume.
|
||
Stream string
|
||
// GroupName sets the name of the consumer group. This will be used when
|
||
// coordinating in Redis. If empty, stream name will be used.
|
||
GroupName string
|
||
// Name sets the name of this consumer. This will be used when fetching from
|
||
// Redis. If empty, the hostname will be used.
|
||
Name string
|
||
|
||
// ConsumerFunc will be called upon task message received from Redis.
|
||
ConsumerFunc ConsumerFunc
|
||
|
||
// VisibilityTimeout dictates the maximum amount of time a message should
|
||
// stay in pending. If there is a message that has been idle for more than
|
||
// this duration, the consumer will attempt to claim it.
|
||
VisibilityTimeout time.Duration
|
||
// BlockingTimeout designates how long the XREADGROUP call blocks for. If
|
||
// this is 0, it will block indefinitely. While this is the most efficient
|
||
// from a polling perspective, if this call never times out, there is no
|
||
// opportunity to yield back to Go at a regular interval. This means it's
|
||
// possible that if no messages are coming in, the consumer cannot
|
||
// gracefully shutdown. Instead, it's recommended to set this to 1-5
|
||
// seconds, or even longer, depending on how long your application can wait
|
||
// for shutdown.
|
||
BlockingTimeout time.Duration
|
||
// ReclaimInterval is the amount of time in between calls to XPENDING to
|
||
// attempt to reclaim jobs that have been idle for more than the visibility
|
||
// timeout. A smaller duration will result in more frequent checks. This
|
||
// will allow messages to be reaped faster, but it will put more load on
|
||
// Redis.
|
||
ReclaimInterval time.Duration
|
||
// Concurrency dictates how many goroutines to spawn to handle the messages.
|
||
Concurrency int
|
||
// UseLegacyReclaim enables using xpending + xclaim instead of xautoclaim.
|
||
UseLegacyReclaim bool
|
||
}
|
||
|
||
// Consumer adds a convenient wrapper around consuming jobs and managing concurrency.
|
||
type Consumer struct {
|
||
// Errors is a channel that you can receive from to centrally handle any
|
||
// errors that may occur either by your ConsumerFunc or by internal
|
||
// processing functions. Because this is an unbuffered channel, you must
|
||
// have a listener on it. If you don't – parts of the consumer could stop
|
||
// functioning when errors occur due to the blocking nature of unbuffered
|
||
// channels.
|
||
Errors chan error
|
||
|
||
options ConsumerOptions
|
||
queue chan *Message
|
||
cond *sync.Cond
|
||
inflight int64
|
||
|
||
stopReclaim chan struct{}
|
||
stopPoll chan struct{}
|
||
stopWorkers chan struct{}
|
||
|
||
shard *redisshard.RedisShard
|
||
}
|
||
|
||
// NewConsumer creates a Consumer with custom ConsumerOptions.
|
||
func NewConsumer(shard *redisshard.RedisShard, options ConsumerOptions) (*Consumer, error) {
|
||
if options.Stream == "" {
|
||
return nil, errors.New("stream required")
|
||
}
|
||
if options.ConsumerFunc == nil {
|
||
return nil, errors.New("ConsumerFunc required")
|
||
}
|
||
if options.GroupName == "" {
|
||
return nil, errors.New("GroupName required")
|
||
}
|
||
if options.Name == "" {
|
||
hostname, _ := os.Hostname()
|
||
options.Name = hostname
|
||
}
|
||
if options.BlockingTimeout == 0 {
|
||
options.BlockingTimeout = 15 * time.Second
|
||
}
|
||
if options.ReclaimInterval == 0 {
|
||
options.ReclaimInterval = 5 * time.Second
|
||
}
|
||
|
||
c := &Consumer{
|
||
Errors: make(chan error),
|
||
|
||
options: options,
|
||
queue: make(chan *Message, options.Concurrency),
|
||
cond: sync.NewCond(&sync.Mutex{}),
|
||
|
||
stopReclaim: make(chan struct{}, 1),
|
||
stopPoll: make(chan struct{}, 1),
|
||
stopWorkers: make(chan struct{}),
|
||
|
||
shard: shard,
|
||
}
|
||
err := CreateConsumerGroup(c.shard, c.options.Stream, c.options.GroupName, "$")
|
||
if err != nil && !strings.Contains(err.Error(), "BUSYGROUP") {
|
||
return nil, fmt.Errorf("error creating consumer group: %w", err)
|
||
}
|
||
return c, nil
|
||
}
|
||
|
||
func (c *Consumer) Options() ConsumerOptions {
|
||
return c.options
|
||
}
|
||
|
||
// Run starts all the worker goroutines and starts processing from the
|
||
// streams that have been registered with Register. All errors will be sent to
|
||
// the Errors channel. If Register was never called, an error will be sent and
|
||
// Run will terminate early. The same will happen if an error occurs when
|
||
// creating the consumer group in Redis. Run will block until Shutdown is called
|
||
// and all the in-flight messages have been processed.
|
||
func (c *Consumer) Run() {
|
||
if c.options.UseLegacyReclaim {
|
||
go c.reclaimLegacy(c.options.Stream)
|
||
} else {
|
||
go c.reclaim(c.options.Stream)
|
||
}
|
||
|
||
var pollWg sync.WaitGroup
|
||
pollWg.Add(1)
|
||
|
||
go func() {
|
||
defer pollWg.Done()
|
||
c.poll()
|
||
}()
|
||
|
||
// NOTE: the upstream library installed its own process-global SIGINT/SIGTERM
|
||
// handler here that os.Exit(1)s on a second signal. That hijacks the embedding
|
||
// server's graceful shutdown (a second signal aborts it, skipping cleanup) and
|
||
// leaks goroutines. Centrifugo drives shutdown via Shutdown() on context
|
||
// cancellation (see consuming.RedisStreamConsumer.Run), so no signal handler
|
||
// is installed here.
|
||
|
||
var workerWg sync.WaitGroup
|
||
workerWg.Add(c.options.Concurrency)
|
||
|
||
for i := 0; i < c.options.Concurrency; i++ {
|
||
go func() {
|
||
defer workerWg.Done()
|
||
c.work()
|
||
}()
|
||
}
|
||
|
||
pollWg.Wait()
|
||
close(c.stopWorkers)
|
||
workerWg.Wait()
|
||
}
|
||
|
||
func (c *Consumer) InflightJobs() int64 {
|
||
return atomic.LoadInt64(&c.inflight)
|
||
}
|
||
|
||
// Shutdown stops new messages from being processed and tells the workers to
|
||
// wait until all in-flight messages have been processed, and then they exit.
|
||
// The order that things stop is 1) the reclaim process (if it's running), 2)
|
||
// the polling process, and 3) the worker processes.
|
||
func (c *Consumer) Shutdown() {
|
||
c.stopReclaim <- struct{}{}
|
||
if c.options.VisibilityTimeout == 0 {
|
||
c.stopPoll <- struct{}{}
|
||
}
|
||
c.cond.L.Lock()
|
||
c.cond.Broadcast()
|
||
c.cond.L.Unlock()
|
||
}
|
||
|
||
// reclaim runs in a separate goroutine and checks the list of pending messages
|
||
// in every stream. For every message, if it's been idle for longer than the
|
||
// VisibilityTimeout, it will attempt to claim that message for this consumer.
|
||
// If VisibilityTimeout is 0, this function returns early and no messages are
|
||
// reclaimed. It checks the list of pending messages according to
|
||
// ReclaimInterval.
|
||
func (c *Consumer) reclaim(stream string) {
|
||
if c.options.VisibilityTimeout == 0 {
|
||
return
|
||
}
|
||
|
||
ticker := time.NewTicker(c.options.ReclaimInterval)
|
||
defer ticker.Stop()
|
||
|
||
for {
|
||
select {
|
||
case <-c.stopReclaim:
|
||
// once the reclaim process has stopped, stop the polling process
|
||
c.stopPoll <- struct{}{}
|
||
return
|
||
case <-ticker.C:
|
||
start := "-"
|
||
for {
|
||
c.cond.L.Lock()
|
||
inflight := atomic.LoadInt64(&c.inflight)
|
||
if inflight >= int64(c.options.Concurrency) {
|
||
c.cond.Wait()
|
||
c.cond.L.Unlock()
|
||
continue
|
||
}
|
||
c.cond.L.Unlock()
|
||
count := c.options.Concurrency - len(c.queue)
|
||
if count <= 0 {
|
||
count = 1
|
||
}
|
||
res := c.shard.RunOp(func(client rueidis.Client) rueidis.RedisResult {
|
||
cmd := client.B().Xautoclaim().
|
||
Key(stream).Group(c.options.GroupName).Consumer(c.options.Name).
|
||
MinIdleTime(strconv.FormatInt(c.options.VisibilityTimeout.Milliseconds(), 10)).
|
||
Start(start).Count(int64(count)).Build()
|
||
return client.Do(context.Background(), cmd)
|
||
})
|
||
entries, newStart, err := NewXAutoClaimCmd(res).Result()
|
||
if err != nil && !rueidis.IsRedisNil(err) {
|
||
c.logError(fmt.Errorf("error listing pending messages: %w", err))
|
||
break
|
||
}
|
||
|
||
if rueidis.IsRedisNil(err) {
|
||
break
|
||
}
|
||
|
||
start = newStart
|
||
|
||
if len(entries) > 0 {
|
||
c.enqueue(entries)
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// reclaimLegacy runs in a separate goroutine and checks the list of pending messages
|
||
// in every stream. For every message, if it's been idle for longer than the
|
||
// VisibilityTimeout, it will attempt to claim that message for this consumer.
|
||
// If VisibilityTimeout is 0, this function returns early and no messages are
|
||
// reclaimed. It checks the list of pending messages according to
|
||
// ReclaimInterval.
|
||
func (c *Consumer) reclaimLegacy(stream string) {
|
||
if c.options.VisibilityTimeout == 0 {
|
||
return
|
||
}
|
||
|
||
ticker := time.NewTicker(c.options.ReclaimInterval)
|
||
|
||
for {
|
||
select {
|
||
case <-c.stopReclaim:
|
||
// once the reclaim process has stopped, stop the polling process
|
||
c.stopPoll <- struct{}{}
|
||
return
|
||
case <-ticker.C:
|
||
start := "-"
|
||
end := "+"
|
||
|
||
for {
|
||
res := xPendingExt(c.shard, stream, c.options.GroupName, start, end, int64(c.options.Concurrency-len(c.queue)))
|
||
if res.Error() != nil {
|
||
c.logError(fmt.Errorf("error getting pending messages: %w", res.Error()))
|
||
break
|
||
}
|
||
xPendingRes := NewXPendingExtCmd(res)
|
||
if xPendingRes.Err != nil {
|
||
c.logError(fmt.Errorf("error parsing pending messages: %w", xPendingRes.Err))
|
||
break
|
||
}
|
||
if len(xPendingRes.Val) == 0 {
|
||
break
|
||
}
|
||
for _, r := range xPendingRes.Val {
|
||
if r.Idle >= c.options.VisibilityTimeout {
|
||
res = c.shard.RunOp(func(client rueidis.Client) rueidis.RedisResult {
|
||
cmd := client.B().Xclaim().
|
||
Key(stream).Group(c.options.GroupName).Consumer(c.options.Name).
|
||
MinIdleTime(strconv.FormatInt(c.options.VisibilityTimeout.Milliseconds(), 10)).
|
||
Id(r.ID).Build()
|
||
return client.Do(context.Background(), cmd)
|
||
})
|
||
if res.Error() != nil {
|
||
c.logError(fmt.Errorf("error claiming pending message: %w", res.Error()))
|
||
break
|
||
}
|
||
|
||
claimedEntries, err := res.AsXRange()
|
||
if err != nil {
|
||
c.logError(fmt.Errorf("error parsing pending messages: %w", err))
|
||
break
|
||
}
|
||
|
||
// If the no message returned, it means that
|
||
// the message no longer exists in the stream.
|
||
// However, it is still in a pending state. This
|
||
// could happen if a message was claimed by a
|
||
// consumer, that consumer died, and the message
|
||
// gets deleted (either through a XDEL call or
|
||
// through MAXLEN). Since the message no longer
|
||
// exists, the only way we can get it out of the
|
||
// pending state is to acknowledge it.
|
||
if len(claimedEntries) == 0 {
|
||
res := c.shard.RunOp(func(client rueidis.Client) rueidis.RedisResult {
|
||
cmd := client.B().Xack().Key(stream).Group(c.options.GroupName).Id(r.ID).Build()
|
||
return client.Do(context.Background(), cmd)
|
||
})
|
||
if res.Error() != nil {
|
||
c.logError(fmt.Errorf("error acknowledging after failed claim for %q stream and %q message: %w", stream, r.ID, res.Error()))
|
||
continue
|
||
}
|
||
}
|
||
c.enqueue(claimedEntries)
|
||
}
|
||
}
|
||
|
||
newID, err := incrementMessageID(xPendingRes.Val[len(xPendingRes.Val)-1].ID)
|
||
if err != nil {
|
||
c.logError(fmt.Errorf("error incrementing message ID: %w", err))
|
||
break
|
||
}
|
||
|
||
start = newID
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func xPendingExt(shard *redisshard.RedisShard, stream, groupName, start, end string, count int64) rueidis.RedisResult {
|
||
return shard.RunOp(func(client rueidis.Client) rueidis.RedisResult {
|
||
cmd := client.B().Xpending().Key(stream).Group(groupName).Start(start).End(end).Count(count).Build()
|
||
return client.Do(context.Background(), cmd)
|
||
})
|
||
}
|
||
|
||
// incrementMessageID takes in a message ID (e.g. 1564886140363-0) and
|
||
// increments the index section (e.g. 1564886140363-1). This is the next valid
|
||
// ID value, and it can be used for paging through messages.
|
||
func incrementMessageID(id string) (string, error) {
|
||
parts := strings.Split(id, "-")
|
||
index := parts[1]
|
||
parsed, err := strconv.ParseInt(index, 10, 64)
|
||
if err != nil {
|
||
return "", fmt.Errorf("error parsing message ID %q: %w", id, err)
|
||
}
|
||
return fmt.Sprintf("%s-%d", parts[0], parsed+1), nil
|
||
}
|
||
|
||
func (c *Consumer) xReadGroup(stream, groupName, name string, count int64, blockingMilliseconds int64) rueidis.RedisResult {
|
||
return c.shard.RunOp(func(client rueidis.Client) rueidis.RedisResult {
|
||
cmd := client.B().Xreadgroup().Group(groupName, name).Count(count).Block(blockingMilliseconds).Streams().Key(stream).Id(">").Build()
|
||
return client.Do(context.Background(), cmd)
|
||
})
|
||
}
|
||
|
||
// poll constantly checks the streams using XREADGROUP to see if there are any
|
||
// messages for this consumer to process. It blocks for up to 5 seconds instead
|
||
// of blocking indefinitely so that it can periodically check to see if Shutdown
|
||
// was called.
|
||
func (c *Consumer) poll() {
|
||
for {
|
||
select {
|
||
case <-c.stopPoll:
|
||
return
|
||
default:
|
||
c.cond.L.Lock()
|
||
inflight := atomic.LoadInt64(&c.inflight)
|
||
if inflight >= int64(c.options.Concurrency) {
|
||
c.cond.Wait()
|
||
c.cond.L.Unlock()
|
||
continue
|
||
}
|
||
c.cond.L.Unlock()
|
||
count := c.options.Concurrency - len(c.queue)
|
||
|
||
res := c.xReadGroup(c.options.Stream, c.options.GroupName, c.options.Name, int64(count), c.options.BlockingTimeout.Milliseconds())
|
||
err := res.Error()
|
||
if err != nil {
|
||
if err, ok := err.(net.Error); ok && err.Timeout() {
|
||
continue
|
||
}
|
||
if rueidis.IsRedisNil(err) {
|
||
continue
|
||
}
|
||
|
||
if strings.Contains(err.Error(), "NOGROUP") {
|
||
err := CreateConsumerGroup(c.shard, c.options.Stream, c.options.GroupName, "$")
|
||
if err != nil {
|
||
c.logError(fmt.Errorf("error creating consumer group: %w", err))
|
||
}
|
||
continue
|
||
}
|
||
|
||
c.logError(fmt.Errorf("error reading redis stream %s: %w", c.options.Stream, err))
|
||
select {
|
||
case <-c.stopPoll:
|
||
return
|
||
case <-time.After(200 * time.Millisecond):
|
||
}
|
||
continue
|
||
}
|
||
|
||
xRead, err := res.AsXRead()
|
||
if err != nil {
|
||
c.logError(fmt.Errorf("error parsing redis stream %s: %w", c.options.Stream, err))
|
||
continue
|
||
}
|
||
|
||
for _, messages := range xRead {
|
||
c.enqueue(messages)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Consumer) logError(err error) {
|
||
select {
|
||
case c.Errors <- err:
|
||
default:
|
||
}
|
||
}
|
||
|
||
// enqueue takes a slice of XMessages, creates corresponding Messages, and sends
|
||
// them on the centralized channel for worker goroutines to process.
|
||
func (c *Consumer) enqueue(messages []rueidis.XRangeEntry) {
|
||
for _, m := range messages {
|
||
msg := &Message{
|
||
ID: m.ID,
|
||
Values: m.FieldValues,
|
||
}
|
||
select {
|
||
case c.queue <- msg:
|
||
atomic.AddInt64(&c.inflight, 1)
|
||
case <-c.stopWorkers:
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
// work is called in a separate goroutine. The number of work goroutines is
|
||
// determined by Concurrency. Once it gets a message from the centralized
|
||
// channel, it calls the corresponding ConsumerFunc depending on the stream it
|
||
// came from. If no error is returned from the ConsumerFunc, the message is
|
||
// acknowledged in Redis.
|
||
func (c *Consumer) work() {
|
||
for {
|
||
select {
|
||
case msg := <-c.queue:
|
||
err := c.process(msg)
|
||
if err != nil {
|
||
c.logError(fmt.Errorf("error calling processing func for %q stream and %q message: %w", c.options.Stream, msg.ID, err))
|
||
continue
|
||
}
|
||
err = c.shard.RunOp(func(client rueidis.Client) rueidis.RedisResult {
|
||
cmd := client.B().Xack().Key(c.options.Stream).Group(c.options.GroupName).Id(msg.ID).Build()
|
||
return client.Do(context.Background(), cmd)
|
||
}).Error()
|
||
if err != nil {
|
||
c.logError(fmt.Errorf("error acknowledging message for %q stream and %q message: %w", c.options.Stream, msg.ID, err))
|
||
continue
|
||
}
|
||
case <-c.stopWorkers:
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
func (c *Consumer) process(msg *Message) (err error) {
|
||
defer func() {
|
||
c.cond.L.Lock()
|
||
atomic.AddInt64(&c.inflight, -1)
|
||
c.cond.Broadcast()
|
||
c.cond.L.Unlock()
|
||
}()
|
||
err = c.options.ConsumerFunc(msg)
|
||
return
|
||
}
|