From d82d2bba20f121b1edebfe637935aaff8d526d9b Mon Sep 17 00:00:00 2001 From: Gregory Wells Date: Thu, 11 Jun 2026 18:13:31 -0400 Subject: [PATCH] redo in memory sesisons to have TTLs --- src/session/session_manager.go | 21 ------ src/store/store_in_memory.go | 114 ++++++++++++++++++++++++++------- 2 files changed, 91 insertions(+), 44 deletions(-) diff --git a/src/session/session_manager.go b/src/session/session_manager.go index 38c83c2..7c1dcba 100644 --- a/src/session/session_manager.go +++ b/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) } diff --git a/src/store/store_in_memory.go b/src/store/store_in_memory.go index dcd6c57..02d2537 100644 --- a/src/store/store_in_memory.go +++ b/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 { - value Value - expiry time.Time +type KeyValue[Value any] struct { + value Value + expireTime time.Time +} + +type ExpireTime struct { + expiryTime time.Time + key string } type MemoryStore[Value any] struct { - Keys map[string]ValueEntry[Value] - - Lock sync.RWMutex + 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]{ - value: value, - expiry: time.Now().Add(duration), + m.Keys[hashedkey] = KeyValue[Value]{ + value: value, + 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]{ - value: value, - expiry: current_key.expiry, + m.Keys[hashedkey] = KeyValue[Value]{ + value: value, + 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) +}