redo in memory sesisons to have TTLs

src/session/session_manager.go

@@ -38,10 +38,6 @@ func (manager *SessionManager) SetStoreType(storeType StoreType, params ...any)
 	case InMemory:
 		{
 			manager.store = store.NewMemoryStore[*SessionData]()
-			// worker.CreateWorker(time.Minute*5, func() {
-			// 	inMemStore, _ := manager.store.(*store.MemoryStore[*SessionData])
-			// 	cleanupInMemoryStore(inMemStore)
-			// })
 			break
 		}
 	case Redis:
@@ -109,23 +105,6 @@ func (manager *SessionManager) GetSession(r *http.Request) (*SessionData, error)
 	return data, nil
 }
 
-// func cleanupInMemoryStore(m *store.MemoryStore[*SessionData]) {
-// 	logging.Debug("Cleaning up memory store sessions")
-// 	now := time.Now()
-
-// 	m.Lock.Lock()
-// 	defer m.Lock.Unlock()
-
-// 	deleted := 0
-// 	for id, session := range m.Sessions {
-// 		if now.After(session.ExpiresAt) {
-// 			delete(m.Sessions, id)
-// 			deleted = deleted + 1
-// 		}
-// 	}
-// 	logging.Infof("Cleaned up %d stale sessions", deleted)
-// }
-
 func (manager *SessionManager) DeleteSession(sessionId string) error {
 	return manager.store.Delete(sessionId)
 }

src/store/store_in_memory.go

@@ -5,58 +5,122 @@ import (
 	"time"
 
 	"astraltech.xyz/accountmanager/src/logging"
+	"astraltech.xyz/accountmanager/src/worker"
 )
 
-type ValueEntry[Value any] struct {
+type KeyValue[Value any] struct {
 	value      Value
-	expiry time.Time
+	expireTime time.Time
+}
+
+type ExpireTime struct {
+	expiryTime time.Time
+	key        string
 }
 
 type MemoryStore[Value any] struct {
-	Keys map[string]ValueEntry[Value]
-
+	Keys        map[string]KeyValue[Value]
 	Lock        sync.RWMutex
+	ExpireTimes []ExpireTime
 }
 
 func NewMemoryStore[Value any]() *MemoryStore[Value] {
 	logging.Debug("Creating new in memory session store")
 	store := &MemoryStore[Value]{
-		Keys: make(map[string]ValueEntry[Value]),
+		Keys:        make(map[string]KeyValue[Value]),
+		Lock:        sync.RWMutex{},
+		ExpireTimes: []ExpireTime{},
 	}
+	worker.CreateWorker(time.Minute*5, store.CleanupExpired)
 	return store
 }
 
 func (m *MemoryStore[Value]) CreateExpiring(key string, value Value, duration time.Duration) error {
-	hashedkey := HashKey(key)
+	var expiry time.Time
+	if duration != 0 {
+		expiry = time.Now().Add(duration)
+	}
 
 	m.Lock.Lock()
 	defer m.Lock.Unlock()
+
+	hashedkey := HashKey(key)
 	_, exist := m.Keys[hashedkey]
 	if exist {
 		return ErrKeyAlreadyExists
 	}
 
-	m.Keys[hashedkey] = ValueEntry[Value]{
+	m.Keys[hashedkey] = KeyValue[Value]{
 		value:      value,
-		expiry: time.Now().Add(duration),
+		expireTime: expiry,
 	}
+
+	if duration != 0 {
+		low := 0
+		high := len(m.ExpireTimes)
+
+		for low < high {
+			mid := (low + high) / 2
+
+			if m.ExpireTimes[mid].expiryTime.Before(expiry) {
+				low = mid + 1
+			} else {
+				high = mid
+			}
+		}
+		insertIndex := low
+
+		m.ExpireTimes = append(m.ExpireTimes, ExpireTime{})
+		copy(m.ExpireTimes[insertIndex+1:], m.ExpireTimes[insertIndex:])
+		m.ExpireTimes[insertIndex] = ExpireTime{
+			key:        key,
+			expiryTime: expiry,
+		}
+	}
+
 	return nil
 }
 
+func (m *MemoryStore[Value]) CleanupExpired() {
+	m.Lock.Lock()
+	defer m.Lock.Unlock()
+
+	now := time.Now()
+	for len(m.ExpireTimes) > 0 && !now.Before(m.ExpireTimes[0].expiryTime) {
+		key, exists := m.Keys[HashKey(m.ExpireTimes[0].key)]
+		if exists && key.expireTime.Equal(m.ExpireTimes[0].expiryTime) {
+			m.deleteWithoutLock(m.ExpireTimes[0].key)
+		}
+		m.ExpireTimes = m.ExpireTimes[1:]
+	}
+}
+
 func (m *MemoryStore[Value]) Create(key string, session Value) error {
 	return m.CreateExpiring(key, session, 0)
 }
 
 func (m *MemoryStore[Value]) Get(key string) (Value, error) {
-	var data ValueEntry[Value]
+	var data KeyValue[Value]
+	var zeroValue Value
+	hashedkey := HashKey(key)
 
 	m.Lock.RLock()
-	hashedkey := HashKey(key)
 	data, exists := m.Keys[hashedkey]
 	m.Lock.RUnlock()
-	if exists == false {
-		return data.value, ErrKeyNotFound
+
+	if !exists {
+		return zeroValue, ErrKeyNotFound
 	}
+
+	if !data.expireTime.IsZero() && !time.Now().Before(data.expireTime) {
+		m.Lock.Lock()
+		if current, ok := m.Keys[hashedkey]; ok && current.expireTime.Equal(data.expireTime) {
+			delete(m.Keys, hashedkey)
+		}
+		m.Lock.Unlock()
+		return zeroValue, ErrKeyNotFound
+	}
+
 	return data.value, nil
 }
 
@@ -65,22 +129,20 @@ func (m *MemoryStore[Value]) Update(sessionID string, value Value) error {
 
 	m.Lock.Lock()
 	defer m.Lock.Unlock()
-	current_key, exist := m.Keys[hashedkey]
+	old_key, exist := m.Keys[hashedkey]
 	if !exist {
 		return ErrKeyNotFound
 	}
-	m.Keys[hashedkey] = ValueEntry[Value]{
+	m.Keys[hashedkey] = KeyValue[Value]{
 		value:      value,
-		expiry: current_key.expiry,
+		expireTime: old_key.expireTime,
 	}
 	return nil
 }
 
-func (m *MemoryStore[Value]) Delete(sessionID string) error {
-	hashedkey := HashKey(sessionID)
+func (m *MemoryStore[Value]) deleteWithoutLock(key string) error {
+	hashedkey := HashKey(key)
 
-	m.Lock.Lock()
-	defer m.Lock.Unlock()
 	_, exist := m.Keys[hashedkey]
 	if !exist {
 		return ErrKeyNotFound
@@ -88,3 +150,9 @@ func (m *MemoryStore[Value]) Delete(sessionID string) error {
 	delete(m.Keys, hashedkey)
 	return nil
 }
+
+func (m *MemoryStore[Value]) Delete(key string) error {
+	m.Lock.Lock()
+	defer m.Lock.Unlock()
+	return m.deleteWithoutLock(key)
+}