chore(storage): storage performance issues (#500)

* chore(storage): storage performance issues

1. close storage connection before shutdown the service
2. update windows image

https://github.com/appleboy/gorush/issues/393

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>

.appveyor.yml

@@ -1,5 +1,5 @@
 version: '{build}'
-image: 'Visual Studio 2017'
+image: 'Visual Studio 2019'
 platform: x64
 
 max_jobs: 1
@@ -17,6 +17,7 @@ branches:
 
 install:
   - ps: |
+      choco install -y mingw
       docker version
       go version
 

docker/Dockerfile.windows.amd64

@@ -1,4 +1,4 @@
-FROM microsoft/nanoserver:10.0.14393.1884
+FROM mcr.microsoft.com/windows/nanoserver:1809-amd64
 
 LABEL maintainer="Bo-Yi Wu <appleboy.tw@gmail.com>" \
   org.label-schema.name="Gorush" \

main.go

@@ -255,8 +255,13 @@ func main() {
 		gorush.LogAccess.Info("close the notification queue channel, current queue len: ", len(gorush.QueueNotification))
 		close(gorush.QueueNotification)
 		wg.Wait()
-		close(finished)
 		gorush.LogAccess.Info("the notification queue has been clear")
+		close(finished)
+		// close the connection with storage
+		gorush.LogAccess.Info("close the storage connection: ", gorush.PushConf.Stat.Engine)
+		if err := gorush.StatStorage.Close(); err != nil {
+			gorush.LogError.Fatal("can't close the storage connection: ", err.Error())
+		}
 	})
 
 	gorush.InitWorkers(ctx, wg, gorush.PushConf.Core.WorkerNum, gorush.PushConf.Core.QueueNum)

storage/badger/badger.go

