chore: Support redis cluster (#6)

* chore: Support redis cluster

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

* update

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

* update

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

Author: appleboy

.github/workflows/testing.yml

@@ -7,25 +7,15 @@ jobs:
     # You must use a Linux environment when using service containers or container jobs
     runs-on: ubuntu-latest
 
-    # Service containers to run with `container-job`
-    services:
-      # Label used to access the service container
-      redis:
-        # Docker Hub image
-        image: redis
-        # Set health checks to wait until redis has started
-        options: >-
-          --health-cmd "redis-cli ping"
-          --health-interval 10s
-          --health-timeout 5s
-          --health-retries 5
-        ports:
-          - 6379:6379
-
     env:
       GO111MODULE: on
       GOPROXY: https://proxy.golang.org
     steps:
+      - uses: vishnudxb/[email protected]
+        with:
+          master1-port: 6379
+          slave1-port: 6380
+
       - name: Set up Go ${{ matrix.go }}
         uses: actions/setup-go@v2
 

conf/redis.conf

@@ -0,0 +1,3 @@
+protected-mode no
+cluster-enabled yes
+appendonly yes

redis.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -20,7 +21,7 @@ type Option func(*Worker)
 // Worker for Redis
 type Worker struct {
 	// redis config
-	rdb              *redis.Client
+	rdb              redis.Cmdable
 	pubsub           *redis.PubSub
 	channel          <-chan *redis.Message
 	addr             string
@@ -29,6 +30,7 @@ type Worker struct {
 	password         string
 	channelName      string
 	channelSize      int
+	cluster          bool
 
 	stopOnce    sync.Once
 	stop        chan struct{}
@@ -52,6 +54,13 @@ func WithDB(db int) Option {
 	}
 }
 
+// WithCluster redis cluster
+func WithCluster(enable bool) Option {
+	return func(w *Worker) {
+		w.cluster = enable
+	}
+}
+
 // WithChannelSize redis channel size
 func WithChannelSize(size int) Option {
 	return func(w *Worker) {
@@ -114,32 +123,41 @@ func NewWorker(opts ...Option) *Worker {
 		opt(w)
 	}
 
-	var options *redis.Options
-
 	if w.connectionString != "" {
-		options, err = redis.ParseURL(w.connectionString)
+		options, err := redis.ParseURL(w.connectionString)
 		if err != nil {
 			w.logger.Fatal(err)
 		}
+		w.rdb = redis.NewClient(options)
 	} else if w.addr != "" {
-		options = &redis.Options{
-			Addr:     w.addr,
-			Password: w.password,
-			DB:       w.db,
+		if w.cluster {
+			w.rdb = redis.NewClusterClient(&redis.ClusterOptions{
+				Addrs:    strings.Split(w.addr, ","),
+				Password: w.password,
+			})
+		} else {
+			options := &redis.Options{
+				Addr:     w.addr,
+				Password: w.password,
+				DB:       w.db,
+			}
+			w.rdb = redis.NewClient(options)
 		}
 	}
 
-	rdb := redis.NewClient(options)
-
-	_, err = rdb.Ping(context.Background()).Result()
+	_, err = w.rdb.Ping(context.Background()).Result()
 	if err != nil {
 		w.logger.Fatal(err)
 	}
 
-	w.rdb = rdb
-
 	ctx := context.Background()
-	w.pubsub = w.rdb.Subscribe(ctx, w.channelName)
+
+	switch v := w.rdb.(type) {
+	case *redis.Client:
+		w.pubsub = v.Subscribe(ctx, w.channelName)
+	case *redis.ClusterClient:
+		w.pubsub = v.Subscribe(ctx, w.channelName)
+	}
 
 	var ropts []redis.ChannelOption
 
@@ -235,7 +253,12 @@ func (w *Worker) Shutdown() error {
 
 	w.stopOnce.Do(func() {
 		w.pubsub.Close()
-		w.rdb.Close()
+		switch v := w.rdb.(type) {
+		case *redis.Client:
+			v.Close()
+		case *redis.ClusterClient:
+			v.Close()
+		}
 		close(w.stop)
 	})
 	return nil

redis_test.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"log"
 	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -77,20 +78,47 @@ func TestCustomFuncAndWait(t *testing.T) {
 			return nil
 		}),
 	)
-	q, err := queue.NewQueue(
+	q := queue.NewPool(
+		5,
 		queue.WithWorker(w),
-		queue.WithWorkerCount(10),
 	)
-	assert.NoError(t, err)
-	q.Start()
 	time.Sleep(100 * time.Millisecond)
 	assert.NoError(t, q.Queue(m))
 	assert.NoError(t, q.Queue(m))
 	assert.NoError(t, q.Queue(m))
 	assert.NoError(t, q.Queue(m))
 	time.Sleep(1000 * time.Millisecond)
-	q.Shutdown()
-	q.Wait()
+	q.Release()
+	// you will see the execute time > 1000ms
+}
+
+func TestRedisCluster(t *testing.T) {
+	m := &mockMessage{
+		Message: "foo",
+	}
+
+	hosts := []string{host + ":6379", host + ":6380"}
+
+	w := NewWorker(
+		WithAddr(strings.Join(hosts, ",")),
+		WithChannel("testCluster"),
+		WithCluster(true),
+		WithRunFunc(func(ctx context.Context, m queue.QueuedMessage) error {
+			time.Sleep(500 * time.Millisecond)
+			return nil
+		}),
+	)
+	q := queue.NewPool(
+		5,
+		queue.WithWorker(w),
+	)
+	time.Sleep(100 * time.Millisecond)
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	assert.NoError(t, q.Queue(m))
+	time.Sleep(1000 * time.Millisecond)
+	q.Release()
 	// you will see the execute time > 1000ms
 }