refactor ingest queue

Author: vm-001

pkg/queue/queue.go

@@ -6,22 +6,23 @@ import (
 )
 
 type Message struct {
-	ID          string
-	Data        []byte
+	Value       []byte
 	Time        time.Time
 	WorkspaceID string
 }
 
-type Options struct {
-	Count   int64
-	Block   bool
-	Timeout time.Duration
-}
+type HandleFunc func(ctx context.Context, messages []*Message) error
 
 type Queue interface {
-	Enqueue(ctx context.Context, message *Message) error
-	Dequeue(ctx context.Context, opts *Options) ([]*Message, error)
-	Delete(ctx context.Context, message []*Message) error
-	Size(ctx context.Context) (int64, error)
+	Producer
+	Consumer
 	Stats() map[string]interface{}
 }
+
+type Producer interface {
+	WriteMessage(ctx context.Context, message *Message) error
+}
+
+type Consumer interface {
+	StartListen(ctx context.Context, handle HandleFunc)
+}

pkg/queue/redis/redis.go

@@ -4,11 +4,9 @@ import (
 	"context"
 	"errors"
 	"github.com/redis/go-redis/v9"
-	"github.com/webhookx-io/webhookx/constants"
-	"github.com/webhookx-io/webhookx/pkg/metrics"
+	"github.com/webhookx-io/webhookx/pkg/loglimiter"
 	"github.com/webhookx-io/webhookx/pkg/queue"
 	"github.com/webhookx-io/webhookx/pkg/tracing"
-	"github.com/webhookx-io/webhookx/utils"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 	"strconv"
@@ -17,66 +15,53 @@ import (
 )
 
 type RedisQueue struct {
-	stream            string
-	group             string
-	consumer          string
-	visibilityTimeout time.Duration
-
+	opts    Options
 	c       *redis.Client
 	log     *zap.SugaredLogger
-	metrics *metrics.Metrics
+	limiter *loglimiter.Limiter
 }
 
-type RedisQueueOptions struct {
+type Options struct {
 	StreamName        string
-	GroupName         string
+	ConsumerGroupName string
 	ConsumerName      string
 	VisibilityTimeout time.Duration
+	Listeners         int
 
 	Client *redis.Client
 }
 
-func NewRedisQueue(opts RedisQueueOptions, logger *zap.SugaredLogger, metrics *metrics.Metrics) (queue.Queue, error) {
+func NewRedisQueue(opts Options, logger *zap.SugaredLogger) (queue.Queue, error) {
 	q := &RedisQueue{
-		stream:            utils.DefaultIfZero(opts.StreamName, constants.QueueRedisQueueName),
-		group:             utils.DefaultIfZero(opts.GroupName, constants.QueueRedisGroupName),
-		consumer:          utils.DefaultIfZero(opts.ConsumerName, constants.QueueRedisConsumerName),
-		visibilityTimeout: utils.DefaultIfZero(opts.VisibilityTimeout, constants.QueueRedisVisibilityTimeout),
-		c:                 opts.Client,
-		log:               logger,
-		metrics:           metrics,
-	}
-
-	go q.process()
-	if metrics.Enabled {
-		go q.monitoring()
+		opts:    opts,
+		c:       opts.Client,
+		log:     logger.Named("queue-redis"),
+		limiter: loglimiter.NewLimiter(time.Second),
 	}
-
 	return q, nil
 }
 
-func (q *RedisQueue) Enqueue(ctx context.Context, message *queue.Message) error {
+func (q *RedisQueue) WriteMessage(ctx context.Context, message *queue.Message) error {
 	ctx, span := tracing.Start(ctx, "redis.queue.enqueue", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
 	args := &redis.XAddArgs{
-		Stream: q.stream,
+		Stream: q.opts.StreamName,
 		ID:     "*",
-		Values: []interface{}{"data", message.Data, "time", message.Time.UnixMilli(), "ws_id", message.WorkspaceID},
-	}
-	res := q.c.XAdd(ctx, args)
-	if res.Err() != nil {
-		return res.Err()
+		Values: []interface{}{
+			"data", message.Value,
+			"time", message.Time.UnixMilli(),
+			"ws_id", message.WorkspaceID,
+		},
 	}
-	message.ID = res.Val()
-	return nil
+	return q.c.XAdd(ctx, args).Err()
 }
 
 func toMessage(values map[string]interface{}) *queue.Message {
 	message := &queue.Message{}
 
 	if data, ok := values["data"].(string); ok {
-		message.Data = []byte(data)
+		message.Value = []byte(data)
 	}
 
 	if timestr, ok := values["time"].(string); ok {
@@ -91,76 +76,97 @@ func toMessage(values map[string]interface{}) *queue.Message {
 	return message
 }
 
-func (q *RedisQueue) Dequeue(ctx context.Context, opt *queue.Options) ([]*queue.Message, error) {
+func (q *RedisQueue) dequeue(ctx context.Context) ([]redis.XMessage, error) {
 	ctx, span := tracing.Start(ctx, "redis.queue.dequeue", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
-	var count int64 = 1
-	if opt != nil && opt.Count != 0 {
-		count = opt.Count
-	}
-	var block time.Duration = -1
-	if opt != nil && opt.Block {
-		block = opt.Timeout
-	}
-
 	args := &redis.XReadGroupArgs{
-		Group:    q.group,
-		Consumer: q.consumer,
-		Streams:  []string{q.stream, ">"},
-		Count:    count,
-		Block:    block,
+		Group:    q.opts.ConsumerGroupName,
+		Consumer: q.opts.ConsumerName,
+		Streams:  []string{q.opts.StreamName, ">"},
+		Count:    20,
+		Block:    time.Second,
 	}
-	res := q.c.XReadGroup(ctx, args)
-	if res.Err() != nil {
-		err := res.Err()
+	streams, err := q.c.XReadGroup(ctx, args).Result()
+	if err != nil {
 		if errors.Is(err, redis.Nil) {
 			err = nil
 		} else if strings.HasPrefix(err.Error(), "NOGROUP") {
-			go q.createConsumerGroup()
 			err = nil
+			go q.createConsumerGroup(q.opts.StreamName, q.opts.ConsumerGroupName)
 		}
 		return nil, err
 	}
 
-	messages := make([]*queue.Message, 0)
-	for _, stream := range res.Val() {
-		for _, xmessage := range stream.Messages {
-			message := toMessage(xmessage.Values)
-			message.ID = xmessage.ID
-			messages = append(messages, message)
-		}
-	}
-
-	return messages, nil
+	return streams[0].Messages, nil
 }
 
-func (q *RedisQueue) Delete(ctx context.Context, messages []*queue.Message) error {
+func (q *RedisQueue) delete(ctx context.Context, xmessages []redis.XMessage) error {
 	ctx, span := tracing.Start(ctx, "redis.queue.delete", trace.WithSpanKind(trace.SpanKindServer))
 	defer span.End()
 
-	ids := make([]string, 0, len(messages))
-	for _, message := range messages {
+	ids := make([]string, 0, len(xmessages))
+	for _, message := range xmessages {
 		ids = append(ids, message.ID)
 	}
+
 	pipeline := q.c.Pipeline()
-	pipeline.XAck(ctx, q.stream, q.group, ids...)
-	pipeline.XDel(ctx, q.stream, ids...)
+	pipeline.XAck(ctx, q.opts.StreamName, q.opts.ConsumerGroupName, ids...)
+	pipeline.XDel(ctx, q.opts.StreamName, ids...)
 	_, err := pipeline.Exec(ctx)
-	if err != nil {
-		return err
+	return err
+}
+
+func (q *RedisQueue) StartListen(ctx context.Context, handle queue.HandleFunc) {
+	q.log.Infof("starting %d listeners", q.opts.Listeners)
+	for i := 0; i < q.opts.Listeners; i++ {
+		go q.listen(ctx, handle)
+	}
+	go q.process(ctx)
+}
+
+func (q *RedisQueue) listen(ctx context.Context, handle queue.HandleFunc) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			xmessages, err := q.dequeue(ctx)
+			if err != nil && q.limiter.Allow(err.Error()) {
+				q.log.Warnf("failed to dequeue: %v", err)
+				time.Sleep(time.Second)
+				continue
+			}
+			if len(xmessages) == 0 {
+				continue
+			}
+
+			messages := make([]*queue.Message, 0, len(xmessages))
+			for _, msg := range xmessages {
+				messages = append(messages, toMessage(msg.Values))
+			}
+
+			err = handle(ctx, messages)
+			if err != nil {
+				q.log.Warnf("failed to handle message: %v", err)
+				continue
+			}
+			err = q.delete(ctx, xmessages)
+			if err != nil {
+				q.log.Warnf("failed to delete message: %v", err)
+			}
+		}
 	}
-	return nil
 }
 
-func (q *RedisQueue) Size(ctx context.Context) (int64, error) {
-	return q.c.XLen(ctx, q.stream).Result()
+func (q *RedisQueue) size(ctx context.Context) (int64, error) {
+	return q.c.XLen(ctx, q.opts.StreamName).Result()
 }
 
 func (q *RedisQueue) Stats() map[string]interface{} {
 	stats := make(map[string]interface{})
 
-	size, err := q.Size(context.TODO())
+	size, err := q.size(context.TODO())
 	if err != nil {
 		q.log.Errorf("failed to retrieve status: %v", err)
 	}
@@ -169,20 +175,18 @@ func (q *RedisQueue) Stats() map[string]interface{} {
 	return stats
 }
 
-func (q *RedisQueue) createConsumerGroup() {
-	res := q.c.XGroupCreateMkStream(context.TODO(), q.stream, q.group, "0")
-	if res.Err() == nil {
-		q.log.Debugf("created default consumer group: %s", q.group)
+func (q *RedisQueue) createConsumerGroup(stream string, group string) {
+	err := q.c.XGroupCreateMkStream(context.TODO(), stream, group, "0").Err()
+	if err != nil {
+		if err.Error() != "BUSYGROUP Consumer Group name already exists" {
+			q.log.Errorf("failed to create Consumer Group '%s': %s", group, err.Error())
+		}
 		return
 	}
-
-	if res.Err().Error() != "BUSYGROUP Consumer Group name already exists" {
-		q.log.Errorf("failed to create the default consumer group: %s", res.Err().Error())
-	}
+	q.log.Debugf("Consumer Group '%s' created", group)
 }
 
-// process re-enqueue invisible messages that reach the visibility timeout
-func (q *RedisQueue) process() {
+func (q *RedisQueue) process(ctx context.Context) {
 	var reenqueueScript = redis.NewScript(`
 		local entries = redis.call('XPENDING', KEYS[1], KEYS[2], 'IDLE', ARGV[1], '-', '+', 1000)
 		local ids = {}
@@ -200,39 +204,25 @@ func (q *RedisQueue) process() {
 		return ids
 	`)
 
-	go func() {
-		ticker := time.NewTicker(time.Second)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-ticker.C:
-				keys := []string{q.stream, q.group}
-				argv := []interface{}{q.visibilityTimeout.Milliseconds()}
-				res, err := reenqueueScript.Run(context.Background(), q.c, keys, argv...).Result()
-				if err != nil {
-					q.log.Errorf("failed to reenqueue: %v", err)
-					continue
-				}
-
-				if ids, ok := res.([]interface{}); ok && len(ids) > 0 {
-					q.log.Debugf("enqueued invisible messages: %d", len(ids))
-				}
-			}
-		}
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
 
-	}()
-}
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			keys := []string{q.opts.StreamName, q.opts.ConsumerGroupName}
+			argv := []interface{}{q.opts.VisibilityTimeout.Milliseconds()}
+			res, err := reenqueueScript.Run(context.TODO(), q.c, keys, argv...).Result()
+			if err != nil {
+				q.log.Errorf("failed to reenqueue: %v", err)
+				continue
+			}
 
-func (q *RedisQueue) monitoring() {
-	ticker := time.NewTicker(q.metrics.Interval)
-	defer ticker.Stop()
-	for range ticker.C {
-		size, err := q.Size(context.TODO())
-		if err != nil {
-			q.log.Errorf("failed to get redis queue size: %v", err)
-			continue
+			if ids, ok := res.([]interface{}); ok && len(ids) > 0 {
+				q.log.Debugf("enqueued invisible messages: %d", len(ids))
+			}
 		}
-		q.metrics.EventPendingGauge.Set(float64(size))
 	}
 }