improve memory storage code and performance

internal/storage/memory/config.go

@@ -0,0 +1,33 @@
+package memory
+
+import "time"
+
+// Config defines the config for storage.
+type Config struct {
+	// Time before deleting expired keys
+	//
+	// Default is 10 * time.Second
+	GCInterval time.Duration
+}
+
+// ConfigDefault is the default config
+var ConfigDefault = Config{
+	GCInterval: 10 * time.Second,
+}
+
+// configDefault is a helper function to set default values
+func configDefault(config ...Config) Config {
+	// Return default config if nothing provided
+	if len(config) < 1 {
+		return ConfigDefault
+	}
+
+	// Override default config
+	cfg := config[0]
+
+	// Set default values
+	if int(cfg.GCInterval.Seconds()) <= 0 {
+		cfg.GCInterval = ConfigDefault.GCInterval
+	}
+	return cfg
+}

internal/storage/memory/memory.go

@@ -2,7 +2,10 @@ package memory
 
 import (
 	"sync"
+	"sync/atomic"
 	"time"
+
+	"github.com/gofiber/fiber/v2/utils"
 )
 
 // Storage interface that is implemented by storage providers
@@ -14,21 +17,25 @@ type Storage struct {
 }
 
 type entry struct {
+	data []byte
 	// max value is 4294967295 -> Sun Feb 07 2106 06:28:15 GMT+0000
 	expiry uint32
-	data   []byte
 }
 
 // New creates a new memory storage