@@ -25,10 +25,12 @@ type Storage struct {
 	config config.ConfYaml
 	opts   badger.Options
 	name   string
+	db     *badger.DB
 }
 
 // Init client storage.
 func (s *Storage) Init() error {
+	var err error
 	s.name = "badger"
 	dbPath := s.config.Stat.BadgerDB.Path
 	if dbPath == "" {
@@ -36,9 +38,20 @@ func (s *Storage) Init() error {
 	}
 	s.opts = badger.DefaultOptions(dbPath)
 
+	s.db, err = badger.Open(s.opts)
+
+	return err
+}
+
+// Close the storage connection
+func (s *Storage) Close() error {
+	if s.db == nil {
 		return nil
 	}
 
+	return s.db.Close()
+}
+
 // Reset Client storage.
 func (s *Storage) Reset() {
 	s.setBadger(storage.TotalCountKey, 0)
@@ -49,21 +62,7 @@ func (s *Storage) Reset() {
 }
 
 func (s *Storage) setBadger(key string, count int64) {
-	db, err := badger.Open(s.opts)
+	err := s.db.Update(func(txn *badger.Txn) error {
-
-	if err != nil {
-		log.Println(s.name, "open error:", err.Error())
-		return
-	}
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println(s.name, "close error:", err.Error())
-		}
-	}()
-
-	err = db.Update(func(txn *badger.Txn) error {
 		value := convert.ToString(count).(string)
 		return txn.Set([]byte(key), []byte(value))
 	})
@@ -74,21 +73,7 @@ func (s *Storage) setBadger(key string, count int64) {
 }
 
 func (s *Storage) getBadger(key string, count *int64) {
-	db, err := badger.Open(s.opts)
+	err := s.db.View(func(txn *badger.Txn) error {
-
-	if err != nil {
-		log.Println(s.name, "open error:", err.Error())
-		return
-	}
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println(s.name, "close error:", err.Error())
-		}
-	}()
-
-	err = db.View(func(txn *badger.Txn) error {
 		item, err := txn.Get([]byte(key))
 		if err != nil {
 			return err

storage/badger/badger_test.go

@@ -44,4 +44,6 @@ func TestBadgerEngine(t *testing.T) {
 	badger.Reset()
 	val = badger.GetAndroidError()
 	assert.Equal(t, int64(0), val)
+
+	assert.NoError(t, badger.Close())
 }

storage/boltdb/boltdb.go

@@ -19,13 +19,25 @@ func New(config config.ConfYaml) *Storage {
 // Storage is interface structure
 type Storage struct {
 	config config.ConfYaml
+	db     *storm.DB
 }
 
 // Init client storage.
 func (s *Storage) Init() error {
+	var err error
+	s.db, err = storm.Open(s.config.Stat.BoltDB.Path)
+	return err
+}
+
+// Close the storage connection
+func (s *Storage) Close() error {
+	if s.db == nil {
 		return nil
 	}
 
+	return s.db.Close()
+}
+
 // Reset Client storage.
 func (s *Storage) Reset() {
 	s.setBoltDB(storage.TotalCountKey, 0)
@@ -36,33 +48,17 @@ func (s *Storage) Reset() {
 }
 
 func (s *Storage) setBoltDB(key string, count int64) {
-	db, _ := storm.Open(s.config.Stat.BoltDB.Path)
+	err := s.db.Set(s.config.Stat.BoltDB.Bucket, key, count)
-	err := db.Set(s.config.Stat.BoltDB.Bucket, key, count)
 	if err != nil {
 		log.Println("BoltDB set error:", err.Error())
 	}
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println("BoltDB error:", err.Error())
-		}
-	}()
 }
 
 func (s *Storage) getBoltDB(key string, count *int64) {
-	db, _ := storm.Open(s.config.Stat.BoltDB.Path)
+	err := s.db.Get(s.config.Stat.BoltDB.Bucket, key, count)
-	err := db.Get(s.config.Stat.BoltDB.Bucket, key, count)
 	if err != nil {
 		log.Println("BoltDB get error:", err.Error())
 	}
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println("BoltDB error:", err.Error())
-		}
-	}()
 }
 
 // AddTotalCount record push notification count.

storage/boltdb/boltdb_test.go

@@ -44,4 +44,6 @@ func TestBoltDBEngine(t *testing.T) {
 	boltDB.Reset()
 	val = boltDB.GetAndroidError()
 	assert.Equal(t, int64(0), val)
+
+	assert.NoError(t, boltDB.Close())
 }

storage/buntdb/buntdb.go

@@ -21,13 +21,25 @@ func New(config config.ConfYaml) *Storage {
 // Storage is interface structure
 type Storage struct {
 	config config.ConfYaml
+	db     *buntdb.DB
 }
 
 // Init client storage.
 func (s *Storage) Init() error {
+	var err error
+	s.db, err = buntdb.Open(s.config.Stat.BuntDB.Path)
+	return err
+}
+
+// Close the storage connection
+func (s *Storage) Close() error {
+	if s.db == nil {
 		return nil
 	}
 
+	return s.db.Close()
+}
+
 // Reset Client storage.
 func (s *Storage) Reset() {
 	s.setBuntDB(storage.TotalCountKey, 0)
@@ -38,9 +50,7 @@ func (s *Storage) Reset() {
 }
 
 func (s *Storage) setBuntDB(key string, count int64) {
-	db, _ := buntdb.Open(s.config.Stat.BuntDB.Path)
+	err := s.db.Update(func(tx *buntdb.Tx) error {
-
-	err := db.Update(func(tx *buntdb.Tx) error {
 		if _, _, err := tx.Set(key, fmt.Sprintf("%d", count), nil); err != nil {
 			return err
 		}
@@ -50,19 +60,10 @@ func (s *Storage) setBuntDB(key string, count int64) {
 	if err != nil {
 		log.Println("BuntDB update error:", err.Error())
 	}
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println("BuntDB error:", err.Error())
-		}
-	}()
 }
 
 func (s *Storage) getBuntDB(key string, count *int64) {
-	db, _ := buntdb.Open(s.config.Stat.BuntDB.Path)
+	err := s.db.View(func(tx *buntdb.Tx) error {
-
-	err := db.View(func(tx *buntdb.Tx) error {
 		val, _ := tx.Get(key)
 		*count, _ = strconv.ParseInt(val, 10, 64)
 		return nil
@@ -71,13 +72,6 @@ func (s *Storage) getBuntDB(key string, count *int64) {
 	if err != nil {
 		log.Println("BuntDB get error:", err.Error())
 	}
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println("BuntDB error:", err.Error())
-		}
-	}()
 }
 
 // AddTotalCount record push notification count.

storage/buntdb/buntdb_test.go

@@ -49,4 +49,6 @@ func TestBuntDBEngine(t *testing.T) {
 	buntDB.Reset()
 	val = buntDB.GetAndroidError()
 	assert.Equal(t, int64(0), val)
+
+	assert.NoError(t, buntDB.Close())
 }

storage/leveldb/leveldb.go

@@ -2,7 +2,6 @@ package leveldb
 
 import (
 	"fmt"
-	"log"
 	"strconv"
 
 	"github.com/appleboy/gorush/config"
@@ -11,34 +10,14 @@ import (
 	"github.com/syndtr/goleveldb/leveldb"
 )
 
-var dbPath string
+func (s *Storage) setLevelDB(key string, count int64) {
-
-func setLevelDB(key string, count int64) {
-	db, _ := leveldb.OpenFile(dbPath, nil)
 	value := fmt.Sprintf("%d", count)
-
+	_ = s.db.Put([]byte(key), []byte(value), nil)
-	_ = db.Put([]byte(key), []byte(value), nil)
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println("LevelDB error:", err.Error())
-		}
-	}()
 }
 
-func getLevelDB(key string, count *int64) {
+func (s *Storage) getLevelDB(key string, count *int64) {
-	db, _ := leveldb.OpenFile(dbPath, nil)
+	data, _ := s.db.Get([]byte(key), nil)
-
-	data, _ := db.Get([]byte(key), nil)
 	*count, _ = strconv.ParseInt(string(data), 10, 64)
-
-	defer func() {
-		err := db.Close()
-		if err != nil {
-			log.Println("LevelDB error:", err.Error())
-		}
-	}()
 }
 
 // New func implements the storage interface for gorush (https://github.com/appleboy/gorush)
@@ -51,57 +30,68 @@ func New(config config.ConfYaml) *Storage {
 // Storage is interface structure
 type Storage struct {
 	config config.ConfYaml
+	db     *leveldb.DB
 }
 
 // Init client storage.
 func (s *Storage) Init() error {
-	dbPath = s.config.Stat.LevelDB.Path
+	var err error
+	s.db, err = leveldb.OpenFile(s.config.Stat.LevelDB.Path, nil)
+	return err
+}
+
+// Close the storage connection
+func (s *Storage) Close() error {
+	if s.db == nil {
 		return nil
 	}
 
+	return s.db.Close()
+}
+
 // Reset Client storage.
 func (s *Storage) Reset() {
-	setLevelDB(storage.TotalCountKey, 0)
+	s.setLevelDB(storage.TotalCountKey, 0)
-	setLevelDB(storage.IosSuccessKey, 0)
+	s.setLevelDB(storage.IosSuccessKey, 0)
-	setLevelDB(storage.IosErrorKey, 0)
+	s.setLevelDB(storage.IosErrorKey, 0)
-	setLevelDB(storage.AndroidSuccessKey, 0)
+	s.setLevelDB(storage.AndroidSuccessKey, 0)
-	setLevelDB(storage.AndroidErrorKey, 0)
+	s.setLevelDB(storage.AndroidErrorKey, 0)
 }
 
 // AddTotalCount record push notification count.
 func (s *Storage) AddTotalCount(count int64) {
 	total := s.GetTotalCount() + count
-	setLevelDB(storage.TotalCountKey, total)
+	s.setLevelDB(storage.TotalCountKey, total)
 }
 
 // AddIosSuccess record counts of success iOS push notification.
 func (s *Storage) AddIosSuccess(count int64) {
 	total := s.GetIosSuccess() + count
-	setLevelDB(storage.IosSuccessKey, total)
+	s.setLevelDB(storage.IosSuccessKey, total)
 }
 
 // AddIosError record counts of error iOS push notification.
 func (s *Storage) AddIosError(count int64) {
 	total := s.GetIosError() + count
-	setLevelDB(storage.IosErrorKey, total)
+	s.setLevelDB(storage.IosErrorKey, total)
 }
 
 // AddAndroidSuccess record counts of success Android push notification.
 func (s *Storage) AddAndroidSuccess(count int64) {
 	total := s.GetAndroidSuccess() + count
-	setLevelDB(storage.AndroidSuccessKey, total)
+	s.setLevelDB(storage.AndroidSuccessKey, total)
 }
 
 // AddAndroidError record counts of error Android push notification.
 func (s *Storage) AddAndroidError(count int64) {
 	total := s.GetAndroidError() + count
-	setLevelDB(storage.AndroidErrorKey, total)
+	s.setLevelDB(storage.AndroidErrorKey, total)
 }
 
 // GetTotalCount show counts of all notification.
 func (s *Storage) GetTotalCount() int64 {
 	var count int64
-	getLevelDB(storage.TotalCountKey, &count)
+	s.getLevelDB(storage.TotalCountKey, &count)
 
 	return count
 }
@@ -109,7 +99,7 @@ func (s *Storage) GetTotalCount() int64 {
 // GetIosSuccess show success counts of iOS notification.
 func (s *Storage) GetIosSuccess() int64 {
 	var count int64
-	getLevelDB(storage.IosSuccessKey, &count)
+	s.getLevelDB(storage.IosSuccessKey, &count)
 
 	return count
 }
@@ -117,7 +107,7 @@ func (s *Storage) GetIosSuccess() int64 {
 // GetIosError show error counts of iOS notification.
 func (s *Storage) GetIosError() int64 {
 	var count int64
-	getLevelDB(storage.IosErrorKey, &count)
+	s.getLevelDB(storage.IosErrorKey, &count)
 
 	return count
 }
@@ -125,7 +115,7 @@ func (s *Storage) GetIosError() int64 {
 // GetAndroidSuccess show success counts of Android notification.
 func (s *Storage) GetAndroidSuccess() int64 {
 	var count int64
-	getLevelDB(storage.AndroidSuccessKey, &count)
+	s.getLevelDB(storage.AndroidSuccessKey, &count)
 
 	return count
 }
@@ -133,7 +123,7 @@ func (s *Storage) GetAndroidSuccess() int64 {
 // GetAndroidError show error counts of Android notification.
 func (s *Storage) GetAndroidError() int64 {
 	var count int64
-	getLevelDB(storage.AndroidErrorKey, &count)
+	s.getLevelDB(storage.AndroidErrorKey, &count)
 
 	return count
 }

storage/leveldb/leveldb_test.go

@@ -49,4 +49,6 @@ func TestLevelDBEngine(t *testing.T) {
 	levelDB.Reset()
 	val = levelDB.GetAndroidError()
 	assert.Equal(t, int64(0), val)
+
+	assert.NoError(t, levelDB.Close())
 }

storage/memory/memory.go

@@ -40,6 +40,11 @@ func (s *Storage) Init() error {
 	return nil
 }
 
+// Close the storage connection
+func (s *Storage) Close() error {
+	return nil
+}
+
 // Reset Client storage.
 func (s *Storage) Reset() {
 	atomic.StoreInt64(&s.stat.TotalCount, 0)

storage/memory/memory_test.go

@@ -41,4 +41,6 @@ func TestMemoryEngine(t *testing.T) {
 	memory.Reset()
 	val = memory.GetTotalCount()
 	assert.Equal(t, int64(0), val)
+
+	assert.NoError(t, memory.Close())
 }

storage/redis/redis.go

@@ -1,7 +1,6 @@
 package redis
 
 import (
-	"log"
 	"strconv"
 
 	"github.com/appleboy/gorush/config"
@@ -10,9 +9,6 @@ import (
 	"gopkg.in/redis.v5"
 )
 
-//
-var redisClient *redis.Client
-
 // New func implements the storage interface for gorush (https://github.com/appleboy/gorush)
 func New(config config.ConfYaml) *Storage {
 	return &Storage{
@@ -20,76 +16,76 @@ func New(config config.ConfYaml) *Storage {
 	}
 }
 
-func getInt64(key string, count *int64) {
+func (s *Storage) getInt64(key string, count *int64) {
-	val, _ := redisClient.Get(key).Result()
+	val, _ := s.client.Get(key).Result()
 	*count, _ = strconv.ParseInt(val, 10, 64)
 }
 
 // Storage is interface structure
 type Storage struct {
 	config config.ConfYaml
+	client *redis.Client
 }
 
 // Init client storage.
 func (s *Storage) Init() error {
-	redisClient = redis.NewClient(&redis.Options{
+	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()
-	_, err := redisClient.Ping().Result()
-
-	if err != nil {
-		// redis server error
-		log.Println("Can't connect redis server: " + err.Error())
 
 	return err
 }
-	// Set initial values
-	s.Reset()
 
+// Close the storage connection
+func (s *Storage) Close() error {
+	if s.client == nil {
 		return nil
 	}
 
+	return s.client.Close()
+}
+
 // Reset Client storage.
 func (s *Storage) Reset() {
-	redisClient.Set(storage.TotalCountKey, int64(0), 0)
+	s.client.Set(storage.TotalCountKey, int64(0), 0)
-	redisClient.Set(storage.IosSuccessKey, int64(0), 0)
+	s.client.Set(storage.IosSuccessKey, int64(0), 0)
-	redisClient.Set(storage.IosErrorKey, int64(0), 0)
+	s.client.Set(storage.IosErrorKey, int64(0), 0)
-	redisClient.Set(storage.AndroidSuccessKey, int64(0), 0)
+	s.client.Set(storage.AndroidSuccessKey, int64(0), 0)
-	redisClient.Set(storage.AndroidErrorKey, int64(0), 0)
+	s.client.Set(storage.AndroidErrorKey, int64(0), 0)
 }
 
 // AddTotalCount record push notification count.
 func (s *Storage) AddTotalCount(count int64) {
-	redisClient.IncrBy(storage.TotalCountKey, count)
+	s.client.IncrBy(storage.TotalCountKey, count)
 }
 
 // AddIosSuccess record counts of success iOS push notification.
 func (s *Storage) AddIosSuccess(count int64) {
-	redisClient.IncrBy(storage.IosSuccessKey, count)
+	s.client.IncrBy(storage.IosSuccessKey, count)
 }
 
 // AddIosError record counts of error iOS push notification.
 func (s *Storage) AddIosError(count int64) {
-	redisClient.IncrBy(storage.IosErrorKey, count)
+	s.client.IncrBy(storage.IosErrorKey, count)
 }
 
 // AddAndroidSuccess record counts of success Android push notification.
 func (s *Storage) AddAndroidSuccess(count int64) {
-	redisClient.IncrBy(storage.AndroidSuccessKey, count)
+	s.client.IncrBy(storage.AndroidSuccessKey, count)
 }
 
 // AddAndroidError record counts of error Android push notification.
 func (s *Storage) AddAndroidError(count int64) {
-	redisClient.IncrBy(storage.AndroidErrorKey, count)
+	s.client.IncrBy(storage.AndroidErrorKey, count)
 }
 
 // GetTotalCount show counts of all notification.
 func (s *Storage) GetTotalCount() int64 {
 	var count int64
-	getInt64(storage.TotalCountKey, &count)
+	s.getInt64(storage.TotalCountKey, &count)
 
 	return count
 }
@@ -97,7 +93,7 @@ func (s *Storage) GetTotalCount() int64 {
 // GetIosSuccess show success counts of iOS notification.
 func (s *Storage) GetIosSuccess() int64 {
 	var count int64
-	getInt64(storage.IosSuccessKey, &count)
+	s.getInt64(storage.IosSuccessKey, &count)
 
 	return count
 }
@@ -105,7 +101,7 @@ func (s *Storage) GetIosSuccess() int64 {
 // GetIosError show error counts of iOS notification.
 func (s *Storage) GetIosError() int64 {
 	var count int64
-	getInt64(storage.IosErrorKey, &count)
+	s.getInt64(storage.IosErrorKey, &count)
 
 	return count
 }
@@ -113,7 +109,7 @@ func (s *Storage) GetIosError() int64 {
 // GetAndroidSuccess show success counts of Android notification.
 func (s *Storage) GetAndroidSuccess() int64 {
 	var count int64
-	getInt64(storage.AndroidSuccessKey, &count)
+	s.getInt64(storage.AndroidSuccessKey, &count)
 
 	return count
 }
@@ -121,7 +117,7 @@ func (s *Storage) GetAndroidSuccess() int64 {
 // GetAndroidError show error counts of Android notification.
 func (s *Storage) GetAndroidError() int64 {
 	var count int64
-	getInt64(storage.AndroidErrorKey, &count)
+	s.getInt64(storage.AndroidErrorKey, &count)
 
 	return count
 }

storage/redis/redis_test.go

@@ -69,4 +69,6 @@ func TestRedisEngine(t *testing.T) {
 	wg.Wait()
 	val = redis.GetTotalCount()
 	assert.Equal(t, int64(10), val)
+
+	assert.NoError(t, redis.Close())
 }

storage/storage.go

@@ -31,4 +31,5 @@ type Storage interface {
 	GetIosError() int64
 	GetAndroidSuccess() int64
 	GetAndroidError() int64
+	Close() error
 }