snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

2026-09-16 01:17:21 +05:00
commit b352a9a30d
472 changed files with 108375 additions and 0 deletions
--- a/internal/redisqueue/consumer.go
+++ b/internal/redisqueue/consumer.go
@@ -0,0 +1,496 @@
+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
+}