Add support for registering consumer at specific Message ID (robinjoseph08#7)

Author: theckman

.go-version

@@ -1 +1 @@
-1.12.7
+1.14.2

.travis.yml

@@ -5,7 +5,7 @@ addons:
   apt:
     packages:
       - redis-server
-go: "1.12.7"
+go: "1.14.2"
 env:
   - GO111MODULE=on
 install:

consumer.go

@@ -14,6 +14,11 @@ import (
 // and process Messages.
 type ConsumerFunc func(*Message) error
 
+type registeredConsumer struct {
+	fn ConsumerFunc
+	id string
+}
+
 // ConsumerOptions provide options to configure the Consumer.
 type ConsumerOptions struct {
 	// Name sets the name of this consumer. This will be used when fetching from
@@ -62,12 +67,12 @@ type Consumer struct {
 	// channels.
 	Errors chan error
 
-	options *ConsumerOptions
-	redis   *redis.Client
-	funcs   map[string]ConsumerFunc
-	streams []string
-	queue   chan *Message
-	wg      *sync.WaitGroup
+	options   *ConsumerOptions
+	redis     *redis.Client
+	consumers map[string]registeredConsumer
+	streams   []string
+	queue     chan *Message
+	wg        *sync.WaitGroup
 
 	stopReclaim chan struct{}
 	stopPoll    chan struct{}
@@ -118,25 +123,43 @@ func NewConsumerWithOptions(options *ConsumerOptions) (*Consumer, error) {
 	return &Consumer{
 		Errors: make(chan error),
 
-		options: options,
-		redis:   r,
-		funcs:   map[string]ConsumerFunc{},
-		streams: make([]string, 0),
-		queue:   make(chan *Message, options.BufferSize),
-		wg:      &sync.WaitGroup{},
+		options:   options,
+		redis:     r,
+		consumers: make(map[string]registeredConsumer),
+		streams:   make([]string, 0),
+		queue:     make(chan *Message, options.BufferSize),
+		wg:        &sync.WaitGroup{},
 
 		stopReclaim: make(chan struct{}, 1),
 		stopPoll:    make(chan struct{}, 1),
 		stopWorkers: make(chan struct{}, options.Concurrency),
 	}, nil
 }
 
+// RegisterWithLastID is the same as Register, except that it also lets you
+// specify the oldest message to receive when first creating the consumer group.
+// This can be any valid message ID, "0" for all messages in the stream, or "$"
+// for only new messages.
+//
+// If the consumer group already exists the id field is ignored, meaning you'll
+// receive unprocessed messages.
+func (c *Consumer) RegisterWithLastID(stream string, id string, fn ConsumerFunc) {
+	if len(id) == 0 {
+		id = "0"
+	}
+
+	c.consumers[stream] = registeredConsumer{
+		fn: fn,
+		id: id,
+	}
+}
+
 // Register takes in a stream name and a ConsumerFunc that will be called when a
 // message comes in from that stream. Register must be called at least once
 // before Run is called. If the same stream name is passed in twice, the first
 // ConsumerFunc is overwritten by the second.
 func (c *Consumer) Register(stream string, fn ConsumerFunc) {
-	c.funcs[stream] = fn
+	c.RegisterWithLastID(stream, "0", fn)
 }
 
 // Run starts all of the worker goroutines and starts processing from the
@@ -146,22 +169,22 @@ func (c *Consumer) Register(stream string, fn ConsumerFunc) {
 // creating the consumer group in Redis. Run will block until Shutdown is called
 // and all of the in-flight messages have been processed.
 func (c *Consumer) Run() {
-	if len(c.funcs) == 0 {
+	if len(c.consumers) == 0 {
 		c.Errors <- errors.New("at least one consumer function needs to be registered")
 		return
 	}
 
-	for stream := range c.funcs {
+	for stream, consumer := range c.consumers {
 		c.streams = append(c.streams, stream)
-		err := c.redis.XGroupCreateMkStream(stream, c.options.GroupName, "0").Err()
+		err := c.redis.XGroupCreateMkStream(stream, c.options.GroupName, consumer.id).Err()
 		// ignoring the BUSYGROUP error makes this a noop
 		if err != nil && err.Error() != "BUSYGROUP Consumer Group name already exists" {
 			c.Errors <- errors.Wrap(err, "error creating consumer group")
 			return
 		}
 	}
 
-	for i := 0; i < len(c.funcs); i++ {
+	for i := 0; i < len(c.consumers); i++ {
 		c.streams = append(c.streams, ">")
 	}
 
@@ -214,7 +237,7 @@ func (c *Consumer) reclaim() {
 			c.stopPoll <- struct{}{}
 			return
 		case <-ticker.C:
-			for stream := range c.funcs {
+			for stream := range c.consumers {
 				start := "-"
 				end := "+"
 
@@ -373,7 +396,6 @@ func (c *Consumer) process(msg *Message) (err error) {
 			err = errors.Errorf("ConsumerFunc panic: %v", r)
 		}
 	}()
-	fn := c.funcs[msg.Stream]
-	err = fn(msg)
+	err = c.consumers[msg.Stream].fn(msg)
 	return
 }

consumer_test.go

@@ -77,10 +77,52 @@ func TestRegister(t *testing.T) {
 
 		c.Register(tt.Name(), fn)
 
-		assert.Len(tt, c.funcs, 1)
+		assert.Len(tt, c.consumers, 1)
 	})
 }
 
+func TestRegisterWithLastID(t *testing.T) {
+	fn := func(msg *Message) error {
+		return nil
+	}
+
+	tests := []struct {
+		name   string
+		stream string
+		id     string
+		want   map[string]registeredConsumer
+	}{
+		{
+			name: "custom_id",
+			id:   "42",
+			want: map[string]registeredConsumer{
+				"test": {id: "42", fn: fn},
+			},
+		},
+		{
+			name: "no_id",
+			id:   "",
+			want: map[string]registeredConsumer{
+				"test": {id: "0", fn: fn},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c, err := NewConsumer()
+			require.NoError(t, err)
+
+			c.RegisterWithLastID("test", tt.id, fn)
+
+			assert.Len(t, c.consumers, 1)
+			assert.Contains(t, c.consumers, "test")
+			assert.Equal(t, c.consumers["test"].id, tt.want["test"].id)
+			assert.NotNil(t, c.consumers["test"].fn)
+		})
+	}
+}
+
 func TestRun(t *testing.T) {
 	t.Run("sends an error if no ConsumerFuncs are registered", func(tt *testing.T) {
 		c, err := NewConsumer()