Allow redis to support cluster mode (#633)

storage/redis/redis.go

@@ -1,97 +1,119 @@
 package redis
 
 import (
+	"context"
+	"fmt"
+	"reflect"
 	"strconv"
+	"strings"
 
 	"github.com/appleboy/gorush/config"
 	"github.com/appleboy/gorush/storage"
 
-	"github.com/go-redis/redis/v7"
+	"github.com/go-redis/redis/v8"
 )
 
 // New func implements the storage interface for gorush (https://github.com/appleboy/gorush)
 func New(config *config.ConfYaml) *Storage {
 	return &Storage{
+		ctx:    context.Background(),
 		config: config,
 	}
 }
 
 func (s *Storage) getInt64(key string, count *int64) {
-	val, _ := s.client.Get(key).Result()
+	val, _ := s.client.Get(s.ctx, key).Result()
 	*count, _ = strconv.ParseInt(val, 10, 64)
 }
 
 // Storage is interface structure
 type Storage struct {
+	ctx    context.Context
 	config *config.ConfYaml
-	client *redis.Client
+	client redis.Cmdable
 }
 
 // Init client storage.
 func (s *Storage) Init() error {
-	s.client = redis.NewClient(&redis.Options{
-		Addr:     s.config.Stat.Redis.Addr,
-		Password: s.config.Stat.Redis.Password,
-		DB:       s.config.Stat.Redis.DB,
-	})
-	_, err := s.client.Ping().Result()
+	if s.config.Stat.Redis.Cluster {
+		s.client = redis.NewClusterClient(&redis.ClusterOptions{
+			Addrs:    strings.Split(s.config.Stat.Redis.Addr, ","),
+			Password: s.config.Stat.Redis.Password,
+		})
+	} else {
+		s.client = redis.NewClient(&redis.Options{
+			Addr:     s.config.Stat.Redis.Addr,
+			Password: s.config.Stat.Redis.Password,
+			DB:       s.config.Stat.Redis.DB,
+		})
+	}
 
-	return err
+	if err := s.client.Ping(s.ctx).Err(); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // Close the storage connection
 func (s *Storage) Close() error {
-	if s.client == nil {
+	switch v := s.client.(type) {
+	case *redis.Client:
+		return v.Close()
+	case *redis.ClusterClient:
+		return v.Close()
+	case nil:
 		return nil
+	default:
+		// this will not happen anyway, unless we mishandle it on `Init`
+		panic(fmt.Sprintf("invalid redis client: %v", reflect.TypeOf(v)))
 	}
-
-	return s.client.Close()
 }
 
 // Reset Client storage.
 func (s *Storage) Reset() {
-	s.client.Set(storage.TotalCountKey, int64(0), 0)
-	s.client.Set(storage.IosSuccessKey, int64(0), 0)
-	s.client.Set(storage.IosErrorKey, int64(0), 0)
-	s.client.Set(storage.AndroidSuccessKey, int64(0), 0)
-	s.client.Set(storage.AndroidErrorKey, int64(0), 0)
-	s.client.Set(storage.HuaweiSuccessKey, int64(0), 0)
-	s.client.Set(storage.HuaweiErrorKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.TotalCountKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.IosSuccessKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.IosErrorKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.AndroidSuccessKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.AndroidErrorKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.HuaweiSuccessKey, int64(0), 0)
+	s.client.Set(s.ctx, storage.HuaweiErrorKey, int64(0), 0)
 }
 
 // AddTotalCount record push notification count.
 func (s *Storage) AddTotalCount(count int64) {
-	s.client.IncrBy(storage.TotalCountKey, count)
+	s.client.IncrBy(s.ctx, storage.TotalCountKey, count)
 }
 
 // AddIosSuccess record counts of success iOS push notification.
 func (s *Storage) AddIosSuccess(count int64) {
-	s.client.IncrBy(storage.IosSuccessKey, count)
+	s.client.IncrBy(s.ctx, storage.IosSuccessKey, count)
 }
 
 // AddIosError record counts of error iOS push notification.
 func (s *Storage) AddIosError(count int64) {
-	s.client.IncrBy(storage.IosErrorKey, count)
+	s.client.IncrBy(s.ctx, storage.IosErrorKey, count)
 }
 
 // AddAndroidSuccess record counts of success Android push notification.
 func (s *Storage) AddAndroidSuccess(count int64) {
-	s.client.IncrBy(storage.AndroidSuccessKey, count)
+	s.client.IncrBy(s.ctx, storage.AndroidSuccessKey, count)
 }
 
 // AddAndroidError record counts of error Android push notification.
 func (s *Storage) AddAndroidError(count int64) {
-	s.client.IncrBy(storage.AndroidErrorKey, count)
+	s.client.IncrBy(s.ctx, storage.AndroidErrorKey, count)
 }
 
 // AddHuaweiSuccess record counts of success Android push notification.
 func (s *Storage) AddHuaweiSuccess(count int64) {
-	s.client.IncrBy(storage.HuaweiSuccessKey, count)
+	s.client.IncrBy(s.ctx, storage.HuaweiSuccessKey, count)
 }
 
 // AddHuaweiError record counts of error Android push notification.
 func (s *Storage) AddHuaweiError(count int64) {
-	s.client.IncrBy(storage.HuaweiErrorKey, count)
+	s.client.IncrBy(s.ctx, storage.HuaweiErrorKey, count)
 }
 
 // GetTotalCount show counts of all notification.