From a0645af1ed365a7fbfd605f3ad880ee9336ad2e3 Mon Sep 17 00:00:00 2001
From: leonklingele <git@leonklingele.de>
Date: Mon, 14 Nov 2022 08:21:03 +0100
Subject: [PATCH] Fix and optimize memory storage (#2207)

* internal/memory: cache timestamp

* internal/memory: ensure to never delete non-expired items

This fixes a TOCTOU problem between a mutex rlock and a mutex lock.

* internal/memory: move costly operations outside of locked area

* internal/storage: cache timestamp

* internal/storage: ensure to never delete non-expired items

This fixes a TOCTOU problem between a mutex rlock and a mutex lock.

* internal/storage: move costly operations outside of locked area
---
 internal/memory/memory.go         | 16 ++++++++++++----
 internal/storage/memory/memory.go | 16 ++++++++++++----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/internal/memory/memory.go b/internal/memory/memory.go
index 0c37aa9c2e..ff2202847b 100644
--- a/internal/memory/memory.go
+++ b/internal/memory/memory.go
@@ -47,8 +47,9 @@ func (s *Storage) Set(key string, val interface{}, ttl time.Duration) {
 	if ttl > 0 {
 		exp = uint32(ttl.Seconds()) + atomic.LoadUint32(&utils.Timestamp)
 	}
+	i := item{exp, val}
 	s.Lock()
-	s.data[key] = item{exp, val}
+	s.data[key] = i
 	s.Unlock()
 }
 
@@ -61,8 +62,9 @@ func (s *Storage) Delete(key string) {
 
 // Reset all keys
 func (s *Storage) Reset() {
+	nd := make(map[string]item)
 	s.Lock()
-	s.data = make(map[string]item)
+	s.data = nd
 	s.Unlock()
 }
 
@@ -74,17 +76,23 @@ func (s *Storage) gc(sleep time.Duration) {
 	for {
 		select {
 		case <-ticker.C:
+			ts := atomic.LoadUint32(&utils.Timestamp)
 			expired = expired[:0]
 			s.RLock()
 			for key, v := range s.data {
-				if v.e != 0 && v.e <= atomic.LoadUint32(&utils.Timestamp) {
+				if v.e != 0 && v.e <= ts {
 					expired = append(expired, key)
 				}
 			}
 			s.RUnlock()
 			s.Lock()
+			// Double-checked locking.
+			// We might have replaced the item in the meantime.
 			for i := range expired {
-				delete(s.data, expired[i])
+				v := s.data[expired[i]]
+				if v.e != 0 && v.e <= ts {
+					delete(s.data, expired[i])
+				}
 			}
 			s.Unlock()
 		}
diff --git a/internal/storage/memory/memory.go b/internal/storage/memory/memory.go
index ff43c30533..1a56106106 100644
--- a/internal/storage/memory/memory.go
+++ b/internal/storage/memory/memory.go
@@ -70,8 +70,9 @@ func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
 		expire = uint32(exp.Seconds()) + atomic.LoadUint32(&utils.Timestamp)
 	}
 
+	e := entry{val, expire}
 	s.mux.Lock()
-	s.db[key] = entry{val, expire}
+	s.db[key] = e
 	s.mux.Unlock()
 	return nil
 }
@@ -90,8 +91,9 @@ func (s *Storage) Delete(key string) error {
 
 // Reset all keys
 func (s *Storage) Reset() error {
+	ndb := make(map[string]entry)
 	s.mux.Lock()
-	s.db = make(map[string]entry)
+	s.db = ndb
 	s.mux.Unlock()
 	return nil
 }
@@ -112,17 +114,23 @@ func (s *Storage) gc() {
 		case <-s.done:
 			return
 		case <-ticker.C:
+			ts := atomic.LoadUint32(&utils.Timestamp)
 			expired = expired[:0]
 			s.mux.RLock()
 			for id, v := range s.db {
-				if v.expiry != 0 && v.expiry < atomic.LoadUint32(&utils.Timestamp) {
+				if v.expiry != 0 && v.expiry <= ts {
 					expired = append(expired, id)
 				}
 			}
 			s.mux.RUnlock()
 			s.mux.Lock()
+			// Double-checked locking.
+			// We might have replaced the item in the meantime.
 			for i := range expired {
-				delete(s.db, expired[i])
+				v := s.db[expired[i]]
+				if v.expiry != 0 && v.expiry <= ts {
+					delete(s.db, expired[i])
+				}
 			}
 			s.mux.Unlock()
 		}