-func New() *Storage {
+func New(config ...Config) *Storage {
+	// Set default config
+	cfg := configDefault(config...)
+
 	// Create storage
 	store := &Storage{
 		db:         make(map[string]entry),
-		gcInterval: 10 * time.Second,
+		gcInterval: cfg.GCInterval,
 		done:       make(chan struct{}),
 	}
 
 	// Start garbage collector
+	utils.StartTimeStampUpdater()
 	go store.gc()
 
 	return store
@@ -42,7 +49,7 @@ func (s *Storage) Get(key string) ([]byte, error) {
 	s.mux.RLock()
 	v, ok := s.db[key]
 	s.mux.RUnlock()
-	if !ok || v.expiry != 0 && v.expiry <= uint32(time.Now().Unix()) {
+	if !ok || v.expiry != 0 && v.expiry <= atomic.LoadUint32(&utils.Timestamp) {
 		return nil, nil
 	}
 
@@ -58,11 +65,11 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 
 	var expire uint32
 	if exp != 0 {
-		expire = uint32(time.Now().Add(exp).Unix())
+		expire = uint32(exp.Seconds()) + atomic.LoadUint32(&utils.Timestamp)
 	}
 
 	s.mux.Lock()
-	s.db[key] = entry{expire, val}
+	s.db[key] = entry{val, expire}
 	s.mux.Unlock()
 	return nil
 }
@@ -96,20 +103,31 @@ func (s *Storage) Close() error {
 func (s *Storage) gc() {
 	ticker := time.NewTicker(s.gcInterval)
 	defer ticker.Stop()
+	var expired []string
 
 	for {
 		select {
 		case <-s.done:
 			return
-		case t := <-ticker.C:
-			now := uint32(t.Unix())
-			s.mux.Lock()
+		case <-ticker.C:
+			expired = expired[:0]
+			s.mux.RLock()
 			for id, v := range s.db {
-				if v.expiry != 0 && v.expiry < now {
-					delete(s.db, id)
+				if v.expiry != 0 && v.expiry < atomic.LoadUint32(&utils.Timestamp) {
+					expired = append(expired, id)
 				}
 			}
+			s.mux.RUnlock()
+			s.mux.Lock()
+			for i := range expired {
+				delete(s.db, expired[i])
+			}
 			s.mux.Unlock()
 		}
 	}
 }
+
+// Return database client
+func (s *Storage) Conn() map[string]entry {
+	return s.db
+}

internal/storage/memory/memory_test.go

@@ -0,0 +1,204 @@
+package memory
+
+import (
+	"testing"
+	"time"
+
+	"github.com/gofiber/fiber/v2/utils"
+)
+
+var testStore = New()
+
+func Test_Memory_Set(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	err := testStore.Set(key, val, 0)
+	utils.AssertEqual(t, nil, err)
+}
+
+func Test_Memory_Set_Override(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	err := testStore.Set(key, val, 0)
+	utils.AssertEqual(t, nil, err)
+
+	err = testStore.Set(key, val, 0)
+	utils.AssertEqual(t, nil, err)
+}
+
+func Test_Memory_Get(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	err := testStore.Set(key, val, 0)
+	utils.AssertEqual(t, nil, err)
+
+	result, err := testStore.Get(key)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, val, result)
+}
+
+func Test_Memory_Set_Expiration(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+		exp = 1 * time.Second
+	)
+
+	err := testStore.Set(key, val, exp)
+	utils.AssertEqual(t, nil, err)
+
+	time.Sleep(1100 * time.Millisecond)
+}
+
+func Test_Memory_Get_Expired(t *testing.T) {
+	var (
+		key = "john"
+	)
+
+	result, err := testStore.Get(key)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, len(result) == 0)
+}
+
+func Test_Memory_Get_NotExist(t *testing.T) {
+
+	result, err := testStore.Get("notexist")
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, len(result) == 0)
+}
+
+func Test_Memory_Delete(t *testing.T) {
+	var (
+		key = "john"
+		val = []byte("doe")
+	)
+
+	err := testStore.Set(key, val, 0)
+	utils.AssertEqual(t, nil, err)
+
+	err = testStore.Delete(key)
+	utils.AssertEqual(t, nil, err)
+
+	result, err := testStore.Get(key)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, len(result) == 0)
+}
+
+func Test_Memory_Reset(t *testing.T) {
+	var (
+		val = []byte("doe")
+	)
+
+	err := testStore.Set("john1", val, 0)
+	utils.AssertEqual(t, nil, err)
+
+	err = testStore.Set("john2", val, 0)
+	utils.AssertEqual(t, nil, err)
+
+	err = testStore.Reset()
+	utils.AssertEqual(t, nil, err)
+
+	result, err := testStore.Get("john1")
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, len(result) == 0)
+
+	result, err = testStore.Get("john2")
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, len(result) == 0)
+}
+
+func Test_Memory_Close(t *testing.T) {
+	utils.AssertEqual(t, nil, testStore.Close())
+}
+
+func Test_Memory_Conn(t *testing.T) {
+	utils.AssertEqual(t, true, testStore.Conn() != nil)
+}
+
+// go test -run Test_Memory -v -race
+
+func Test_Memory(t *testing.T) {
+	var store = New()
+	var (
+		key = "john"
+		val = []byte("doe")
+		exp = 1 * time.Second
+	)
+
+	store.Set(key, val, 0)
+
+	result, error := store.Get(key)
+	utils.AssertEqual(t, val, result)
+	utils.AssertEqual(t, nil, error)
+
+	result, error = store.Get("empty")
+	utils.AssertEqual(t, nil, result)
+	utils.AssertEqual(t, nil, error)
+
+	store.Set(key, val, exp)
+	time.Sleep(1100 * time.Millisecond)
+
+	result, error = store.Get(key)
+	utils.AssertEqual(t, nil, result)
+	utils.AssertEqual(t, nil, error)
+
+	store.Set(key, val, 0)
+	result, error = store.Get(key)
+	utils.AssertEqual(t, val, result)
+	utils.AssertEqual(t, nil, error)
+
+	store.Delete(key)
+	result, error = store.Get(key)
+	utils.AssertEqual(t, nil, result)
+	utils.AssertEqual(t, nil, error)
+
+	store.Set("john", val, 0)
+	store.Set("doe", val, 0)
+	store.Reset()
+
+	result, error = store.Get("john")
+	utils.AssertEqual(t, nil, result)
+	utils.AssertEqual(t, nil, error)
+
+	result, error = store.Get("doe")
+	utils.AssertEqual(t, nil, result)
+	utils.AssertEqual(t, nil, error)
+
+}
+
+// go test -v -run=^$ -bench=Benchmark_Memory -benchmem -count=4
+func Benchmark_Memory(b *testing.B) {
+	keyLength := 1000
+	keys := make([]string, keyLength)
+	for i := 0; i < keyLength; i++ {
+		keys[i] = utils.UUID()
+	}
+	value := []byte("joe")
+
+	ttl := 2 * time.Second
+	b.Run("fiber_memory", func(b *testing.B) {
+		d := New()
+		b.ReportAllocs()
+		b.ResetTimer()
+		for n := 0; n < b.N; n++ {
+			for _, key := range keys {
+				d.Set(key, value, ttl)
+			}
+			for _, key := range keys {
+				_, _ = d.Get(key)
+			}
+			for _, key := range keys {
+				d.Delete(key)
+			}
+		}
+	})
+}

utils/time.go

@@ -0,0 +1,28 @@
+package utils
+
+import (
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	timestampTimer sync.Once
+	Timestamp      uint32
+)
+
+func StartTimeStampUpdater() {
+	timestampTimer.Do(func() {
+		go func(sleep time.Duration) {
+			ticker := time.NewTicker(sleep)
+			defer ticker.Stop()
+			for {
+				select {
+				case t := <-ticker.C:
+					// update timestamp
+					atomic.StoreUint32(&Timestamp, uint32(t.Unix()))
+				}
+			}
+		}(1 * time.Second) // duration
+	})
+}