diff --git a/go.mod b/go.mod index 29a3567..cabb15f 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 3cf0b49..ff9d69f 100644 --- a/go.sum +++ b/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= diff --git a/status/status.go b/status/status.go index 68b423d..e432948 100644 --- a/status/status.go +++ b/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()) diff --git a/status/storage.go b/status/storage.go new file mode 100644 index 0000000..d0dab45 --- /dev/null +++ b/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) +} diff --git a/storage/badger/badger.go b/storage/badger/badger.go index b0b02d1..1c40d90 100644 --- a/storage/badger/badger.go +++ b/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 } diff --git a/storage/badger/badger_test.go b/storage/badger/badger_test.go index 3c26b27..11d6009 100644 --- a/storage/badger/badger_test.go +++ b/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) - val = badger.GetTotalCount() + badger.Add(storage.HuaweiSuccessKey, 10) + val = badger.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(10), val) - badger.AddTotalCount(10) - val = badger.GetTotalCount() + badger.Add(storage.HuaweiSuccessKey, 10) + val = badger.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(20), val) - badger.AddIosSuccess(20) - val = badger.GetIosSuccess() - 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() + badger.Set(storage.HuaweiSuccessKey, 0) + val = badger.Get(storage.HuaweiSuccessKey) 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()) } diff --git a/storage/boltdb/boltdb.go b/storage/boltdb/boltdb.go index 1ae26a9..5d0cce1 100644 --- a/storage/boltdb/boltdb.go +++ b/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) { - err := s.db.Get(s.config.Stat.BoltDB.Bucket, key, count) +func (s *Storage) getBoltDB(key string) int64 { + 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 } diff --git a/storage/boltdb/boltdb_test.go b/storage/boltdb/boltdb_test.go index 07e1687..3115407 100644 --- a/storage/boltdb/boltdb_test.go +++ b/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) - val = boltDB.GetTotalCount() + boltDB.Add(storage.HuaweiSuccessKey, 10) + val = boltDB.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(10), val) - boltDB.AddTotalCount(10) - val = boltDB.GetTotalCount() + boltDB.Add(storage.HuaweiSuccessKey, 10) + val = boltDB.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(20), val) - boltDB.AddIosSuccess(20) - val = boltDB.GetIosSuccess() - 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() + boltDB.Set(storage.HuaweiSuccessKey, 0) + val = boltDB.Get(storage.HuaweiSuccessKey) 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()) } diff --git a/storage/buntdb/buntdb.go b/storage/buntdb/buntdb.go index 5bcd30d..6c3411c 100644 --- a/storage/buntdb/buntdb.go +++ b/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 } diff --git a/storage/buntdb/buntdb_test.go b/storage/buntdb/buntdb_test.go index 022ce42..35bd833 100644 --- a/storage/buntdb/buntdb_test.go +++ b/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) - val = buntDB.GetTotalCount() + buntDB.Add(storage.HuaweiSuccessKey, 10) + val = buntDB.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(10), val) - buntDB.AddTotalCount(10) - val = buntDB.GetTotalCount() + buntDB.Add(storage.HuaweiSuccessKey, 10) + val = buntDB.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(20), val) - buntDB.AddIosSuccess(20) - val = buntDB.GetIosSuccess() - 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() + buntDB.Set(storage.HuaweiSuccessKey, 0) + val = buntDB.Get(storage.HuaweiSuccessKey) 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()) } diff --git a/storage/leveldb/leveldb.go b/storage/leveldb/leveldb.go index 8164fc0..f9e285d 100644 --- a/storage/leveldb/leveldb.go +++ b/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 -} diff --git a/storage/leveldb/leveldb_test.go b/storage/leveldb/leveldb_test.go index 8268214..9cd7fa2 100644 --- a/storage/leveldb/leveldb_test.go +++ b/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) - val = levelDB.GetTotalCount() + levelDB.Add(storage.HuaweiSuccessKey, 10) + val = levelDB.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(10), val) - levelDB.AddTotalCount(10) - val = levelDB.GetTotalCount() + levelDB.Add(storage.HuaweiSuccessKey, 10) + val = levelDB.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(20), val) - levelDB.AddIosSuccess(20) - val = levelDB.GetIosSuccess() - 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() + levelDB.Set(storage.HuaweiSuccessKey, 0) + val = levelDB.Get(storage.HuaweiSuccessKey) 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()) } diff --git a/storage/memory/memory.go b/storage/memory/memory.go index 4b3c5ff..fa46b2d 100644 --- a/storage/memory/memory.go +++ b/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{ - stat: &statApp{}, - } + return &Storage{} } // 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 -} diff --git a/storage/memory/memory_test.go b/storage/memory/memory_test.go index 2e088c1..3e92f81 100644 --- a/storage/memory/memory_test.go +++ b/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) - val = memory.GetTotalCount() - 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() + memory.Set(storage.HuaweiSuccessKey, 0) + val = memory.Get(storage.HuaweiSuccessKey) 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()) } diff --git a/storage/redis/redis.go b/storage/redis/redis.go index 45001a1..bd76be5 100644 --- a/storage/redis/redis.go +++ b/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 -} diff --git a/storage/redis/redis_test.go b/storage/redis/redis_test.go index 0f31072..94f440b 100644 --- a/storage/redis/redis_test.go +++ b/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) - val = redis.GetTotalCount() + redis.Add(storage.HuaweiSuccessKey, 10) + val = redis.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(10), val) - redis.AddTotalCount(10) - val = redis.GetTotalCount() + redis.Add(storage.HuaweiSuccessKey, 10) + val = redis.Get(storage.HuaweiSuccessKey) assert.Equal(t, int64(20), val) - redis.AddIosSuccess(20) - val = redis.GetIosSuccess() - 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() + redis.Set(storage.HuaweiSuccessKey, 0) + val = redis.Get(storage.HuaweiSuccessKey) 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()) diff --git a/storage/storage.go b/storage/storage.go index c89c514..6eb9fa5 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -26,20 +26,8 @@ const ( // Storage interface type Storage interface { Init() error - Reset() - AddTotalCount(int64) - AddIosSuccess(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 + Add(key string, count int64) + Set(key string, count int64) + Get(key string) int64 Close() error }