chore: implement sub message in redis

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

go.mod

@@ -5,4 +5,5 @@ go 1.16
 require (
 	github.com/go-redis/redis/v8 v8.11.3
 	github.com/golang-queue/queue v0.0.7
+	github.com/stretchr/testify v1.7.0
 )

go.sum

@@ -92,6 +92,7 @@ google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzi
 google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=

redis.go

@@ -2,6 +2,8 @@ package redisdb
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -20,7 +22,7 @@ type Option func(*Worker)
 type Worker struct {
 	// redis config
 	rdb              *redis.Client
-	pubsub           *redis.Pubsub
+	pubsub           *redis.PubSub
 	addr             string
 	db               int
 	connectionString string
@@ -229,14 +231,8 @@ func (s *Worker) Queue(job queue.QueuedMessage) error {
 
 	ctx := context.Background()
 
-	// Wait for confirmation that subscription is created before publishing anything.
-	_, err := s.pubsub.Receive(ctx)
-	if err != nil {
-		return err
-	}
-
 	// Publish a message.
-	err = s.rdb.Publish(ctx, s.channel, job).Err()
+	err := s.rdb.Publish(ctx, s.channel, job.Bytes()).Err()
 	if err != nil {
 		return err
 	}
@@ -253,5 +249,37 @@ func (s *Worker) Run() error {
 	default:
 	}
 
-	return nil
+	var options []redis.ChannelOption
+	ctx := context.Background()
+
+	if s.channelSize > 1 {
+		options = append(options, redis.WithChannelSize(s.channelSize))
+	}
+
+	ch := s.pubsub.Channel(options...)
+	// make sure the connection is successful
+	err := s.pubsub.Ping(ctx)
+	if err != nil {
+		return err
+	}
+
+	for {
+		select {
+		case m, ok := <-ch:
+			if !ok {
+				return fmt.Errorf("redis pubsub: channel=%s closed", s.channel)
+			}
+
+			var data queue.Job
+			if err := json.Unmarshal([]byte(m.Payload), &data); err != nil {
+				s.logger.Error("json unmarshal error: ", err)
+				continue
+			}
+			if err := s.handle(data); err != nil {
+				s.logger.Error("handle job error: ", err)
+			}
+		case <-s.stop:
+			return nil
+		}
+	}
 }

redis_test.go

@@ -1 +1,88 @@
 package redisdb
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/golang-queue/queue"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var host = "127.0.0.1"
+
+type mockMessage struct {
+	Message string
+}
+
+func (m mockMessage) Bytes() []byte {
+	return []byte(m.Message)
+}
+
+func TestRedisDefaultFlow(t *testing.T) {
+	m := &mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithAddr(host+":6379"),
+		WithChannel("test"),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(100 * time.Millisecond)
+	assert.NoError(t, q.Queue(m))
+	m.Message = "bar"
+	assert.NoError(t, q.Queue(m))
+	q.Shutdown()
+	q.Wait()
+}
+
+func TestNSQShutdown(t *testing.T) {
+	w := NewWorker(
+		WithAddr(host+":6379"),
+		WithChannel("test2"),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(1 * time.Second)
+	q.Shutdown()
+	// check shutdown once
+	assert.Error(t, w.Shutdown())
+	assert.Equal(t, queue.ErrQueueShutdown, w.Shutdown())
+	q.Wait()
+}
+
+func TestGoroutinePanic(t *testing.T) {
+	m := mockMessage{
+		Message: "foo",
+	}
+	w := NewWorker(
+		WithAddr(host+":6379"),
+		WithChannel("GoroutinePanic"),
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			panic("missing something")
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(50 * time.Millisecond)
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	time.Sleep(200 * time.Millisecond)
+	q.Shutdown()
+	assert.Error(t, q.Queue(m))
+	q.Wait()
+}