153 lines
3.1 KiB
Go
153 lines
3.1 KiB
Go
package sessions
|
|
|
|
import (
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/goccy/go-json"
|
|
)
|
|
|
|
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
|
|
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,
|
|
maxValueSize: DefaultMaxValueSize,
|
|
}
|
|
}
|
|
|
|
// Get retrieves a value from the session
|
|
func (s *Session) Get(key string) any {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
return s.Data[key]
|
|
}
|
|
|
|
// Set stores a value in the session
|
|
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()
|
|
}
|
|
|
|
// Clear removes all data from the session
|
|
func (s *Session) Clear() {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.Data = make(map[string]any)
|
|
s.totalDataSize = 0
|
|
s.UpdatedAt = time.Now()
|
|
}
|
|
|
|
// GetAll returns a copy of all session data
|
|
func (s *Session) GetAll() map[string]any {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
|
|
// Create a copy to avoid concurrent map access issues
|
|
copy := make(map[string]any, len(s.Data))
|
|
for k, v := range s.Data {
|
|
copy[k] = v
|
|
}
|
|
|
|
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
|
|
}
|