Merge pull request #63 from appleboy/boltdb

Fixed #58 support Boltdb engine.

.gitignore

@@ -32,3 +32,4 @@ coverage.out
 gorush/log/*.log
 build.tar.gz
 gorush.tar.gz
+gorush.db

README.md

@@ -21,7 +21,7 @@ A push notification server using [Gin](https://github.com/gin-gonic/gin) framewo
 * Support notification queue and multiple workers.
 * Support `/api/stat/app` show notification success and failure counts.
 * Support `/api/config` show your yml config.
-* Support store app stat to [redis](http://redis.io/) or memory.
+* Support store app stat to memory, [redis](http://redis.io/) or [BoltDB](https://github.com/boltdb/bolt).
 
 See the [YAML config example](config/config.yml):
 
@@ -60,11 +60,14 @@ log:
   error_level: "error"
 
 stat:
-  engine: "memory" # support memory or redis
+  engine: "memory" # support memory, redis or boltdb
   redis:
     addr: "localhost:6379"
     password: ""
     db: 0
+  boltdb:
+    path: "gorush.db"
+    bucket: "gorush"
 ```
 
 ## Basic Usage

config/config.yml

@@ -37,3 +37,6 @@ stat:
     addr: "localhost:6379"
     password: ""
     db: 0
+  boltdb:
+    path: "gorush.db"
+    bucket: "gorush"

gorush/config.go

@@ -61,8 +61,9 @@ type SectionLog struct {
 
 // SectionStat is sub seciont of config.
 type SectionStat struct {
-	Engine string       `yaml:"service"`
+	Engine string        `yaml:"engine"`
 	Redis  SectionRedis  `yaml:"redis"`
+	BoltDB SectionBoltDB `yaml:"boltdb"`
 }
 
 // SectionRedis is sub seciont of config.
@@ -72,6 +73,12 @@ type SectionRedis struct {
 	DB       int64  `yaml:"db"`
 }
 
+// SectionBoltDB is sub seciont of config.
+type SectionBoltDB struct {
+	Path   string `yaml:"path"`
+	Bucket string `yaml:"bucket"`
+}
+
 // BuildDefaultPushConf is default config setting.
 func BuildDefaultPushConf() ConfYaml {
 	var conf ConfYaml
@@ -114,6 +121,9 @@ func BuildDefaultPushConf() ConfYaml {
 	conf.Stat.Redis.Password = ""
 	conf.Stat.Redis.DB = 0
 
+	conf.Stat.BoltDB.Path = "gorush.db"
+	conf.Stat.BoltDB.Bucket = "gorush"
+
 	return conf
 }
 

gorush/status.go

@@ -1,6 +1,7 @@
 package gorush
 
 import (
+	"github.com/asdine/storm"
 	"github.com/gin-gonic/gin"
 	"gopkg.in/redis.v3"
 	"net/http"
@@ -29,13 +30,6 @@ type IosStatus struct {
 	PushError   int64 `json:"push_error"`
 }
 
-func getRedisInt64Result(key string) int64 {
-	val, _ := RedisClient.Get(key).Result()
-	count, _ := strconv.ParseInt(val, 10, 64)
-
-	return count
-}
-
 func initApp() {
 	RushStatus.TotalCount = 0
 	RushStatus.Ios.PushSuccess = 0
@@ -69,6 +63,14 @@ func initRedis() error {
 	return nil
 }
 
+func initBoltDB() {
+	RushStatus.TotalCount = getTotalCount()
+	RushStatus.Ios.PushSuccess = getIosSuccess()
+	RushStatus.Ios.PushError = getIosError()
+	RushStatus.Android.PushSuccess = getAndroidSuccess()
+	RushStatus.Android.PushError = getAndroidError()
+}
+
 // InitAppStatus for initialize app status
 func InitAppStatus() {
 	switch PushConf.Stat.Engine {
@@ -76,18 +78,39 @@ func InitAppStatus() {
 		initApp()
 	case "redis":
 		initRedis()
+	case "boltdb":
+		initBoltDB()
 	default:
 		initApp()
 	}
 
 }
 
+func getRedisInt64Result(key string, count *int64) {
+	val, _ := RedisClient.Get(key).Result()
+	*count, _ = strconv.ParseInt(val, 10, 64)
+}
+
+func boltdbSet(key string, count int64) {
+	db, _ := storm.Open(PushConf.Stat.BoltDB.Path)
+	db.Set(PushConf.Stat.BoltDB.Bucket, key, count)
+	defer db.Close()
+}
+
+func boltdbGet(key string, count *int64) {
+	db, _ := storm.Open(PushConf.Stat.BoltDB.Path)
+	db.Get(PushConf.Stat.BoltDB.Bucket, key, count)
+	defer db.Close()
+}
+
 func addTotalCount(count int64) {
 	switch PushConf.Stat.Engine {
 	case "memory":
 		atomic.AddInt64(&RushStatus.TotalCount, count)
 	case "redis":
 		RedisClient.Set(gorushTotalCount, strconv.Itoa(int(count)), 0)
+	case "boltdb":
+		boltdbSet(gorushTotalCount, count)
 	default:
 		atomic.AddInt64(&RushStatus.TotalCount, count)
 	}
@@ -99,6 +122,8 @@ func addIosSuccess(count int64) {
 		atomic.AddInt64(&RushStatus.Ios.PushSuccess, count)
 	case "redis":
 		RedisClient.Set(gorushIosSuccess, strconv.Itoa(int(count)), 0)
+	case "boltdb":
+		boltdbSet(gorushIosSuccess, count)
 	default:
 		atomic.AddInt64(&RushStatus.Ios.PushSuccess, count)
 	}
@@ -110,6 +135,8 @@ func addIosError(count int64) {
 		atomic.AddInt64(&RushStatus.Ios.PushError, count)
 	case "redis":
 		RedisClient.Set(gorushIosError, strconv.Itoa(int(count)), 0)
+	case "boltdb":
+		boltdbSet(gorushIosError, count)
 	default:
 		atomic.AddInt64(&RushStatus.Ios.PushError, count)
 	}
@@ -120,8 +147,9 @@ func addAndroidSuccess(count int64) {
 	case "memory":
 		atomic.AddInt64(&RushStatus.Android.PushSuccess, count)
 	case "redis":
-
 		RedisClient.Set(gorushAndroidSuccess, strconv.Itoa(int(count)), 0)
+	case "boltdb":
+		boltdbSet(gorushAndroidSuccess, count)
 	default:
 		atomic.AddInt64(&RushStatus.Android.PushSuccess, count)
 	}
@@ -133,6 +161,8 @@ func addAndroidError(count int64) {
 		atomic.AddInt64(&RushStatus.Android.PushError, count)
 	case "redis":
 		RedisClient.Set(gorushAndroidError, strconv.Itoa(int(count)), 0)
+	case "boltdb":
+		boltdbSet(gorushAndroidError, count)
 	default:
 		atomic.AddInt64(&RushStatus.Android.PushError, count)
 	}
@@ -144,7 +174,9 @@ func getTotalCount() int64 {
 	case "memory":
 		count = atomic.LoadInt64(&RushStatus.TotalCount)
 	case "redis":
-		count = getRedisInt64Result(gorushAndroidError)
+		getRedisInt64Result(gorushTotalCount, &count)
+	case "boltdb":
+		boltdbGet(gorushTotalCount, &count)
 	default:
 		count = atomic.LoadInt64(&RushStatus.TotalCount)
 	}
@@ -158,7 +190,9 @@ func getIosSuccess() int64 {
 	case "memory":
 		count = atomic.LoadInt64(&RushStatus.Ios.PushSuccess)
 	case "redis":
-		count = getRedisInt64Result(gorushAndroidError)
+		getRedisInt64Result(gorushIosSuccess, &count)
+	case "boltdb":
+		boltdbGet(gorushIosSuccess, &count)
 	default:
 		count = atomic.LoadInt64(&RushStatus.Ios.PushSuccess)
 	}
@@ -172,7 +206,9 @@ func getIosError() int64 {
 	case "memory":
 		count = atomic.LoadInt64(&RushStatus.Ios.PushError)
 	case "redis":
-		count = getRedisInt64Result(gorushAndroidError)
+		getRedisInt64Result(gorushIosError, &count)
+	case "boltdb":
+		boltdbGet(gorushIosError, &count)
 	default:
 		count = atomic.LoadInt64(&RushStatus.Ios.PushError)
 	}
@@ -186,7 +222,9 @@ func getAndroidSuccess() int64 {
 	case "memory":
 		count = atomic.LoadInt64(&RushStatus.Android.PushSuccess)
 	case "redis":
-		count = getRedisInt64Result(gorushAndroidError)
+		getRedisInt64Result(gorushAndroidSuccess, &count)
+	case "boltdb":
+		boltdbGet(gorushAndroidSuccess, &count)
 	default:
 		count = atomic.LoadInt64(&RushStatus.Android.PushSuccess)
 	}
@@ -200,7 +238,9 @@ func getAndroidError() int64 {
 	case "memory":
 		count = atomic.LoadInt64(&RushStatus.Android.PushError)
 	case "redis":
-		count = getRedisInt64Result(gorushAndroidError)
+		getRedisInt64Result(gorushAndroidError, &count)
+	case "boltdb":
+		boltdbGet(gorushAndroidError, &count)
 	default:
 		count = atomic.LoadInt64(&RushStatus.Android.PushError)
 	}

gorush/status_test.go

@@ -73,22 +73,22 @@ func TestStatForRedisEngine(t *testing.T) {
 	PushConf.Stat.Redis.Addr = "localhost:6379"
 	InitAppStatus()
 
-	addTotalCount(1000)
+	addTotalCount(10)
-	addIosSuccess(1000)
+	addIosSuccess(20)
-	addIosError(1000)
+	addIosError(30)
-	addAndroidSuccess(1000)
+	addAndroidSuccess(40)
-	addAndroidError(1000)
+	addAndroidError(50)
 
 	val = getTotalCount()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(10), val)
 	val = getIosSuccess()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(20), val)
 	val = getIosError()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(30), val)
 	val = getAndroidSuccess()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(40), val)
 	val = getAndroidError()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(50), val)
 }
 
 func TestDefaultEngine(t *testing.T) {
@@ -96,20 +96,43 @@ func TestDefaultEngine(t *testing.T) {
 	PushConf.Stat.Engine = "test"
 	InitAppStatus()
 
-	addTotalCount(1000)
+	addTotalCount(1)
-	addIosSuccess(1000)
+	addIosSuccess(2)
-	addIosError(1000)
+	addIosError(3)
-	addAndroidSuccess(1000)
+	addAndroidSuccess(4)
-	addAndroidError(1000)
+	addAndroidError(5)
 
 	val = getTotalCount()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(1), val)
 	val = getIosSuccess()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(2), val)
 	val = getIosError()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(3), val)
 	val = getAndroidSuccess()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(4), val)
 	val = getAndroidError()
-	assert.Equal(t, int64(1000), val)
+	assert.Equal(t, int64(5), val)
+}
+
+func TestStatForBoltDBEngine(t *testing.T) {
+	var val int64
+	PushConf.Stat.Engine = "boltdb"
+	InitAppStatus()
+
+	addTotalCount(100)
+	addIosSuccess(200)
+	addIosError(300)
+	addAndroidSuccess(400)
+	addAndroidError(500)
+
+	val = getTotalCount()
+	assert.Equal(t, int64(100), val)
+	val = getIosSuccess()
+	assert.Equal(t, int64(200), val)
+	val = getIosError()
+	assert.Equal(t, int64(300), val)
+	val = getAndroidSuccess()
+	assert.Equal(t, int64(400), val)
+	val = getAndroidError()
+	assert.Equal(t, int64(500), val)
 }