90CNServer/server/internal/platform/replay/memory.go

package replay

import (
	"container/heap"
	"container/list"
	"context"
	"sync"
	"time"
)

const defaultMemMaxEntries = 1024

type memoryEntry struct {
	key        string
	expiresAt  time.Time
	generation uint64
}

// MemoryStore 是单进程 LRU + TTL nonce 表。它只适合测试、单节点后台或降级路径；
// 多 gateway / logic 实例必须换成共享 Redis，否则不同进程会各自接受同一 nonce。
type MemoryStore struct {
	mu         sync.Mutex
	maxEntries int
	items      map[string]*list.Element
	order      *list.List
	expiry     expiryHeap
	generation uint64
	now        func() time.Time
}

// NewMemoryStore 创建进程内 nonce 表。maxEntries<=0 时使用保守默认值，避免调用方
// 因配置漏填得到一个永远无法拒重的空表。
func NewMemoryStore(maxEntries int) *MemoryStore {
	if maxEntries <= 0 {
		maxEntries = defaultMemMaxEntries
	}
	return &MemoryStore{
		maxEntries: maxEntries,
		items:      make(map[string]*list.Element, maxEntries),
		order:      list.New(),
		expiry:     expiryHeap{},
		now:        time.Now,
	}
}

// Seen 记录 nonce 并返回它是否首次出现；重复 nonce 会返回 false。
func (s *MemoryStore) Seen(ctx context.Context, key string, ttl time.Duration) (bool, error) {
	_, key, err := normalizeInput(ctx, key, ttl)
	if err != nil {
		return false, err
	}
	if s == nil {
		return false, ErrReplayStoreUnavailable
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.ensureLocked()
	now := s.now().UTC()
	s.pruneExpiredLocked(now)
	if existing := s.items[key]; existing != nil {
		entry := existing.Value.(memoryEntry)
		if entry.expiresAt.After(now) {
			s.order.MoveToFront(existing)
			return false, nil
		}
		s.removeLocked(existing)
	}
	if len(s.items) >= s.maxEntries {
		return false, ErrReplayStoreFull
	}
	s.generation++
	entry := memoryEntry{key: key, expiresAt: now.Add(ttl), generation: s.generation}
	s.items[key] = s.order.PushFront(entry)
	heap.Push(&s.expiry, expiryItem{key: key, expiresAt: entry.expiresAt, generation: entry.generation})
	s.compactExpiryLocked()
	return true, nil
}

func (s *MemoryStore) ensureLocked() {
	if s.maxEntries <= 0 {
		s.maxEntries = defaultMemMaxEntries
	}
	if s.items == nil {
		s.items = make(map[string]*list.Element, s.maxEntries)
	}
	if s.order == nil {
		s.order = list.New()
	}
	if s.expiry == nil {
		s.expiry = expiryHeap{}
	}
	if s.now == nil {
		s.now = time.Now
	}
}

func (s *MemoryStore) pruneExpiredLocked(now time.Time) {
	for s.expiry.Len() > 0 {
		next := s.expiry[0]
		if next.expiresAt.After(now) {
			return
		}
		heap.Pop(&s.expiry)
		element := s.items[next.key]
		if element == nil {
			continue
		}
		entry := element.Value.(memoryEntry)
		if entry.generation != next.generation || !entry.expiresAt.Equal(next.expiresAt) {
			continue
		}
		s.removeLocked(element)
	}
}

func (s *MemoryStore) compactExpiryLocked() {
	if s.expiry.Len() <= len(s.items)*2+defaultMemMaxEntries {
		return
	}
	// expiry 是惰性堆：更新会留下旧 generation，直到 prune 时才跳过。
	// 从 LRU 链表重建可以防止长进程积累无界的过期记录尾巴。
	expiry := make(expiryHeap, 0, len(s.items))
	for element := s.order.Front(); element != nil; element = element.Next() {
		entry := element.Value.(memoryEntry)
		expiry = append(expiry, expiryItem(entry))
	}
	heap.Init(&expiry)
	s.expiry = expiry
}

func (s *MemoryStore) removeLocked(element *list.Element) {
	if element == nil {
		return
	}
	entry := element.Value.(memoryEntry)
	delete(s.items, entry.key)
	s.order.Remove(element)
}

type expiryItem struct {
	key        string
	expiresAt  time.Time
	generation uint64
}

type expiryHeap []expiryItem

// Len 返回过期堆里的候选记录数量。
func (h expiryHeap) Len() int { return len(h) }

// Less 按过期时间和 generation 排序，保证最早失效的记录先被裁剪。
func (h expiryHeap) Less(i, j int) bool {
	if !h[i].expiresAt.Equal(h[j].expiresAt) {
		return h[i].expiresAt.Before(h[j].expiresAt)
	}
	return h[i].generation < h[j].generation
}

// Swap 交换两个堆节点。
func (h expiryHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

// Push 把新的过期候选记录加入堆。
func (h *expiryHeap) Push(x any) {
	*h = append(*h, x.(expiryItem))
}

// Pop 移除并返回堆尾节点，供 container/heap 维护最小堆。
func (h *expiryHeap) Pop() any {
	old := *h
	n := len(old)
	item := old[n-1]
	*h = old[:n-1]
	return item
}