Separate storage logic from actual business logic (#649)

go.mod

@@ -35,6 +35,7 @@ require (
 	github.com/thoas/stats v0.0.0-20190407194641-965cb2de1678
 	github.com/tidwall/buntdb v1.2.0
 	github.com/tidwall/gjson v1.12.1 // indirect
+	go.uber.org/atomic v1.5.0
 	golang.org/x/crypto v0.0.0-20210616213533-5ff15b29337e
 	golang.org/x/net v0.0.0-20210428140749-89ef3d95e781
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c

go.sum

@@ -538,6 +538,7 @@ go.opencensus.io v0.22.5 h1:dntmOdLpSpHlVqbW5Eay97DelsZHe+55D+xC6i0dDS0=
 go.opencensus.io v0.22.5/go.mod h1:5pWMHQbX5EPX2/62yrJeAkowc+lfs/XD7Uxpq3pI6kk=
 go.uber.org/atomic v1.3.2/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
 go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
+go.uber.org/atomic v1.5.0 h1:OI5t8sDa1Or+q8AeE+yKeB/SDYioSHAgcVljj9JIETY=
 go.uber.org/atomic v1.5.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
 go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
@@ -580,6 +581,7 @@ golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKG
 golang.org/x/mod v0.1.0/go.mod h1:0QHyrYULN0/3qlju5TqG8bIK38QM8yzMo5ekMj3DlcY=
 golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/mod v0.3.0 h1:RM4zey1++hCTbCVQfnWeKs9/IEsaBLA8vTkd0WVtmH4=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo=
 golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
@@ -705,6 +707,7 @@ golang.org/x/tools v0.0.0-20200103221440-774c71fcf114/go.mod h1:TB2adYChydJhpapK
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
+golang.org/x/tools v0.1.0 h1:po9/4sTYwZU9lPhi1tOrb4hCv3qrhiQ77LZfGa2OjwY=
 golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
 golang.org/x/tools v0.1.5 h1:ouewzE6p+/VEB31YYnTbEJdi8pFqKp4P4n85vwo3DHA=
 golang.org/x/tools v0.1.5/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=

status/status.go

@@ -20,7 +20,7 @@ import (
 var Stats *stats.Stats
 
 // StatStorage implements the storage interface
-var StatStorage storage.Storage
+var StatStorage *StateStorage
 
 // App is status structure
 type App struct {
@@ -54,24 +54,28 @@ type HuaweiStatus struct {
 // InitAppStatus for initialize app status
 func InitAppStatus(conf *config.ConfYaml) error {
 	logx.LogAccess.Info("Init App Status Engine as ", conf.Stat.Engine)
+
+	var store storage.Storage
 	switch conf.Stat.Engine {
 	case "memory":
-		StatStorage = memory.New()
+		store = memory.New()
 	case "redis":
-		StatStorage = redis.New(conf)
+		store = redis.New(conf)
 	case "boltdb":
-		StatStorage = boltdb.New(conf)
+		store = boltdb.New(conf)
 	case "buntdb":
-		StatStorage = buntdb.New(conf)
+		store = buntdb.New(conf)
 	case "leveldb":
-		StatStorage = leveldb.New(conf)
+		store = leveldb.New(conf)
 	case "badger":
-		StatStorage = badger.New(conf)
+		store = badger.New(conf)
 	default:
 		logx.LogError.Error("storage error: can't find storage driver")
 		return errors.New("can't find storage driver")
 	}
 
+	StatStorage = NewStateStorage(store)
+
 	if err := StatStorage.Init(); err != nil {
 		logx.LogError.Error("storage error: " + err.Error())
 

status/storage.go

@@ -0,0 +1,102 @@
+package status
+
+import "github.com/appleboy/gorush/storage"
+
+type StateStorage struct {
+	store storage.Storage
+}
+
+func NewStateStorage(store storage.Storage) *StateStorage {
+	return &StateStorage{
+		store: store,
+	}
+}
+
+func (s *StateStorage) Init() error {
+	return s.store.Init()
+}
+
+func (s *StateStorage) Close() error {
+	return s.store.Close()
+}
+
+// Reset Client storage.
+func (s *StateStorage) Reset() {
+	s.store.Set(storage.TotalCountKey, 0)
+	s.store.Set(storage.IosSuccessKey, 0)
+	s.store.Set(storage.IosErrorKey, 0)
+	s.store.Set(storage.AndroidSuccessKey, 0)
+	s.store.Set(storage.AndroidErrorKey, 0)
+	s.store.Set(storage.HuaweiSuccessKey, 0)
+	s.store.Set(storage.HuaweiErrorKey, 0)
+}
+
+// AddTotalCount record push notification count.
+func (s *StateStorage) AddTotalCount(count int64) {
+	s.store.Add(storage.TotalCountKey, count)
+}
+
+// AddIosSuccess record counts of success iOS push notification.
+func (s *StateStorage) AddIosSuccess(count int64) {
+	s.store.Add(storage.IosSuccessKey, count)
+}
+
+// AddIosError record counts of error iOS push notification.
+func (s *StateStorage) AddIosError(count int64) {
+	s.store.Add(storage.IosErrorKey, count)
+}
+
+// AddAndroidSuccess record counts of success Android push notification.
+func (s *StateStorage) AddAndroidSuccess(count int64) {
+	s.store.Add(storage.AndroidSuccessKey, count)
+}
+
+// AddAndroidError record counts of error Android push notification.
+func (s *StateStorage) AddAndroidError(count int64) {
+	s.store.Add(storage.AndroidErrorKey, count)
+}
+
+// AddHuaweiSuccess record counts of success Huawei push notification.
+func (s *StateStorage) AddHuaweiSuccess(count int64) {
+	s.store.Add(storage.HuaweiSuccessKey, count)
+}
+
+// AddHuaweiError record counts of error Huawei push notification.
+func (s *StateStorage) AddHuaweiError(count int64) {
+	s.store.Add(storage.HuaweiErrorKey, count)
+}
+
+// GetTotalCount show counts of all notification.
+func (s *StateStorage) GetTotalCount() int64 {
+	return s.store.Get(storage.TotalCountKey)
+}
+
+// GetIosSuccess show success counts of iOS notification.
+func (s *StateStorage) GetIosSuccess() int64 {
+	return s.store.Get(storage.IosSuccessKey)
+}
+
+// GetIosError show error counts of iOS notification.
+func (s *StateStorage) GetIosError() int64 {
+	return s.store.Get(storage.IosErrorKey)
+}
+
+// GetAndroidSuccess show success counts of Android notification.
+func (s *StateStorage) GetAndroidSuccess() int64 {
+	return s.store.Get(storage.AndroidSuccessKey)
+}
+
+// GetAndroidError show error counts of Android notification.
+func (s *StateStorage) GetAndroidError() int64 {
+	return s.store.Get(storage.AndroidErrorKey)
+}
+
+// GetHuaweiSuccess show success counts of Huawei notification.
+func (s *StateStorage) GetHuaweiSuccess() int64 {
+	return s.store.Get(storage.HuaweiSuccessKey)
+}
+
+// GetHuaweiError show error counts of Huawei notification.
+func (s *StateStorage) GetHuaweiError() int64 {
+	return s.store.Get(storage.HuaweiErrorKey)
+}

storage/badger/badger.go

@@ -4,10 +4,9 @@ import (
 	"log"
 	"os"
 	"strconv"
+	"sync"
 
 	"github.com/appleboy/gorush/config"
-	"github.com/appleboy/gorush/storage"
-
 	"github.com/dgraph-io/badger/v3"
 )
 
@@ -24,6 +23,26 @@ type Storage struct {
 	opts   badger.Options
 	name   string
 	db     *badger.DB
+
+	lock sync.RWMutex
+}
+
+func (s *Storage) Add(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setBadger(key, s.getBadger(key)+count)
+}
+
+func (s *Storage) Set(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setBadger(key, count)
+}
+
+func (s *Storage) Get(key string) int64 {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return s.getBadger(key)
 }
 
 // Init client storage.
@@ -50,17 +69,6 @@ func (s *Storage) Close() error {
 	return s.db.Close()
 }
 
-// Reset Client storage.
-func (s *Storage) Reset() {
-	s.setBadger(storage.TotalCountKey, 0)
-	s.setBadger(storage.IosSuccessKey, 0)
-	s.setBadger(storage.IosErrorKey, 0)
-	s.setBadger(storage.AndroidSuccessKey, 0)
-	s.setBadger(storage.AndroidErrorKey, 0)
-	s.setBadger(storage.HuaweiSuccessKey, 0)
-	s.setBadger(storage.HuaweiErrorKey, 0)
-}
-
 func (s *Storage) setBadger(key string, count int64) {
 	err := s.db.Update(func(txn *badger.Txn) error {
 		value := strconv.FormatInt(count, 10)
@@ -71,126 +79,28 @@ func (s *Storage) setBadger(key string, count int64) {
 	}
 }
 
-func (s *Storage) getBadger(key string, count *int64) {
+func (s *Storage) getBadger(key string) int64 {
+	var count int64
 	err := s.db.View(func(txn *badger.Txn) error {
 		item, err := txn.Get([]byte(key))
 		if err != nil {
 			return err
 		}
-		dst := []byte{}
+		var dst []byte
 		val, err := item.ValueCopy(dst)
 		if err != nil {
 			return err
 		}
 
-		i, err := strconv.ParseInt(string(val), 10, 64)
+		count, err = strconv.ParseInt(string(val), 10, 64)
 		if err != nil {
 			return err
 		}
 
-		*count = i
-
 		return nil
 	})
 	if err != nil {
 		log.Println(s.name, "get error:", err.Error())
 	}
-}
-
-// AddTotalCount record push notification count.
-func (s *Storage) AddTotalCount(count int64) {
-	total := s.GetTotalCount() + count
-	s.setBadger(storage.TotalCountKey, total)
-}
-
-// AddIosSuccess record counts of success iOS push notification.
-func (s *Storage) AddIosSuccess(count int64) {
-	total := s.GetIosSuccess() + count
-	s.setBadger(storage.IosSuccessKey, total)
-}
-
-// AddIosError record counts of error iOS push notification.
-func (s *Storage) AddIosError(count int64) {
-	total := s.GetIosError() + count
-	s.setBadger(storage.IosErrorKey, total)
-}
-
-// AddAndroidSuccess record counts of success Android push notification.
-func (s *Storage) AddAndroidSuccess(count int64) {
-	total := s.GetAndroidSuccess() + count
-	s.setBadger(storage.AndroidSuccessKey, total)
-}
-
-// AddAndroidError record counts of error Android push notification.
-func (s *Storage) AddAndroidError(count int64) {
-	total := s.GetAndroidError() + count
-	s.setBadger(storage.AndroidErrorKey, total)
-}
-
-// AddHuaweiSuccess record counts of success Huawei push notification.
-func (s *Storage) AddHuaweiSuccess(count int64) {
-	total := s.GetHuaweiSuccess() + count
-	s.setBadger(storage.HuaweiSuccessKey, total)
-}
-
-// AddHuaweiError record counts of error Huawei push notification.
-func (s *Storage) AddHuaweiError(count int64) {
-	total := s.GetHuaweiError() + count
-	s.setBadger(storage.HuaweiErrorKey, total)
-}
-
-// GetTotalCount show counts of all notification.
-func (s *Storage) GetTotalCount() int64 {
-	var count int64
-	s.getBadger(storage.TotalCountKey, &count)
-
-	return count
-}
-
-// GetIosSuccess show success counts of iOS notification.
-func (s *Storage) GetIosSuccess() int64 {
-	var count int64
-	s.getBadger(storage.IosSuccessKey, &count)
-
-	return count
-}
-
-// GetIosError show error counts of iOS notification.
-func (s *Storage) GetIosError() int64 {
-	var count int64
-	s.getBadger(storage.IosErrorKey, &count)
-
-	return count
-}
-
-// GetAndroidSuccess show success counts of Android notification.
-func (s *Storage) GetAndroidSuccess() int64 {
-	var count int64
-	s.getBadger(storage.AndroidSuccessKey, &count)
-
-	return count
-}
-
-// GetAndroidError show error counts of Android notification.
-func (s *Storage) GetAndroidError() int64 {
-	var count int64
-	s.getBadger(storage.AndroidErrorKey, &count)
-
-	return count
-}
-
-// GetHuaweiSuccess show success counts of Huawei notification.
-func (s *Storage) GetHuaweiSuccess() int64 {
-	var count int64
-	s.getBadger(storage.HuaweiSuccessKey, &count)
-
-	return count
-}
-
-// GetHuaweiError show error counts of Huawei notification.
-func (s *Storage) GetHuaweiError() int64 {
-	var count int64
-	s.getBadger(storage.HuaweiErrorKey, &count)
-
 	return count
 }

storage/badger/badger_test.go

@@ -1,8 +1,11 @@
 package badger
 
 import (
+	"sync"
 	"testing"
 
+	"github.com/appleboy/gorush/storage"
+
 	"github.com/appleboy/gorush/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -15,35 +18,30 @@ func TestBadgerEngine(t *testing.T) {
 	badger := New(cfg)
 	err := badger.Init()
 	assert.Nil(t, err)
-	badger.Reset()
 
-	badger.AddTotalCount(10)
+	badger.Add(storage.HuaweiSuccessKey, 10)
-	val = badger.GetTotalCount()
+	val = badger.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(10), val)
-	badger.AddTotalCount(10)
+	badger.Add(storage.HuaweiSuccessKey, 10)
-	val = badger.GetTotalCount()
+	val = badger.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(20), val)
 
-	badger.AddIosSuccess(20)
+	badger.Set(storage.HuaweiSuccessKey, 0)
-	val = badger.GetIosSuccess()
+	val = badger.Get(storage.HuaweiSuccessKey)
-	assert.Equal(t, int64(20), val)
-
-	badger.AddIosError(30)
-	val = badger.GetIosError()
-	assert.Equal(t, int64(30), val)
-
-	badger.AddAndroidSuccess(40)
-	val = badger.GetAndroidSuccess()
-	assert.Equal(t, int64(40), val)
-
-	badger.AddAndroidError(50)
-	val = badger.GetAndroidError()
-	assert.Equal(t, int64(50), val)
-
-	// test reset db
-	badger.Reset()
-	val = badger.GetAndroidError()
 	assert.Equal(t, int64(0), val)
 
+	// test concurrency issues
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			badger.Add(storage.HuaweiSuccessKey, 1)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	val = badger.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(10), val)
+
 	assert.NoError(t, badger.Close())
 }

storage/boltdb/boltdb.go

@@ -2,10 +2,9 @@ package boltdb
 
 import (
 	"log"
+	"sync"
 
 	"github.com/appleboy/gorush/config"
-	"github.com/appleboy/gorush/storage"
-
 	"github.com/asdine/storm/v3"
 )
 
@@ -20,6 +19,25 @@ func New(config *config.ConfYaml) *Storage {
 type Storage struct {
 	config *config.ConfYaml
 	db     *storm.DB
+	lock   sync.RWMutex
+}
+
+func (s *Storage) Add(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setBoltDB(key, s.getBoltDB(key)+count)
+}
+
+func (s *Storage) Set(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setBoltDB(key, count)
+}
+
+func (s *Storage) Get(key string) int64 {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return s.getBoltDB(key)
 }
 
 // Init client storage.
@@ -38,17 +56,6 @@ func (s *Storage) Close() error {
 	return s.db.Close()
 }
 
-// Reset Client storage.
-func (s *Storage) Reset() {
-	s.setBoltDB(storage.TotalCountKey, 0)
-	s.setBoltDB(storage.IosSuccessKey, 0)
-	s.setBoltDB(storage.IosErrorKey, 0)
-	s.setBoltDB(storage.AndroidSuccessKey, 0)
-	s.setBoltDB(storage.AndroidErrorKey, 0)
-	s.setBoltDB(storage.HuaweiSuccessKey, 0)
-	s.setBoltDB(storage.HuaweiErrorKey, 0)
-}
-
 func (s *Storage) setBoltDB(key string, count int64) {
 	err := s.db.Set(s.config.Stat.BoltDB.Bucket, key, count)
 	if err != nil {
@@ -56,107 +63,11 @@ func (s *Storage) setBoltDB(key string, count int64) {
 	}
 }
 
-func (s *Storage) getBoltDB(key string, count *int64) {
+func (s *Storage) getBoltDB(key string) int64 {
-	err := s.db.Get(s.config.Stat.BoltDB.Bucket, key, count)
+	var count int64
+	err := s.db.Get(s.config.Stat.BoltDB.Bucket, key, &count)
 	if err != nil {
 		log.Println("BoltDB get error:", err.Error())
 	}
-}
-
-// AddTotalCount record push notification count.
-func (s *Storage) AddTotalCount(count int64) {
-	total := s.GetTotalCount() + count
-	s.setBoltDB(storage.TotalCountKey, total)
-}
-
-// AddIosSuccess record counts of success iOS push notification.
-func (s *Storage) AddIosSuccess(count int64) {
-	total := s.GetIosSuccess() + count
-	s.setBoltDB(storage.IosSuccessKey, total)
-}
-
-// AddIosError record counts of error iOS push notification.
-func (s *Storage) AddIosError(count int64) {
-	total := s.GetIosError() + count
-	s.setBoltDB(storage.IosErrorKey, total)
-}
-
-// AddAndroidSuccess record counts of success Android push notification.
-func (s *Storage) AddAndroidSuccess(count int64) {
-	total := s.GetAndroidSuccess() + count
-	s.setBoltDB(storage.AndroidSuccessKey, total)
-}
-
-// AddAndroidError record counts of error Android push notification.
-func (s *Storage) AddAndroidError(count int64) {
-	total := s.GetAndroidError() + count
-	s.setBoltDB(storage.AndroidErrorKey, total)
-}
-
-// AddHuaweiSuccess record counts of success Huawei push notification.
-func (s *Storage) AddHuaweiSuccess(count int64) {
-	total := s.GetHuaweiSuccess() + count
-	s.setBoltDB(storage.HuaweiSuccessKey, total)
-}
-
-// AddHuaweiError record counts of error Huawei push notification.
-func (s *Storage) AddHuaweiError(count int64) {
-	total := s.GetHuaweiError() + count
-	s.setBoltDB(storage.HuaweiErrorKey, total)
-}
-
-// GetTotalCount show counts of all notification.
-func (s *Storage) GetTotalCount() int64 {
-	var count int64
-	s.getBoltDB(storage.TotalCountKey, &count)
-
-	return count
-}
-
-// GetIosSuccess show success counts of iOS notification.
-func (s *Storage) GetIosSuccess() int64 {
-	var count int64
-	s.getBoltDB(storage.IosSuccessKey, &count)
-
-	return count
-}
-
-// GetIosError show error counts of iOS notification.
-func (s *Storage) GetIosError() int64 {
-	var count int64
-	s.getBoltDB(storage.IosErrorKey, &count)
-
-	return count
-}
-
-// GetAndroidSuccess show success counts of Android notification.
-func (s *Storage) GetAndroidSuccess() int64 {
-	var count int64
-	s.getBoltDB(storage.AndroidSuccessKey, &count)
-
-	return count
-}
-
-// GetAndroidError show error counts of Android notification.
-func (s *Storage) GetAndroidError() int64 {
-	var count int64
-	s.getBoltDB(storage.AndroidErrorKey, &count)
-
-	return count
-}
-
-// GetHuaweiSuccess show success counts of Huawei notification.
-func (s *Storage) GetHuaweiSuccess() int64 {
-	var count int64
-	s.getBoltDB(storage.HuaweiSuccessKey, &count)
-
-	return count
-}
-
-// GetHuaweiError show error counts of Huawei notification.
-func (s *Storage) GetHuaweiError() int64 {
-	var count int64
-	s.getBoltDB(storage.HuaweiErrorKey, &count)
-
 	return count
 }

storage/boltdb/boltdb_test.go

@@ -1,8 +1,11 @@
 package boltdb
 
 import (
+	"sync"
 	"testing"
 
+	"github.com/appleboy/gorush/storage"
+
 	"github.com/appleboy/gorush/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -15,35 +18,30 @@ func TestBoltDBEngine(t *testing.T) {
 	boltDB := New(cfg)
 	err := boltDB.Init()
 	assert.Nil(t, err)
-	boltDB.Reset()
 
-	boltDB.AddTotalCount(10)
+	boltDB.Add(storage.HuaweiSuccessKey, 10)
-	val = boltDB.GetTotalCount()
+	val = boltDB.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(10), val)
-	boltDB.AddTotalCount(10)
+	boltDB.Add(storage.HuaweiSuccessKey, 10)
-	val = boltDB.GetTotalCount()
+	val = boltDB.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(20), val)
 
-	boltDB.AddIosSuccess(20)
+	boltDB.Set(storage.HuaweiSuccessKey, 0)
-	val = boltDB.GetIosSuccess()
+	val = boltDB.Get(storage.HuaweiSuccessKey)
-	assert.Equal(t, int64(20), val)
-
-	boltDB.AddIosError(30)
-	val = boltDB.GetIosError()
-	assert.Equal(t, int64(30), val)
-
-	boltDB.AddAndroidSuccess(40)
-	val = boltDB.GetAndroidSuccess()
-	assert.Equal(t, int64(40), val)
-
-	boltDB.AddAndroidError(50)
-	val = boltDB.GetAndroidError()
-	assert.Equal(t, int64(50), val)
-
-	// test reset db
-	boltDB.Reset()
-	val = boltDB.GetAndroidError()
 	assert.Equal(t, int64(0), val)
 
+	// test concurrency issues
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			boltDB.Add(storage.HuaweiSuccessKey, 1)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	val = boltDB.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(10), val)
+
 	assert.NoError(t, boltDB.Close())
 }

storage/buntdb/buntdb.go

@@ -4,10 +4,9 @@ import (
 	"fmt"
 	"log"
 	"strconv"
+	"sync"
 
 	"github.com/appleboy/gorush/config"
-	"github.com/appleboy/gorush/storage"
-
 	"github.com/tidwall/buntdb"
 )
 
@@ -22,6 +21,25 @@ func New(config *config.ConfYaml) *Storage {
 type Storage struct {
 	config *config.ConfYaml
 	db     *buntdb.DB
+	lock   sync.RWMutex
+}
+
+func (s *Storage) Add(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setBuntDB(key, s.getBuntDB(key)+count)
+}
+
+func (s *Storage) Set(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setBuntDB(key, count)
+}
+
+func (s *Storage) Get(key string) int64 {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return s.getBuntDB(key)
 }
 
 // Init client storage.
@@ -40,17 +58,6 @@ func (s *Storage) Close() error {
 	return s.db.Close()
 }
 
-// Reset Client storage.
-func (s *Storage) Reset() {
-	s.setBuntDB(storage.TotalCountKey, 0)
-	s.setBuntDB(storage.IosSuccessKey, 0)
-	s.setBuntDB(storage.IosErrorKey, 0)
-	s.setBuntDB(storage.AndroidSuccessKey, 0)
-	s.setBuntDB(storage.AndroidErrorKey, 0)
-	s.setBuntDB(storage.HuaweiSuccessKey, 0)
-	s.setBuntDB(storage.HuaweiErrorKey, 0)
-}
-
 func (s *Storage) setBuntDB(key string, count int64) {
 	err := s.db.Update(func(tx *buntdb.Tx) error {
 		if _, _, err := tx.Set(key, fmt.Sprintf("%d", count), nil); err != nil {
@@ -63,111 +70,16 @@ func (s *Storage) setBuntDB(key string, count int64) {
 	}
 }
 
-func (s *Storage) getBuntDB(key string, count *int64) {
+func (s *Storage) getBuntDB(key string) int64 {
+	var count int64
 	err := s.db.View(func(tx *buntdb.Tx) error {
 		val, _ := tx.Get(key)
-		*count, _ = strconv.ParseInt(val, 10, 64)
+		count, _ = strconv.ParseInt(val, 10, 64)
 		return nil
 	})
 	if err != nil {
 		log.Println("BuntDB get error:", err.Error())
 	}
-}
-
-// AddTotalCount record push notification count.
-func (s *Storage) AddTotalCount(count int64) {
-	total := s.GetTotalCount() + count
-	s.setBuntDB(storage.TotalCountKey, total)
-}
-
-// AddIosSuccess record counts of success iOS push notification.
-func (s *Storage) AddIosSuccess(count int64) {
-	total := s.GetIosSuccess() + count
-	s.setBuntDB(storage.IosSuccessKey, total)
-}
-
-// AddIosError record counts of error iOS push notification.
-func (s *Storage) AddIosError(count int64) {
-	total := s.GetIosError() + count
-	s.setBuntDB(storage.IosErrorKey, total)
-}
-
-// AddAndroidSuccess record counts of success Android push notification.
-func (s *Storage) AddAndroidSuccess(count int64) {
-	total := s.GetAndroidSuccess() + count
-	s.setBuntDB(storage.AndroidSuccessKey, total)
-}
-
-// AddAndroidError record counts of error Android push notification.
-func (s *Storage) AddAndroidError(count int64) {
-	total := s.GetAndroidError() + count
-	s.setBuntDB(storage.AndroidErrorKey, total)
-}
-
-// AddHuaweiSuccess record counts of success Huawei push notification.
-func (s *Storage) AddHuaweiSuccess(count int64) {
-	total := s.GetHuaweiSuccess() + count
-	s.setBuntDB(storage.HuaweiSuccessKey, total)
-}
-
-// AddHuaweiError record counts of error Huawei push notification.
-func (s *Storage) AddHuaweiError(count int64) {
-	total := s.GetHuaweiError() + count
-	s.setBuntDB(storage.HuaweiErrorKey, total)
-}
-
-// GetTotalCount show counts of all notification.
-func (s *Storage) GetTotalCount() int64 {
-	var count int64
-	s.getBuntDB(storage.TotalCountKey, &count)
-
-	return count
-}
-
-// GetIosSuccess show success counts of iOS notification.
-func (s *Storage) GetIosSuccess() int64 {
-	var count int64
-	s.getBuntDB(storage.IosSuccessKey, &count)
-
-	return count
-}
-
-// GetIosError show error counts of iOS notification.
-func (s *Storage) GetIosError() int64 {
-	var count int64
-	s.getBuntDB(storage.IosErrorKey, &count)
-
-	return count
-}
-
-// GetAndroidSuccess show success counts of Android notification.
-func (s *Storage) GetAndroidSuccess() int64 {
-	var count int64
-	s.getBuntDB(storage.AndroidSuccessKey, &count)
-
-	return count
-}
-
-// GetAndroidError show error counts of Android notification.
-func (s *Storage) GetAndroidError() int64 {
-	var count int64
-	s.getBuntDB(storage.AndroidErrorKey, &count)
-
-	return count
-}
-
-// GetHuaweiSuccess show success counts of Huawei notification.
-func (s *Storage) GetHuaweiSuccess() int64 {
-	var count int64
-	s.getBuntDB(storage.HuaweiSuccessKey, &count)
-
-	return count
-}
-
-// GetHuaweiError show error counts of Huawei notification.
-func (s *Storage) GetHuaweiError() int64 {
-	var count int64
-	s.getBuntDB(storage.HuaweiErrorKey, &count)
 
 	return count
 }

storage/buntdb/buntdb_test.go

@@ -2,8 +2,11 @@ package buntdb
 
 import (
 	"os"
+	"sync"
 	"testing"
 
+	"github.com/appleboy/gorush/storage"
+
 	"github.com/appleboy/gorush/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -21,34 +24,30 @@ func TestBuntDBEngine(t *testing.T) {
 	buntDB := New(cfg)
 	err := buntDB.Init()
 	assert.Nil(t, err)
-	buntDB.Reset()
 
-	buntDB.AddTotalCount(10)
+	buntDB.Add(storage.HuaweiSuccessKey, 10)
-	val = buntDB.GetTotalCount()
+	val = buntDB.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(10), val)
-	buntDB.AddTotalCount(10)
+	buntDB.Add(storage.HuaweiSuccessKey, 10)
-	val = buntDB.GetTotalCount()
+	val = buntDB.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(20), val)
 
-	buntDB.AddIosSuccess(20)
+	buntDB.Set(storage.HuaweiSuccessKey, 0)
-	val = buntDB.GetIosSuccess()
+	val = buntDB.Get(storage.HuaweiSuccessKey)
-	assert.Equal(t, int64(20), val)
-
-	buntDB.AddIosError(30)
-	val = buntDB.GetIosError()
-	assert.Equal(t, int64(30), val)
-
-	buntDB.AddAndroidSuccess(40)
-	val = buntDB.GetAndroidSuccess()
-	assert.Equal(t, int64(40), val)
-
-	buntDB.AddAndroidError(50)
-	val = buntDB.GetAndroidError()
-	assert.Equal(t, int64(50), val)
-
-	buntDB.Reset()
-	val = buntDB.GetAndroidError()
 	assert.Equal(t, int64(0), val)
 
+	// test concurrency issues
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			buntDB.Add(storage.HuaweiSuccessKey, 1)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	val = buntDB.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(10), val)
+
 	assert.NoError(t, buntDB.Close())
 }

storage/leveldb/leveldb.go

@@ -3,10 +3,9 @@ package leveldb
 import (
 	"fmt"
 	"strconv"
+	"sync"
 
 	"github.com/appleboy/gorush/config"
-	"github.com/appleboy/gorush/storage"
-
 	"github.com/syndtr/goleveldb/leveldb"
 )
 
@@ -15,9 +14,10 @@ func (s *Storage) setLevelDB(key string, count int64) {
 	_ = s.db.Put([]byte(key), []byte(value), nil)
 }
 
-func (s *Storage) getLevelDB(key string, count *int64) {
+func (s *Storage) getLevelDB(key string) int64 {
 	data, _ := s.db.Get([]byte(key), nil)
-	*count, _ = strconv.ParseInt(string(data), 10, 64)
+	count, _ := strconv.ParseInt(string(data), 10, 64)
+	return count
 }
 
 // New func implements the storage interface for gorush (https://github.com/appleboy/gorush)
@@ -31,6 +31,25 @@ func New(config *config.ConfYaml) *Storage {
 type Storage struct {
 	config *config.ConfYaml
 	db     *leveldb.DB
+	lock   sync.RWMutex
+}
+
+func (s *Storage) Add(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setLevelDB(key, s.getLevelDB(key)+count)
+}
+
+func (s *Storage) Set(key string, count int64) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.setLevelDB(key, count)
+}
+
+func (s *Storage) Get(key string) int64 {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return s.getLevelDB(key)
 }
 
 // Init client storage.
@@ -48,112 +67,3 @@ func (s *Storage) Close() error {
 
 	return s.db.Close()
 }
-
-// Reset Client storage.
-func (s *Storage) Reset() {
-	s.setLevelDB(storage.TotalCountKey, 0)
-	s.setLevelDB(storage.IosSuccessKey, 0)
-	s.setLevelDB(storage.IosErrorKey, 0)
-	s.setLevelDB(storage.AndroidSuccessKey, 0)
-	s.setLevelDB(storage.AndroidErrorKey, 0)
-	s.setLevelDB(storage.HuaweiSuccessKey, 0)
-	s.setLevelDB(storage.HuaweiErrorKey, 0)
-}
-
-// AddTotalCount record push notification count.
-func (s *Storage) AddTotalCount(count int64) {
-	total := s.GetTotalCount() + count
-	s.setLevelDB(storage.TotalCountKey, total)
-}
-
-// AddIosSuccess record counts of success iOS push notification.
-func (s *Storage) AddIosSuccess(count int64) {
-	total := s.GetIosSuccess() + count
-	s.setLevelDB(storage.IosSuccessKey, total)
-}
-
-// AddIosError record counts of error iOS push notification.
-func (s *Storage) AddIosError(count int64) {
-	total := s.GetIosError() + count
-	s.setLevelDB(storage.IosErrorKey, total)
-}
-
-// AddAndroidSuccess record counts of success Android push notification.
-func (s *Storage) AddAndroidSuccess(count int64) {
-	total := s.GetAndroidSuccess() + count
-	s.setLevelDB(storage.AndroidSuccessKey, total)
-}
-
-// AddAndroidError record counts of error Android push notification.
-func (s *Storage) AddAndroidError(count int64) {
-	total := s.GetAndroidError() + count
-	s.setLevelDB(storage.AndroidErrorKey, total)
-}
-
-// AddHuaweiSuccess record counts of success Huawei push notification.
-func (s *Storage) AddHuaweiSuccess(count int64) {
-	total := s.GetHuaweiSuccess() + count
-	s.setLevelDB(storage.HuaweiSuccessKey, total)
-}
-
-// AddHuaweiError record counts of error Huawei push notification.
-func (s *Storage) AddHuaweiError(count int64) {
-	total := s.GetHuaweiError() + count
-	s.setLevelDB(storage.HuaweiErrorKey, total)
-}
-
-// GetTotalCount show counts of all notification.
-func (s *Storage) GetTotalCount() int64 {
-	var count int64
-	s.getLevelDB(storage.TotalCountKey, &count)
-
-	return count
-}
-
-// GetIosSuccess show success counts of iOS notification.
-func (s *Storage) GetIosSuccess() int64 {
-	var count int64
-	s.getLevelDB(storage.IosSuccessKey, &count)
-
-	return count
-}
-
-// GetIosError show error counts of iOS notification.
-func (s *Storage) GetIosError() int64 {
-	var count int64
-	s.getLevelDB(storage.IosErrorKey, &count)
-
-	return count
-}
-
-// GetAndroidSuccess show success counts of Android notification.
-func (s *Storage) GetAndroidSuccess() int64 {
-	var count int64
-	s.getLevelDB(storage.AndroidSuccessKey, &count)
-
-	return count
-}
-
-// GetAndroidError show error counts of Android notification.
-func (s *Storage) GetAndroidError() int64 {
-	var count int64
-	s.getLevelDB(storage.AndroidErrorKey, &count)
-
-	return count
-}
-
-// GetHuaweiSuccess show success counts of Huawei notification.
-func (s *Storage) GetHuaweiSuccess() int64 {
-	var count int64
-	s.getLevelDB(storage.HuaweiSuccessKey, &count)
-
-	return count
-}
-
-// GetHuaweiError show error counts of Huawei notification.
-func (s *Storage) GetHuaweiError() int64 {
-	var count int64
-	s.getLevelDB(storage.HuaweiErrorKey, &count)
-
-	return count
-}

storage/leveldb/leveldb_test.go

@@ -2,8 +2,11 @@ package leveldb
 
 import (
 	"os"
+	"sync"
 	"testing"
 
+	"github.com/appleboy/gorush/storage"
+
 	"github.com/appleboy/gorush/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -21,34 +24,30 @@ func TestLevelDBEngine(t *testing.T) {
 	levelDB := New(cfg)
 	err := levelDB.Init()
 	assert.Nil(t, err)
-	levelDB.Reset()
 
-	levelDB.AddTotalCount(10)
+	levelDB.Add(storage.HuaweiSuccessKey, 10)
-	val = levelDB.GetTotalCount()
+	val = levelDB.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(10), val)
-	levelDB.AddTotalCount(10)
+	levelDB.Add(storage.HuaweiSuccessKey, 10)
-	val = levelDB.GetTotalCount()
+	val = levelDB.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(20), val)
 
-	levelDB.AddIosSuccess(20)
+	levelDB.Set(storage.HuaweiSuccessKey, 0)
-	val = levelDB.GetIosSuccess()
+	val = levelDB.Get(storage.HuaweiSuccessKey)
-	assert.Equal(t, int64(20), val)
-
-	levelDB.AddIosError(30)
-	val = levelDB.GetIosError()
-	assert.Equal(t, int64(30), val)
-
-	levelDB.AddAndroidSuccess(40)
-	val = levelDB.GetAndroidSuccess()
-	assert.Equal(t, int64(40), val)
-
-	levelDB.AddAndroidError(50)
-	val = levelDB.GetAndroidError()
-	assert.Equal(t, int64(50), val)
-
-	levelDB.Reset()
-	val = levelDB.GetAndroidError()
 	assert.Equal(t, int64(0), val)
 
+	// test concurrency issues
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			levelDB.Add(storage.HuaweiSuccessKey, 1)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	val = levelDB.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(10), val)
+
 	assert.NoError(t, levelDB.Close())
 }

storage/memory/memory.go

@@ -1,148 +1,48 @@
 package memory
 
 import (
-	"sync/atomic"
+	"sync"
+
+	"go.uber.org/atomic"
 )
 
-// statApp is app status structure
-type statApp struct {
-	TotalCount int64         `json:"total_count"`
-	Ios        IosStatus     `json:"ios"`
-	Android    AndroidStatus `json:"android"`
-	Huawei     HuaweiStatus  `json:"huawei"`
-}
-
-// AndroidStatus is android structure
-type AndroidStatus struct {
-	PushSuccess int64 `json:"push_success"`
-	PushError   int64 `json:"push_error"`
-}
-
-// IosStatus is iOS structure
-type IosStatus struct {
-	PushSuccess int64 `json:"push_success"`
-	PushError   int64 `json:"push_error"`
-}
-
-// HuaweiStatus is android structure
-type HuaweiStatus struct {
-	PushSuccess int64 `json:"push_success"`
-	PushError   int64 `json:"push_error"`
-}
-
 // New func implements the storage interface for gorush (https://github.com/appleboy/gorush)
 func New() *Storage {
-	return &Storage{
+	return &Storage{}
-		stat: &statApp{},
-	}
 }
 
 // Storage is interface structure
 type Storage struct {
-	stat *statApp
+	mem sync.Map
+}
+
+func (s *Storage) getValueBtKey(key string) *atomic.Int64 {
+	if val, ok := s.mem.Load(key); ok {
+		return val.(*atomic.Int64)
+	}
+	val := atomic.NewInt64(0)
+	s.mem.Store(key, val)
+	return val
+}
+
+func (s *Storage) Add(key string, count int64) {
+	s.getValueBtKey(key).Add(count)
+}
+
+func (s *Storage) Set(key string, count int64) {
+	s.getValueBtKey(key).Store(count)
+}
+
+func (s *Storage) Get(key string) int64 {
+	return s.getValueBtKey(key).Load()
 }
 
 // Init client storage.
-func (s *Storage) Init() error {
+func (*Storage) Init() error {
 	return nil
 }
 
 // Close the storage connection
-func (s *Storage) Close() error {
+func (*Storage) Close() error {
 	return nil
 }
-
-// Reset Client storage.
-func (s *Storage) Reset() {
-	atomic.StoreInt64(&s.stat.TotalCount, 0)
-	atomic.StoreInt64(&s.stat.Ios.PushSuccess, 0)
-	atomic.StoreInt64(&s.stat.Ios.PushError, 0)
-	atomic.StoreInt64(&s.stat.Android.PushSuccess, 0)
-	atomic.StoreInt64(&s.stat.Android.PushError, 0)
-	atomic.StoreInt64(&s.stat.Huawei.PushSuccess, 0)
-	atomic.StoreInt64(&s.stat.Huawei.PushError, 0)
-}
-
-// AddTotalCount record push notification count.
-func (s *Storage) AddTotalCount(count int64) {
-	atomic.AddInt64(&s.stat.TotalCount, count)
-}
-
-// AddIosSuccess record counts of success iOS push notification.
-func (s *Storage) AddIosSuccess(count int64) {
-	atomic.AddInt64(&s.stat.Ios.PushSuccess, count)
-}
-
-// AddIosError record counts of error iOS push notification.
-func (s *Storage) AddIosError(count int64) {
-	atomic.AddInt64(&s.stat.Ios.PushError, count)
-}
-
-// AddAndroidSuccess record counts of success Android push notification.
-func (s *Storage) AddAndroidSuccess(count int64) {
-	atomic.AddInt64(&s.stat.Android.PushSuccess, count)
-}
-
-// AddAndroidError record counts of error Android push notification.
-func (s *Storage) AddAndroidError(count int64) {
-	atomic.AddInt64(&s.stat.Android.PushError, count)
-}
-
-// AddHuaweiSuccess record counts of success Huawei push notification.
-func (s *Storage) AddHuaweiSuccess(count int64) {
-	atomic.AddInt64(&s.stat.Huawei.PushSuccess, count)
-}
-
-// AddHuaweiError record counts of error Huawei push notification.
-func (s *Storage) AddHuaweiError(count int64) {
-	atomic.AddInt64(&s.stat.Huawei.PushError, count)
-}
-
-// GetTotalCount show counts of all notification.
-func (s *Storage) GetTotalCount() int64 {
-	count := atomic.LoadInt64(&s.stat.TotalCount)
-
-	return count
-}
-
-// GetIosSuccess show success counts of iOS notification.
-func (s *Storage) GetIosSuccess() int64 {
-	count := atomic.LoadInt64(&s.stat.Ios.PushSuccess)
-
-	return count
-}
-
-// GetIosError show error counts of iOS notification.
-func (s *Storage) GetIosError() int64 {
-	count := atomic.LoadInt64(&s.stat.Ios.PushError)
-
-	return count
-}
-
-// GetAndroidSuccess show success counts of Android notification.
-func (s *Storage) GetAndroidSuccess() int64 {
-	count := atomic.LoadInt64(&s.stat.Android.PushSuccess)
-
-	return count
-}
-
-// GetAndroidError show error counts of Android notification.
-func (s *Storage) GetAndroidError() int64 {
-	count := atomic.LoadInt64(&s.stat.Android.PushError)
-
-	return count
-}
-
-// GetHuaweiSuccess show success counts of Huawei notification.
-func (s *Storage) GetHuaweiSuccess() int64 {
-	count := atomic.LoadInt64(&s.stat.Huawei.PushSuccess)
-
-	return count
-}
-
-// GetHuaweiError show error counts of Huawei notification.
-func (s *Storage) GetHuaweiError() int64 {
-	count := atomic.LoadInt64(&s.stat.Huawei.PushError)
-
-	return count
-}

storage/memory/memory_test.go

@@ -1,8 +1,11 @@
 package memory
 
 import (
+	"sync"
 	"testing"
 
+	"github.com/appleboy/gorush/storage"
+
 	"github.com/stretchr/testify/assert"
 )
 
@@ -10,37 +13,32 @@ func TestMemoryEngine(t *testing.T) {
 	var val int64
 
 	memory := New()
+	err := memory.Init()
+	assert.Nil(t, err)
 
-	assert.Nil(t, memory.Init())
+	memory.Add(storage.HuaweiSuccessKey, 10)
+	val = memory.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(10), val)
+	memory.Add(storage.HuaweiSuccessKey, 10)
+	val = memory.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(20), val)
 
-	memory.AddTotalCount(1)
+	memory.Set(storage.HuaweiSuccessKey, 0)
-	val = memory.GetTotalCount()
+	val = memory.Get(storage.HuaweiSuccessKey)
-	assert.Equal(t, int64(1), val)
-
-	memory.AddTotalCount(100)
-	val = memory.GetTotalCount()
-	assert.Equal(t, int64(101), val)
-
-	memory.AddIosSuccess(2)
-	val = memory.GetIosSuccess()
-	assert.Equal(t, int64(2), val)
-
-	memory.AddIosError(3)
-	val = memory.GetIosError()
-	assert.Equal(t, int64(3), val)
-
-	memory.AddAndroidSuccess(4)
-	val = memory.GetAndroidSuccess()
-	assert.Equal(t, int64(4), val)
-
-	memory.AddAndroidError(5)
-	val = memory.GetAndroidError()
-	assert.Equal(t, int64(5), val)
-
-	// test reset db
-	memory.Reset()
-	val = memory.GetTotalCount()
 	assert.Equal(t, int64(0), val)
 
+	// test concurrency issues
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			memory.Add(storage.HuaweiSuccessKey, 1)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	val = memory.Get(storage.HuaweiSuccessKey)
+	assert.Equal(t, int64(10), val)
+
 	assert.NoError(t, memory.Close())
 }

storage/redis/redis.go

@@ -8,8 +8,6 @@ import (
 	"strings"
 
 	"github.com/appleboy/gorush/config"
-	"github.com/appleboy/gorush/storage"
-
 	"github.com/go-redis/redis/v8"
 )
 
@@ -21,11 +19,6 @@ func New(config *config.ConfYaml) *Storage {
 	}
 }
 
-func (s *Storage) getInt64(key string, count *int64) {
-	val, _ := s.client.Get(s.ctx, key).Result()
-	*count, _ = strconv.ParseInt(val, 10, 64)
-}
-
 // Storage is interface structure
 type Storage struct {
 	ctx    context.Context
@@ -33,6 +26,20 @@ type Storage struct {
 	client redis.Cmdable
 }
 
+func (s *Storage) Add(key string, count int64) {
+	s.client.IncrBy(s.ctx, key, count)
+}
+
+func (s *Storage) Set(key string, count int64) {
+	s.client.Set(s.ctx, key, count, 0)
+}
+
+func (s *Storage) Get(key string) int64 {
+	val, _ := s.client.Get(s.ctx, key).Result()
+	count, _ := strconv.ParseInt(val, 10, 64)
+	return count
+}
+
 // Init client storage.
 func (s *Storage) Init() error {
 	if s.config.Stat.Redis.Cluster {
@@ -69,105 +76,3 @@ func (s *Storage) Close() error {
 		panic(fmt.Sprintf("invalid redis client: %v", reflect.TypeOf(v)))
 	}
 }
-
-// Reset Client storage.
-func (s *Storage) Reset() {
-	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(s.ctx, storage.TotalCountKey, count)
-}
-
-// AddIosSuccess record counts of success iOS push notification.
-func (s *Storage) AddIosSuccess(count int64) {
-	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(s.ctx, storage.IosErrorKey, count)
-}
-
-// AddAndroidSuccess record counts of success Android push notification.
-func (s *Storage) AddAndroidSuccess(count int64) {
-	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(s.ctx, storage.AndroidErrorKey, count)
-}
-
-// AddHuaweiSuccess record counts of success Android push notification.
-func (s *Storage) AddHuaweiSuccess(count int64) {
-	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(s.ctx, storage.HuaweiErrorKey, count)
-}
-
-// GetTotalCount show counts of all notification.
-func (s *Storage) GetTotalCount() int64 {
-	var count int64
-	s.getInt64(storage.TotalCountKey, &count)
-
-	return count
-}
-
-// GetIosSuccess show success counts of iOS notification.
-func (s *Storage) GetIosSuccess() int64 {
-	var count int64
-	s.getInt64(storage.IosSuccessKey, &count)
-
-	return count
-}
-
-// GetIosError show error counts of iOS notification.
-func (s *Storage) GetIosError() int64 {
-	var count int64
-	s.getInt64(storage.IosErrorKey, &count)
-
-	return count
-}
-
-// GetAndroidSuccess show success counts of Android notification.
-func (s *Storage) GetAndroidSuccess() int64 {
-	var count int64
-	s.getInt64(storage.AndroidSuccessKey, &count)
-
-	return count
-}
-
-// GetAndroidError show error counts of Android notification.
-func (s *Storage) GetAndroidError() int64 {
-	var count int64
-	s.getInt64(storage.AndroidErrorKey, &count)
-
-	return count
-}
-
-// GetHuaweiSuccess show success counts of Huawei notification.
-func (s *Storage) GetHuaweiSuccess() int64 {
-	var count int64
-	s.getInt64(storage.HuaweiSuccessKey, &count)
-
-	return count
-}
-
-// GetHuaweiError show error counts of Huawei notification.
-func (s *Storage) GetHuaweiError() int64 {
-	var count int64
-	s.getInt64(storage.HuaweiErrorKey, &count)
-
-	return count
-}

storage/redis/redis_test.go

@@ -4,6 +4,8 @@ import (
 	"sync"
 	"testing"
 
+	"github.com/appleboy/gorush/storage"
+
 	"github.com/appleboy/gorush/config"
 	"github.com/stretchr/testify/assert"
 )
@@ -27,34 +29,16 @@ func TestRedisEngine(t *testing.T) {
 	redis := New(cfg)
 	err := redis.Init()
 	assert.Nil(t, err)
-	redis.Reset()
 
-	redis.AddTotalCount(10)
+	redis.Add(storage.HuaweiSuccessKey, 10)
-	val = redis.GetTotalCount()
+	val = redis.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(10), val)
-	redis.AddTotalCount(10)
+	redis.Add(storage.HuaweiSuccessKey, 10)
-	val = redis.GetTotalCount()
+	val = redis.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(20), val)
 
-	redis.AddIosSuccess(20)
+	redis.Set(storage.HuaweiSuccessKey, 0)
-	val = redis.GetIosSuccess()
+	val = redis.Get(storage.HuaweiSuccessKey)
-	assert.Equal(t, int64(20), val)
-
-	redis.AddIosError(30)
-	val = redis.GetIosError()
-	assert.Equal(t, int64(30), val)
-
-	redis.AddAndroidSuccess(40)
-	val = redis.GetAndroidSuccess()
-	assert.Equal(t, int64(40), val)
-
-	redis.AddAndroidError(50)
-	val = redis.GetAndroidError()
-	assert.Equal(t, int64(50), val)
-
-	// test reset db
-	redis.Reset()
-	val = redis.GetAndroidError()
 	assert.Equal(t, int64(0), val)
 
 	// test concurrency issues
@@ -62,12 +46,12 @@ func TestRedisEngine(t *testing.T) {
 	for i := 0; i < 10; i++ {
 		wg.Add(1)
 		go func() {
-			redis.AddTotalCount(1)
+			redis.Add(storage.HuaweiSuccessKey, 1)
 			wg.Done()
 		}()
 	}
 	wg.Wait()
-	val = redis.GetTotalCount()
+	val = redis.Get(storage.HuaweiSuccessKey)
 	assert.Equal(t, int64(10), val)
 
 	assert.NoError(t, redis.Close())

storage/storage.go

@@ -26,20 +26,8 @@ const (
 // Storage interface
 type Storage interface {
 	Init() error
-	Reset()
+	Add(key string, count int64)
-	AddTotalCount(int64)
+	Set(key string, count int64)
-	AddIosSuccess(int64)
+	Get(key string) int64
-	AddIosError(int64)
-	AddAndroidSuccess(int64)
-	AddAndroidError(int64)
-	AddHuaweiSuccess(int64)
-	AddHuaweiError(int64)
-	GetTotalCount() int64
-	GetIosSuccess() int64
-	GetIosError() int64
-	GetAndroidSuccess() int64
-	GetAndroidError() int64
-	GetHuaweiSuccess() int64
-	GetHuaweiError() int64
 	Close() error
 }