visionA/visionA-backend/internal/session/inmemory_store.go

package session

import (
	"context"
	"sync"
	"time"
)

// InMemoryStore 是 Store 的單節點記憶體實作，只部署在 remote-proxy binary。
//
// 語意重點（對齊 tunnel.md §2.3 / §5.2 + Q5 裁決）：
//   - 同 token 後連覆蓋前連 — Register 時若已存在，先 Close() 舊 handle 再寫入。
//   - Heartbeat 僅更新 Summary.LastHeartbeat 時間戳。
//   - CleanupExpired 掃描並移除逾時者（同時 Close 對應 handle）。
//   - 所有操作並發安全（sync.RWMutex）。
type InMemoryStore struct {
	mu       sync.RWMutex
	sessions map[string]Handle
}

// NewInMemoryStore 建立一個空的記憶體 session store。
func NewInMemoryStore() *InMemoryStore {
	return &InMemoryStore{
		sessions: make(map[string]Handle),
	}
}

// Register 註冊一個 session；同 token 的舊 session 會先被 Close 再覆蓋（Q5 裁決）。
func (s *InMemoryStore) Register(ctx context.Context, token string, h Handle) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if old, ok := s.sessions[token]; ok {
		// 後連覆蓋前連：關閉舊的 handle 以釋放 yamux / WS 資源。
		// Close 錯誤忽略 — 舊的可能已經斷線，這不影響新連線的註冊。
		_ = old.Close()
	}
	s.sessions[token] = h
	return nil
}

// Unregister 移除指定 token；不存在為 no-op。
func (s *InMemoryStore) Unregister(ctx context.Context, token string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.sessions, token)
	return nil
}

// Lookup 回傳指定 token 的 handle；不存在回 ErrSessionNotFound。
func (s *InMemoryStore) Lookup(ctx context.Context, token string) (Handle, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	h, ok := s.sessions[token]
	if !ok {
		return nil, ErrSessionNotFound
	}
	return h, nil
}

// Exists 判斷 token 是否有 active session。
func (s *InMemoryStore) Exists(ctx context.Context, token string) (bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	_, ok := s.sessions[token]
	return ok, nil
}

// List 回傳所有 active session 的 summary。
func (s *InMemoryStore) List(ctx context.Context) ([]*Summary, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	out := make([]*Summary, 0, len(s.sessions))
	for _, h := range s.sessions {
		if sum := h.Summary(); sum != nil {
			// 複製 Summary 避免 caller 誤改內部狀態
			cp := *sum
			out = append(out, &cp)
		}
	}
	return out, nil
}

// Heartbeat 更新 session 的 LastHeartbeat 時間。
//
// 修 B2 Review M1（race condition）：改為呼叫 Handle.RecordHeartbeat，
// 由各實作自行用 mutex / atomic 保護 LastHeartbeat 欄位，
// 避免 Store.Heartbeat 與 Store.CleanupExpired / Store.List 對同一 Summary pointer
// 的並發讀寫被 race detector 捕捉。
func (s *InMemoryStore) Heartbeat(ctx context.Context, token string) error {
	s.mu.RLock()
	h, ok := s.sessions[token]
	s.mu.RUnlock()

	if !ok {
		return ErrSessionNotFound
	}
	h.RecordHeartbeat(time.Now().UTC())
	return nil
}

// CleanupExpired 清除 LastHeartbeat 超過 expireAfter 的 session。
//
// 實作步驟：
//  1. 在讀鎖下找出所有過期 token（避免長時間持寫鎖）
//  2. 升級為寫鎖，逐一移除（二次檢查避免 race）
//  3. Close 對應 handle 釋放資源
func (s *InMemoryStore) CleanupExpired(ctx context.Context, expireAfter time.Duration) (int, error) {
	cutoff := time.Now().UTC().Add(-expireAfter)

	s.mu.Lock()
	defer s.mu.Unlock()

	removed := 0
	for token, h := range s.sessions {
		sum := h.Summary()
		if sum == nil {
			continue
		}
		if sum.LastHeartbeat.Before(cutoff) {
			_ = h.Close()
			delete(s.sessions, token)
			removed++
		}
	}
	return removed, nil
}

// 編譯時檢查：確保 InMemoryStore 實作 Store。
var _ Store = (*InMemoryStore)(nil)