🚤 Improve memory storage

memory/memory.go

@@ -11,6 +11,7 @@ type Storage struct {
 	mux        sync.RWMutex
 	db         map[string]entry
 	gcInterval time.Duration
+	done       chan struct{}
 }
 
 // Common storage errors
@@ -30,6 +31,7 @@ func New(config ...Config) *Storage {
 	store := &Storage{
 		db:         make(map[string]entry),
 		gcInterval: cfg.GCInterval,
+		done:       make(chan struct{}),
 	}
 
 	// Start garbage collector
@@ -83,20 +85,29 @@ func (s *Storage) Reset() error {
 	return nil
 }
 
-// Close the database
+// Close the memory storage
 func (s *Storage) Close() error {
+	s.done <- struct{}{}
 	return nil
 }
 
 func (s *Storage) gc() {
-	for t := range time.NewTicker(s.gcInterval).C {
+	ticker := time.NewTicker(s.gcInterval)
-		now := t.Unix()
+	defer ticker.Stop()
-		s.mux.Lock()
+
-		for id, v := range s.db {
+	for {
-			if v.expiry != 0 && v.expiry < now {
+		select {
-				delete(s.db, id)
+		case <-s.done:
+			return
+		case t := <-ticker.C:
+			now := t.Unix()
+			s.mux.Lock()
+			for id, v := range s.db {
+				if v.expiry != 0 && v.expiry < now {
+					delete(s.db, id)
+				}
 			}
+			s.mux.Unlock()
 		}
-		s.mux.Unlock()
 	}
 }

memory/memory_test.go

@@ -115,3 +115,11 @@ func Test_Memory_Clear(t *testing.T) {
 	utils.AssertEqual(t, ErrNotExist, err)
 	utils.AssertEqual(t, true, len(result) == 0)
 }
+
+func Test_Memory_Close(t *testing.T) {
+	s := Storage{done: make(chan struct{}, 1)}
+
+	err := s.Close()
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, 1, len(s.done))
+}