MetaLab/internal/session/memory_store.go

package session

import (
	"log"
	"sync"
	"time"
)

// MemoryStore 基于内存的会话存储（sync.Map + RWMutex，适合单实例部署）
type MemoryStore struct {
	mu              sync.RWMutex
	sessions        map[string]*Session // sid → session
	byUID           map[uint][]string   // uid → []sid（用于批量删除）
	idleTimeout     time.Duration       // 不记住我：空闲超时
	rememberTimeout time.Duration       // 记住我：空闲超时
}

// NewMemoryStore 创建内存存储实例
func NewMemoryStore(idleTimeout, rememberTimeout time.Duration) *MemoryStore {
	return &MemoryStore{
		sessions:        make(map[string]*Session),
		byUID:           make(map[uint][]string),
		idleTimeout:     idleTimeout,
		rememberTimeout: rememberTimeout,
	}
}

// Get 获取会话（根据 RememberMe 选择对应的空闲超时检查过期）
func (ms *MemoryStore) Get(id string) (*Session, error) {
	ms.mu.RLock()
	s, ok := ms.sessions[id]
	ms.mu.RUnlock()
	if !ok {
		return nil, nil
	}
	// 记住我 → 使用更长的超时
	timeout := ms.idleTimeout
	if s.RememberMe {
		timeout = ms.rememberTimeout
	}
	if s.IsExpired(timeout) {
		ms.Delete(id)
		return nil, nil
	}
	return s, nil
}

// Set 写入会话（同步更新 byUID 索引）
func (ms *MemoryStore) Set(s *Session) error {
	ms.mu.Lock()
	defer ms.mu.Unlock()

	// 如果 sid 已存在且 uid 不同，先从旧 uid 索引移除
	if old, ok := ms.sessions[s.ID]; ok && old.UserID != s.UserID {
		ms.removeFromUIDIndex(old.UserID, s.ID)
	}

	ms.sessions[s.ID] = s
	ms.byUID[s.UserID] = append(ms.byUID[s.UserID], s.ID)
	return nil
}

// Delete 删除会话（同步更新 byUID 索引）
func (ms *MemoryStore) Delete(id string) error {
	ms.mu.Lock()
	defer ms.mu.Unlock()

	s, ok := ms.sessions[id]
	if !ok {
		return nil
	}
	ms.removeFromUIDIndex(s.UserID, id)
	delete(ms.sessions, id)
	return nil
}

// DeleteByUID 删除某用户的所有会话
func (ms *MemoryStore) DeleteByUID(uid uint) error {
	ms.mu.Lock()
	defer ms.mu.Unlock()

	sids, ok := ms.byUID[uid]
	if !ok {
		return nil
	}
	for _, sid := range sids {
		delete(ms.sessions, sid)
	}
	delete(ms.byUID, uid)
	log.Printf("[MemoryStore] DeleteByUID uid=%d count=%d", uid, len(sids))
	return nil
}

// Cleanup 清理过期会话（根据 RememberMe 选择对应超时）
func (ms *MemoryStore) Cleanup() int {
	ms.mu.Lock()
	defer ms.mu.Unlock()

	var expiredSids []string
	for sid, s := range ms.sessions {
		timeout := ms.idleTimeout
		if s.RememberMe {
			timeout = ms.rememberTimeout
		}
		if s.IsExpired(timeout) {
			expiredSids = append(expiredSids, sid)
		}
	}
	for _, sid := range expiredSids {
		s := ms.sessions[sid]
		ms.removeFromUIDIndex(s.UserID, sid)
		delete(ms.sessions, sid)
	}
	if len(expiredSids) > 0 {
		log.Printf("[MemoryStore] Cleanup removed=%d remaining=%d", len(expiredSids), len(ms.sessions))
	}
	return len(expiredSids)
}

// removeFromUIDIndex 从 byUID 索引中移除指定 sid（调用者需持有写锁）
func (ms *MemoryStore) removeFromUIDIndex(uid uint, sid string) {
	sids := ms.byUID[uid]
	for i, s := range sids {
		if s == sid {
			ms.byUID[uid] = append(sids[:i], sids[i+1:]...)
			break
		}
	}
	if len(ms.byUID[uid]) == 0 {
		delete(ms.byUID, uid)
	}
}

// StartCleanup 启动后台清理 goroutine（应在 main 中调用）
func (ms *MemoryStore) StartCleanup(interval time.Duration) {
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for range ticker.C {
			ms.Cleanup()
		}
	}()
}