From d516147238b4bddb060c071f92908d99bce80caf Mon Sep 17 00:00:00 2001 From: Sky Johnson Date: Sat, 29 Mar 2025 18:20:31 -0500 Subject: [PATCH] sessions 3 --- core/sessions/Session.go | 100 ++++++++++++++++++++--- core/sessions/SessionManager.go | 136 +++++++++++++++++++++++++++++++- 2 files changed, 225 insertions(+), 11 deletions(-) diff --git a/core/sessions/Session.go b/core/sessions/Session.go index 4864920..efd3775 100644 --- a/core/sessions/Session.go +++ b/core/sessions/Session.go @@ -1,27 +1,40 @@ package sessions import ( + "encoding/json" + "errors" "sync" "time" ) +const ( + DefaultMaxValueSize = 256 * 1024 // 256KB per value +) + +var ( + ErrValueTooLarge = errors.New("session value exceeds size limit") +) + // Session stores data for a single user session type Session struct { - ID string - Data map[string]any - CreatedAt time.Time - UpdatedAt time.Time - mu sync.RWMutex // Protect concurrent access to Data + ID string + Data map[string]any + CreatedAt time.Time + UpdatedAt time.Time + mu sync.RWMutex // Protect concurrent access to Data + maxValueSize int // Maximum size of individual values in bytes + totalDataSize int // Track total size of all data } // NewSession creates a new session with the given ID func NewSession(id string) *Session { now := time.Now() return &Session{ - ID: id, - Data: make(map[string]any), - CreatedAt: now, - UpdatedAt: now, + ID: id, + Data: make(map[string]any), + CreatedAt: now, + UpdatedAt: now, + maxValueSize: DefaultMaxValueSize, } } @@ -33,17 +46,65 @@ func (s *Session) Get(key string) any { } // Set stores a value in the session -func (s *Session) Set(key string, value any) { +func (s *Session) Set(key string, value any) error { + // Estimate value size + size, err := estimateSize(value) + if err != nil { + return err + } + + // Check against limit + if size > s.maxValueSize { + return ErrValueTooLarge + } + s.mu.Lock() defer s.mu.Unlock() + + // If replacing, subtract old value size + if oldVal, exists := s.Data[key]; exists { + oldSize, _ := estimateSize(oldVal) + s.totalDataSize -= oldSize + } + s.Data[key] = value + s.totalDataSize += size s.UpdatedAt = time.Now() + return nil +} + +// SetMaxValueSize changes the maximum allowed value size +func (s *Session) SetMaxValueSize(bytes int) { + s.mu.Lock() + defer s.mu.Unlock() + s.maxValueSize = bytes +} + +// GetMaxValueSize returns the current max value size +func (s *Session) GetMaxValueSize() int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.maxValueSize +} + +// GetTotalSize returns the estimated total size of all session data +func (s *Session) GetTotalSize() int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.totalDataSize } // Delete removes a value from the session func (s *Session) Delete(key string) { s.mu.Lock() defer s.mu.Unlock() + + // Update size tracking + if oldVal, exists := s.Data[key]; exists { + oldSize, _ := estimateSize(oldVal) + s.totalDataSize -= oldSize + } + delete(s.Data, key) s.UpdatedAt = time.Now() } @@ -53,6 +114,7 @@ func (s *Session) Clear() { s.mu.Lock() defer s.mu.Unlock() s.Data = make(map[string]any) + s.totalDataSize = 0 s.UpdatedAt = time.Now() } @@ -69,3 +131,21 @@ func (s *Session) GetAll() map[string]any { return copy } + +// estimateSize approximates the memory footprint of a value +func estimateSize(v any) (int, error) { + // Fast path for common types + switch val := v.(type) { + case string: + return len(val), nil + case []byte: + return len(val), nil + } + + // For other types, use JSON serialization as approximation + data, err := json.Marshal(v) + if err != nil { + return 0, err + } + return len(data), nil +} diff --git a/core/sessions/SessionManager.go b/core/sessions/SessionManager.go index d75453c..28f68cb 100644 --- a/core/sessions/SessionManager.go +++ b/core/sessions/SessionManager.go @@ -1,12 +1,20 @@ package sessions import ( + "container/list" "crypto/rand" "encoding/base64" "sync" + "sync/atomic" "time" ) +const ( + // Default limits + DefaultMaxMemory int64 = 100 * 1024 * 1024 // 100MB default + DefaultMaxSessions int = 10000 // 10K sessions max +) + // SessionManager handles multiple sessions type SessionManager struct { sessions map[string]*Session @@ -18,6 +26,13 @@ type SessionManager struct { cookieHTTPOnly bool cookieMaxAge int gcInterval time.Duration + + // Memory management + maxMemory int64 // Maximum memory limit in bytes + currentMemory int64 // Current estimated memory usage + maxSessions int // Maximum number of sessions + lruList *list.List // For tracking session access order + lruMap map[string]*list.Element // For fast lookups } // NewSessionManager creates a new session manager @@ -29,6 +44,10 @@ func NewSessionManager() *SessionManager { cookieHTTPOnly: true, cookieMaxAge: 86400, // 1 day gcInterval: time.Hour, + maxMemory: DefaultMaxMemory, + maxSessions: DefaultMaxSessions, + lruList: list.New(), + lruMap: make(map[string]*list.Element), } // Start the garbage collector @@ -37,6 +56,61 @@ func NewSessionManager() *SessionManager { return sm } +// trackAccess moves a session to the front of the LRU list +func (sm *SessionManager) trackAccess(id string) { + // No need for sm.mu lock here as this is always called from + // methods that already hold the lock + + if elem, ok := sm.lruMap[id]; ok { + sm.lruList.MoveToFront(elem) + return + } + + // If not in list, add it + elem := sm.lruList.PushFront(id) + sm.lruMap[id] = elem +} + +// evictOldest removes the least recently used session +func (sm *SessionManager) evictOldest() { + // No need for sm.mu lock here as this is always called from + // methods that already hold the lock + + if sm.lruList.Len() == 0 { + return + } + + // Get the oldest session + elem := sm.lruList.Back() + id := elem.Value.(string) + + // Remove from list and map + sm.lruList.Remove(elem) + delete(sm.lruMap, id) + + // Remove session and update memory + if session, exists := sm.sessions[id]; exists { + sessionSize := int64(session.GetTotalSize() + 256) // Base size + data + atomic.AddInt64(&sm.currentMemory, -sessionSize) + delete(sm.sessions, id) + } +} + +// ensureCapacity evicts sessions if we're over the limit +func (sm *SessionManager) ensureCapacity() { + // No lock needed - called from methods that have the lock + + // Check session count limit + for len(sm.sessions) >= sm.maxSessions { + sm.evictOldest() + } + + // Check memory limit + for atomic.LoadInt64(&sm.currentMemory) > sm.maxMemory && sm.lruList.Len() > 0 { + sm.evictOldest() + } +} + // generateSessionID creates a cryptographically secure random session ID func (sm *SessionManager) generateSessionID() string { b := make([]byte, 32) @@ -53,6 +127,10 @@ func (sm *SessionManager) GetSession(id string) *Session { sm.mu.RUnlock() if exists { + // Update LRU status + sm.mu.Lock() + sm.trackAccess(id) + sm.mu.Unlock() return session } @@ -62,11 +140,21 @@ func (sm *SessionManager) GetSession(id string) *Session { // Double check to avoid race conditions if session, exists = sm.sessions[id]; exists { + sm.trackAccess(id) return session } + // Check capacity before creating new session + sm.ensureCapacity() + session = NewSession(id) sm.sessions[id] = session + sm.trackAccess(id) + + // Update memory tracking (estimate base session size + initial map) + sessionSize := int64(256) + atomic.AddInt64(&sm.currentMemory, sessionSize) + return session } @@ -77,8 +165,17 @@ func (sm *SessionManager) CreateSession() *Session { sm.mu.Lock() defer sm.mu.Unlock() + // Ensure we have capacity + sm.ensureCapacity() + session := NewSession(id) sm.sessions[id] = session + sm.trackAccess(id) + + // Update memory tracking (estimate base session size + initial map) + sessionSize := int64(256) + atomic.AddInt64(&sm.currentMemory, sessionSize) + return session } @@ -86,7 +183,35 @@ func (sm *SessionManager) CreateSession() *Session { func (sm *SessionManager) DestroySession(id string) { sm.mu.Lock() defer sm.mu.Unlock() - delete(sm.sessions, id) + + if session, exists := sm.sessions[id]; exists { + // Update memory tracking + sessionSize := int64(session.GetTotalSize() + 256) // Base size + data + atomic.AddInt64(&sm.currentMemory, -sessionSize) + + // Remove from tracking structures + delete(sm.sessions, id) + if elem, ok := sm.lruMap[id]; ok { + sm.lruList.Remove(elem) + delete(sm.lruMap, id) + } + } +} + +// SetMaxMemory sets the maximum memory limit +func (sm *SessionManager) SetMaxMemory(bytes int64) { + atomic.StoreInt64(&sm.maxMemory, bytes) +} + +// SetMaxSessions sets the maximum number of sessions +func (sm *SessionManager) SetMaxSessions(max int) { + sm.mu.Lock() + defer sm.mu.Unlock() + + sm.maxSessions = max + + // If we're over the new limit, evict oldest sessions + sm.ensureCapacity() } // startGC starts the garbage collector to clean up expired sessions @@ -112,7 +237,16 @@ func (sm *SessionManager) gc() { session.mu.RUnlock() if lastUpdated.Before(expiry) { + // Update memory tracking + sessionSize := int64(session.GetTotalSize() + 256) + atomic.AddInt64(&sm.currentMemory, -sessionSize) + + // Remove from tracking structures delete(sm.sessions, id) + if elem, ok := sm.lruMap[id]; ok { + sm.lruList.Remove(elem) + delete(sm.lruMap, id) + } } } }