backend/internal/application/clientkey/cache.go

package clientkey

import (
	"sync"
	"time"

	clientkeydomain "github.com/chenyme/grok2api/backend/internal/domain/clientkey"
)

const (
	keyTouchInterval       = time.Minute
	touchTrackerMaxEntries = 10000
	keyAuthCacheTTL        = time.Second
	keyAuthCacheMaxEntries = 10000
)

type cachedAuthKey struct {
	value     clientkeydomain.Key
	expiresAt time.Time
}

type authKeyCache struct {
	mu       sync.RWMutex
	byPrefix map[string]cachedAuthKey
}

func newAuthKeyCache() *authKeyCache {
	return &authKeyCache{byPrefix: make(map[string]cachedAuthKey)}
}

func (c *authKeyCache) get(prefix string, now time.Time) (clientkeydomain.Key, bool) {
	c.mu.RLock()
	entry, ok := c.byPrefix[prefix]
	c.mu.RUnlock()
	if !ok || !now.Before(entry.expiresAt) {
		if ok {
			c.mu.Lock()
			delete(c.byPrefix, prefix)
			c.mu.Unlock()
		}
		return clientkeydomain.Key{}, false
	}
	value := entry.value
	value.AllowedModels = append([]uint64(nil), entry.value.AllowedModels...)
	return value, true
}

func (c *authKeyCache) put(prefix string, value clientkeydomain.Key, now time.Time) {
	if prefix == "" || value.BillingLimitUSDTicks > 0 {
		return
	}
	value.EncryptedSecret = ""
	value.AllowedModels = append([]uint64(nil), value.AllowedModels...)
	c.mu.Lock()
	defer c.mu.Unlock()
	c.byPrefix[prefix] = cachedAuthKey{value: value, expiresAt: now.Add(keyAuthCacheTTL)}
	if len(c.byPrefix) <= keyAuthCacheMaxEntries {
		return
	}
	for candidate, entry := range c.byPrefix {
		if !now.Before(entry.expiresAt) {
			delete(c.byPrefix, candidate)
		}
	}
	for len(c.byPrefix) > keyAuthCacheMaxEntries {
		for candidate := range c.byPrefix {
			delete(c.byPrefix, candidate)
			break
		}
	}
}

func (c *authKeyCache) deleteID(id uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for prefix, entry := range c.byPrefix {
		if entry.value.ID == id {
			delete(c.byPrefix, prefix)
		}
	}
}

func (c *authKeyCache) deleteIDs(ids []uint64) {
	set := make(map[uint64]struct{}, len(ids))
	for _, id := range ids {
		set[id] = struct{}{}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	for prefix, entry := range c.byPrefix {
		if _, ok := set[entry.value.ID]; ok {
			delete(c.byPrefix, prefix)
		}
	}
}

func (c *authKeyCache) clear() {
	c.mu.Lock()
	clear(c.byPrefix)
	c.mu.Unlock()
}

// touchTracker 合并非关键的最近使用时间写入。
type touchTracker struct {
	mu          sync.Mutex
	lastTouched map[uint64]time.Time
}

func newTouchTracker() *touchTracker {
	return &touchTracker{lastTouched: make(map[uint64]time.Time)}
}

func (c *touchTracker) deleteID(id uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.lastTouched, id)
}

func (c *touchTracker) deleteIDs(ids []uint64) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, id := range ids {
		delete(c.lastTouched, id)
	}
}

func (c *touchTracker) shouldTouch(id uint64, now time.Time) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if last := c.lastTouched[id]; !last.IsZero() && now.Sub(last) < keyTouchInterval {
		return false
	}
	c.lastTouched[id] = now
	if len(c.lastTouched) > touchTrackerMaxEntries {
		var oldestID uint64
		var oldest time.Time
		for candidateID, touchedAt := range c.lastTouched {
			if oldestID == 0 || touchedAt.Before(oldest) {
				oldestID = candidateID
				oldest = touchedAt
			}
		}
		delete(c.lastTouched, oldestID)
	}
	return true
}