MCPX MCP Gateway

package utils import ( "fmt" "lunar/toolkit-core/clock" "sync" "time" "github.com/rs/zerolog/log" ) type Cache[K any, V any] interface { Has(key K) bool Get(key K) (V, bool) Set(key K, value V, ttlSec float64) error Del(key K) WithMaxCacheSize(calculateSizeFunc func(K, V) float64, maxCacheSize float64) } type MemoryCache[K comparable, V any] struct { cache map[K]ValueWrapper[V] mutex sync.RWMutex clock clock.Clock calculateCacheSize bool maxCacheSize float64 currentCacheSize float64 calculateSizeFunc func(key K, value V) float64 } type ValueWrapper[V any] struct { value V expirationTimeNano int64 } func (cache *MemoryCache[K, V]) Has(key K) bool { log.Trace().Msgf("Checking cache key %v", key) ensureCacheInitialized(cache) cache.mutex.RLock() valueWrapper, found := cache.cache[key] cache.mutex.RUnlock() if valueExpired(cache, key, valueWrapper.expirationTimeNano) { return false } if found { log.Trace().Msgf("Cache key %v found", key) } else { log.Trace().Msgf("Cache key %v not found", key) } return found } func (cache *MemoryCache[K, V]) Get(key K) (V, bool) { log.Trace().Msgf("Getting cache key %v", key) ensureCacheInitialized(cache) var value V cache.mutex.RLock() valueWrapper, found := cache.cache[key] cache.mutex.RUnlock() if !found { log.Trace().Msgf("Cache miss for key %v", key) return value, false } if valueExpired(cache, key, valueWrapper.expirationTimeNano) { return value, false } log.Trace().Msgf("Cache hit for key %v", key) return valueWrapper.value, true } func valueExpired[K comparable, V any]( cache *MemoryCache[K, V], key K, expirationTimeNano int64, ) bool { if cache.clock.Now().UnixNano() > expirationTimeNano { log.Trace().Msgf("Cache expired for key %v", key) return true } return false } func (cache *MemoryCache[K, V]) Set(key K, value V, ttlSec float64) error { log.Debug().Msgf( "Setting cache key %v to value %+v with ttl %v", key, value, ttlSec) ensureCacheInitialized(cache) itemSize := float64(0) if cache.calculateCacheSize && cache.calculateSizeFunc != nil { // We might miss here a momentary case of adding 2 messages when // there is enough place for only one. // The choice to put the check outside of the lock is intentional, // we are 'saving' value allocation and lock by checking the size first itemSize = cache.calculateSizeFunc(key, value) if cache.currentCacheSize+itemSize > cache.maxCacheSize { return fmt.Errorf( "cannot add item: max cache size would be exceeded."+ " Current cache size is %v", cache.currentCacheSize) } } ttlDuration := time.Duration(float64(time.Second) * ttlSec) expirationTimeNano := cache.clock.Now().UnixNano() + ttlDuration.Nanoseconds() cache.mutex.Lock() cache.cache[key] = ValueWrapper[V]{value, expirationTimeNano} if cache.calculateCacheSize { cache.currentCacheSize += itemSize } cache.mutex.Unlock() go func() { cache.clock.Sleep(ttlDuration) clearKey(cache, key) }() return nil } func (cache *MemoryCache[K, V]) Del(key K) { log.Trace().Msgf("Deleting cache key %v", key) ensureCacheInitialized(cache) clearKey(cache, key) } func (cache *MemoryCache[K, V]) WithMaxCacheSize( calculateSizeFunc func(K, V) float64, maxCacheSize float64, ) { cache.calculateCacheSize = true cache.calculateSizeFunc = calculateSizeFunc cache.maxCacheSize = maxCacheSize } // This function is used for printing the cache in debug logs // It is needed to avoid concurrent map read and map write errors func (cache *MemoryCache[K, V]) String() string { ensureCacheInitialized(cache) cache.mutex.RLock() defer cache.mutex.RUnlock() return fmt.Sprintf("%+v", cache.cache) } func NewMemoryCache[K comparable, V any](clock clock.Clock) *MemoryCache[K, V] { return &MemoryCache[K, V]{ cache: map[K]ValueWrapper[V]{}, mutex: sync.RWMutex{}, clock: clock, calculateCacheSize: false, maxCacheSize: 1, // TBD - get default value currentCacheSize: 0, calculateSizeFunc: nil, } } func ensureCacheInitialized[K comparable, V any](cache *MemoryCache[K, V]) { cache.mutex.Lock() if cache.cache == nil { cache.cache = map[K]ValueWrapper[V]{} } cache.mutex.Unlock() } func clearKey[K comparable, V any](cache *MemoryCache[K, V], key K) { cache.mutex.Lock() if cache.calculateCacheSize { valueWrapper, found := cache.cache[key] if found { cache.currentCacheSize -= cache.calculateSizeFunc(key, valueWrapper.value) } } delete(cache.cache, key) cache.mutex.Unlock() }