From b58386704c4cc4a90f14afb85275c3e11c3611d3 Mon Sep 17 00:00:00 2001
From: "Yaroslav \"Zorg\" Zborovsky" <yaronius@users.noreply.github.com>
Date: Sat, 29 Feb 2020 11:14:44 +0200
Subject: [PATCH] Fix Redis stats storage concurrency issues (#478)

---
 storage/redis/redis.go      | 27 ++++++++++++---------------
 storage/redis/redis_test.go | 14 ++++++++++++++
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/storage/redis/redis.go b/storage/redis/redis.go
index a40bbc7..ca843aa 100644
--- a/storage/redis/redis.go
+++ b/storage/redis/redis.go
@@ -46,47 +46,44 @@ func (s *Storage) Init() error {
 
 		return err
 	}
+	// Set initial values
+	s.Reset()
 
 	return nil
 }
 
 // Reset Client storage.
 func (s *Storage) Reset() {
-	redisClient.Set(storage.TotalCountKey, strconv.Itoa(0), 0)
-	redisClient.Set(storage.IosSuccessKey, strconv.Itoa(0), 0)
-	redisClient.Set(storage.IosErrorKey, strconv.Itoa(0), 0)
-	redisClient.Set(storage.AndroidSuccessKey, strconv.Itoa(0), 0)
-	redisClient.Set(storage.AndroidErrorKey, strconv.Itoa(0), 0)
+	redisClient.Set(storage.TotalCountKey, int64(0), 0)
+	redisClient.Set(storage.IosSuccessKey, int64(0), 0)
+	redisClient.Set(storage.IosErrorKey, int64(0), 0)
+	redisClient.Set(storage.AndroidSuccessKey, int64(0), 0)
+	redisClient.Set(storage.AndroidErrorKey, int64(0), 0)
 }
 
 // AddTotalCount record push notification count.
 func (s *Storage) AddTotalCount(count int64) {
-	total := s.GetTotalCount() + count
-	redisClient.Set(storage.TotalCountKey, strconv.Itoa(int(total)), 0)
+	redisClient.IncrBy(storage.TotalCountKey, count)
 }
 
 // AddIosSuccess record counts of success iOS push notification.
 func (s *Storage) AddIosSuccess(count int64) {
-	total := s.GetIosSuccess() + count
-	redisClient.Set(storage.IosSuccessKey, strconv.Itoa(int(total)), 0)
+	redisClient.IncrBy(storage.IosSuccessKey, count)
 }
 
 // AddIosError record counts of error iOS push notification.
 func (s *Storage) AddIosError(count int64) {
-	total := s.GetIosError() + count
-	redisClient.Set(storage.IosErrorKey, strconv.Itoa(int(total)), 0)
+	redisClient.IncrBy(storage.IosErrorKey, count)
 }
 
 // AddAndroidSuccess record counts of success Android push notification.
 func (s *Storage) AddAndroidSuccess(count int64) {
-	total := s.GetAndroidSuccess() + count
-	redisClient.Set(storage.AndroidSuccessKey, strconv.Itoa(int(total)), 0)
+	redisClient.IncrBy(storage.AndroidSuccessKey, count)
 }
 
 // AddAndroidError record counts of error Android push notification.
 func (s *Storage) AddAndroidError(count int64) {
-	total := s.GetAndroidError() + count
-	redisClient.Set(storage.AndroidErrorKey, strconv.Itoa(int(total)), 0)
+	redisClient.IncrBy(storage.AndroidErrorKey, count)
 }
 
 // GetTotalCount show counts of all notification.
diff --git a/storage/redis/redis_test.go b/storage/redis/redis_test.go
index 1e8a705..d3363a6 100644
--- a/storage/redis/redis_test.go
+++ b/storage/redis/redis_test.go
@@ -1,6 +1,7 @@
 package redis
 
 import (
+	"sync"
 	"testing"
 
 	c "github.com/appleboy/gorush/config"
@@ -55,4 +56,17 @@ func TestRedisEngine(t *testing.T) {
 	redis.Reset()
 	val = redis.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() {
+			redis.AddTotalCount(1)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+	val = redis.GetTotalCount()
+	assert.Equal(t, int64(10), val)
 }