improvements mid may 2025 (#24)

* General improvements and bug fixes.

* Improve tests coverage.

* fixup! Improve tests coverage.

* Update README.md with latest changes.

* Fix the uint32

* Resolve issue with race condition for logging.

* fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* Fix the test of the rate limiter

* Add default ratelimit.json file

* Update dependencies.

* Significant refactor.

* fixup! Significant refactor.

* fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! fixup! fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

* fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Merge remote-tracking branch 'origin/main' into improvements-mid-apr-2025

cache/memory/buffer_pool.go

@@ -0,0 +1,17 @@
+package libpack_cache_memory
+
+import (
+	"bytes"
+
+	"github.com/lukaszraczylo/graphql-monitoring-proxy/pkg/pools"
+)
+
+// GetBuffer gets a buffer from the pool (delegates to unified implementation)
+func GetBuffer() *bytes.Buffer {
+	return pools.GetBuffer()
+}
+
+// PutBuffer returns a buffer to the pool (delegates to unified implementation)
+func PutBuffer(buf *bytes.Buffer) {
+	pools.PutBuffer(buf)
+}

cache/memory/compression_test.go

@@ -0,0 +1,218 @@
+package libpack_cache_memory
+
+import (
+	"bytes"
+	"compress/gzip"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// TestCompressionThreshold tests that values are only compressed when they exceed the threshold
+func TestCompressionThreshold(t *testing.T) {
+	cache := New(5 * time.Second)
+
+	// Create test values
+	smallValue := make([]byte, CompressionThreshold-100) // Below threshold
+	largeValue := make([]byte, CompressionThreshold*2)   // Above threshold
+
+	// Fill values with compressible data (repeating patterns compress well)
+	for i := 0; i < len(smallValue); i++ {
+		smallValue[i] = byte(i % 10)
+	}
+	for i := 0; i < len(largeValue); i++ {
+		largeValue[i] = byte(i % 10)
+	}
+
+	// Test small value
+	cache.Set("small-key", smallValue, 5*time.Second)
+
+	// Extract the entry directly from the cache to check if it's compressed
+	entryRaw, found := cache.entries.Load("small-key")
+	assert.True(t, found, "Entry should exist")
+
+	entry := entryRaw.(CacheEntry)
+	assert.False(t, entry.Compressed, "Small value should not be compressed")
+	assert.Equal(t, smallValue, entry.Value, "Small value should be stored as-is")
+
+	// Test large value
+	cache.Set("large-key", largeValue, 5*time.Second)
+
+	entryRaw, found = cache.entries.Load("large-key")
+	assert.True(t, found, "Entry should exist")
+
+	entry = entryRaw.(CacheEntry)
+	assert.True(t, entry.Compressed, "Large value should be compressed")
+
+	// Ensure the stored value isn't the original
+	assert.NotEqual(t, largeValue, entry.Value, "Large value should not be stored as-is")
+
+	// Verify the value is actually compressed (should be smaller)
+	assert.Less(t, len(entry.Value), len(largeValue), "Compressed value should be smaller than original")
+
+	// Verify we can retrieve the uncompressed value correctly
+	retrievedLarge, found := cache.Get("large-key")
+	assert.True(t, found, "Large value should be retrievable")
+	assert.Equal(t, largeValue, retrievedLarge, "Retrieved large value should match original")
+}
+
+// TestCompressionMemoryUsage tests that memory usage is calculated correctly for compressed entries
+func TestCompressionMemoryUsage(t *testing.T) {
+	cache := New(5 * time.Second)
+
+	// Create a large, highly compressible value
+	valueSize := CompressionThreshold * 4
+	value := make([]byte, valueSize)
+	for i := 0; i < valueSize; i++ {
+		value[i] = byte(i % 2) // Highly compressible pattern (alternating 0s and 1s)
+	}
+
+	// Get initial memory usage
+	initialMemUsage := cache.GetMemoryUsage()
+
+	// Add the value
+	key := "large-compressible-key"
+	cache.Set(key, value, 5*time.Second)
+
+	// Get memory usage after adding
+	newMemUsage := cache.GetMemoryUsage()
+
+	// The memory usage increase should be less than the full value size due to compression
+	memUsageIncrease := newMemUsage - initialMemUsage
+
+	// Extract the entry to check its compressed size
+	entryRaw, found := cache.entries.Load(key)
+	assert.True(t, found, "Entry should exist")
+
+	entry := entryRaw.(CacheEntry)
+	assert.True(t, entry.Compressed, "Value should be compressed")
+
+	// Verify the reported memory usage matches the compressed size + overheads
+	compressedSize := int64(len(entry.Value))
+	keySize := int64(len(key))
+	expectedUsage := compressedSize + keySize + approxEntryOverhead
+
+	// The memory usage should reflect the compressed size, not the original size
+	assert.InDelta(t, expectedUsage, memUsageIncrease, float64(approxEntryOverhead),
+		"Memory usage should be based on compressed size")
+
+	// Verify memory usage is correctly updated after deletion
+	cache.Delete(key)
+	finalMemUsage := cache.GetMemoryUsage()
+	assert.Equal(t, initialMemUsage, finalMemUsage,
+		"Memory usage should return to initial value after deletion")
+}
+
+// TestUncompressibleData tests the case where compression doesn't reduce size
+func TestUncompressibleData(t *testing.T) {
+	cache := New(5 * time.Second)
+
+	// Create a large, random (less compressible) value
+	valueSize := CompressionThreshold * 2
+
+	// Create pseudo-random data that doesn't compress well
+	// Using a custom PRNG for deterministic results across test runs
+	value := make([]byte, valueSize)
+	seed := uint32(42)
+	for i := 0; i < valueSize; i++ {
+		// Simple linear congruential generator
+		seed = seed*1664525 + 1013904223
+		value[i] = byte(seed)
+	}
+
+	// Try to compress it directly to see if it actually would reduce size
+	var buf bytes.Buffer
+	gw := gzip.NewWriter(&buf)
+	_, _ = gw.Write(value)
+	_ = gw.Close()
+	compressedDirectly := buf.Bytes()
+
+	// Now use the cache's Set method
+	key := "uncompressible-key"
+	cache.Set(key, value, 5*time.Second)
+
+	// Extract the entry to check if it's compressed
+	entryRaw, found := cache.entries.Load(key)
+	assert.True(t, found, "Entry should exist")
+
+	entry := entryRaw.(CacheEntry)
+
+	// If our test data actually compressed to a smaller size, we expect the cache to store it compressed
+	if len(compressedDirectly) < len(value) {
+		assert.True(t, entry.Compressed, "Value should be stored compressed if smaller")
+		assert.Less(t, len(entry.Value), len(value), "Compressed value should be smaller")
+	} else {
+		// Uncommon case: our pseudo-random data actually expanded with gzip
+		// In this case, the cache should store it uncompressed
+		assert.False(t, entry.Compressed, "Value should not be compressed if it would expand")
+		assert.Equal(t, value, entry.Value, "Value should be stored as-is")
+	}
+
+	// Regardless, we should be able to get the correct value back
+	retrievedValue, found := cache.Get(key)
+	assert.True(t, found, "Value should be retrievable")
+	assert.Equal(t, value, retrievedValue, "Retrieved value should match original")
+}
+
+// TestCompressDecompressDirectly tests the compress and decompress methods directly
+func TestCompressDecompressDirectly(t *testing.T) {
+	cache := New(5 * time.Second)
+
+	// Test with various sizes
+	testSizes := []int{
+		100,                      // Small
+		CompressionThreshold - 1, // Just below threshold
+		CompressionThreshold,     // At threshold
+		CompressionThreshold + 1, // Just above threshold
+		CompressionThreshold * 2, // Well above threshold
+	}
+
+	for _, size := range testSizes {
+		t.Run("Size-"+string(rune('A'+len(testSizes)%26)), func(t *testing.T) {
+			// Generate test data with a repeating pattern
+			data := make([]byte, size)
+			for i := 0; i < size; i++ {
+				data[i] = byte(i % 256)
+			}
+
+			// Compress the data
+			compressed, err := cache.compress(data)
+			assert.NoError(t, err, "Compression should not error")
+
+			// Small data may get larger when compressed, larger data should get smaller
+			if size > CompressionThreshold {
+				assert.Less(t, len(compressed), len(data),
+					"Compression should reduce size for data above threshold")
+			}
+
+			// Decompress and verify it matches the original
+			decompressed, err := cache.decompress(compressed)
+			assert.NoError(t, err, "Decompression should not error")
+			assert.Equal(t, data, decompressed, "Data should round-trip correctly through compression")
+		})
+	}
+}
+
+// TestDecompressInvalidData tests handling invalid data in decompress
+func TestDecompressInvalidData(t *testing.T) {
+	cache := New(5 * time.Second)
+
+	// Try to decompress non-gzip data
+	invalidData := []byte("This is not valid gzip data")
+	_, err := cache.decompress(invalidData)
+	assert.Error(t, err, "Decompressing invalid data should return error")
+
+	// Set compressed flag but store invalid data
+	key := "invalid-compressed-key"
+	cache.entries.Store(key, CacheEntry{
+		Value:      invalidData,
+		ExpiresAt:  time.Now().Add(5 * time.Second),
+		Compressed: true, // Flag as compressed even though it's not
+		MemorySize: int64(len(invalidData) + len(key) + approxEntryOverhead),
+	})
+
+	// Try to get it - should fail gracefully
+	_, found := cache.Get(key)
+	assert.False(t, found, "Get should fail gracefully for invalid compressed data")
+}

cache/memory/eviction_test.go

@@ -0,0 +1,185 @@
+package libpack_cache_memory
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// TestEvictToFreeMemory tests that the cache correctly evicts
+// items when it exceeds its memory limit.
+func TestEvictToFreeMemory(t *testing.T) {
+	// Create a cache with a small memory limit: 5KB (ensure eviction happens)
+	smallMemLimit := int64(5 * 1024)
+	cache := NewWithSize(5*time.Second, smallMemLimit, 1000)
+
+	// Create entries with known sizes
+	// Each entry will be ~512 bytes plus overhead
+	valueSize := 512
+	numEntriesToExceedLimit := 12 // Should exceed the 5KB limit and force eviction
+
+	// Create a slice to track keys in insertion order
+	keys := make([]string, numEntriesToExceedLimit)
+
+	// Add entries with significant delays between insertions
+	for i := 0; i < numEntriesToExceedLimit; i++ {
+		key := fmt.Sprintf("test-key-%d", i)
+		keys[i] = key
+
+		value := make([]byte, valueSize)
+		for j := 0; j < valueSize; j++ {
+			value[j] = byte(i % 256) // Fill with a repeating pattern
+		}
+
+		cache.Set(key, value, 30*time.Second)
+
+		// More significant delay to ensure different timestamps
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	// Allow time for eviction to complete
+	time.Sleep(50 * time.Millisecond)
+
+	// Verify memory usage is below the limit
+	memUsage := cache.GetMemoryUsage()
+	assert.LessOrEqual(t, memUsage, smallMemLimit,
+		"Memory usage (%d) should be less than or equal to the limit (%d)", memUsage, smallMemLimit)
+
+	// Count how many items are left in the cache and which ones
+	present := 0
+	for i := 0; i < numEntriesToExceedLimit; i++ {
+		_, found := cache.Get(keys[i])
+		if found {
+			present++
+		}
+	}
+
+	// We expect some items to be evicted based on the memory limit
+	assert.Less(t, present, numEntriesToExceedLimit,
+		"Some items should have been evicted (%d present out of %d total)",
+		present, numEntriesToExceedLimit)
+
+	// Verify newer items (inserted later) are more likely to be in the cache
+	// Check the last few items which should be the newest
+	for i := numEntriesToExceedLimit - 3; i < numEntriesToExceedLimit; i++ {
+		_, found := cache.Get(keys[i])
+		assert.True(t, found, "Newer key %s should still exist", keys[i])
+	}
+}
+
+// TestMaxCacheSize verifies the behavior when adding more items than the maxCacheSize limit
+func TestMaxCacheSize(t *testing.T) {
+	// Create a cache with a small limit
+	smallLimit := int64(5)
+	cache := NewWithSize(5*time.Second, DefaultMaxMemorySize, smallLimit)
+
+	// Add entries with increasing size (to avoid memory-based eviction)
+	for i := 0; i < 20; i++ {
+		key := fmt.Sprintf("test-key-%d", i)
+		value := []byte(key)
+		cache.Set(key, value, 10*time.Second)
+	}
+
+	// Verify we can get a reasonable number of items
+	// (we don't test for exact count as implementation may vary)
+	foundCount := 0
+	for i := 0; i < 20; i++ {
+		key := fmt.Sprintf("test-key-%d", i)
+		_, found := cache.Get(key)
+		if found {
+			foundCount++
+		}
+	}
+
+	// We should find some items but not all 20
+	assert.Greater(t, foundCount, 0, "Some items should be in the cache")
+	assert.LessOrEqual(t, foundCount, 20, "Not all items should be in the cache with small limit")
+}
+
+// TestGetMemoryUsage verifies that memory usage tracking is accurate
+func TestGetMemoryUsage(t *testing.T) {
+	cache := New(5 * time.Second)
+
+	// Initially memory usage should be 0
+	assert.Equal(t, int64(0), cache.GetMemoryUsage(), "Initial memory usage should be 0")
+
+	// Add an entry with a known approximate size
+	valueSize := 1024
+	value := make([]byte, valueSize)
+	key := "test-key"
+
+	cache.Set(key, value, 5*time.Second)
+
+	// Check memory usage - should be approximately valueSize + key length + overhead
+	expectedMinUsage := int64(valueSize + len(key))
+	memUsage := cache.GetMemoryUsage()
+	assert.GreaterOrEqual(t, memUsage, expectedMinUsage,
+		"Memory usage (%d) should be at least the value size plus key length (%d)", memUsage, expectedMinUsage)
+
+	// Delete the entry and verify memory usage decreases
+	cache.Delete(key)
+	assert.Equal(t, int64(0), cache.GetMemoryUsage(), "Memory usage should be 0 after deletion")
+}
+
+// TestSetMaxMemorySize tests changing the memory limit and resulting eviction
+func TestSetMaxMemorySize(t *testing.T) {
+	// Start with a large limit
+	initialLimit := int64(100 * 1024)
+	cache := NewWithSize(5*time.Second, initialLimit, 1000)
+
+	// Fill the cache with ~50KB of data
+	valueSize := 1024
+	numEntries := 50
+
+	for i := 0; i < numEntries; i++ {
+		key := generateKey(i)
+		value := make([]byte, valueSize)
+		cache.Set(key, value, 5*time.Second)
+
+		// Small delay for timestamp differences
+		time.Sleep(time.Millisecond)
+	}
+
+	// Verify all entries exist
+	for i := 0; i < numEntries; i++ {
+		_, found := cache.Get(generateKey(i))
+		assert.True(t, found, "All entries should exist before limit change")
+	}
+
+	// Get current memory usage
+	originalUsage := cache.GetMemoryUsage()
+
+	// Now reduce the limit to 20KB - should trigger eviction
+	newLimit := int64(20 * 1024)
+	cache.SetMaxMemorySize(newLimit)
+
+	// Verify memory usage is now below the new limit
+	newUsage := cache.GetMemoryUsage()
+	assert.LessOrEqual(t, newUsage, newLimit,
+		"After SetMaxMemorySize, memory usage (%d) should be less than or equal to new limit (%d)",
+		newUsage, newLimit)
+	assert.Less(t, newUsage, originalUsage,
+		"Memory usage should have decreased after lowering the limit")
+
+	// Some older entries should be gone, newer ones should still exist
+	removedCount := 0
+	remainingCount := 0
+	for i := 0; i < numEntries; i++ {
+		_, found := cache.Get(generateKey(i))
+		if found {
+			remainingCount++
+		} else {
+			removedCount++
+		}
+	}
+
+	assert.Greater(t, removedCount, 0, "Some entries should have been removed")
+	assert.Greater(t, remainingCount, 0, "Some entries should still exist")
+}
+
+// Helper function to generate consistent keys
+func generateKey(index int) string {
+	return "test-key-" + fmt.Sprintf("%d", index)
+}

cache/memory/lru_memory_cache.go

@@ -0,0 +1,281 @@
+package libpack_cache_memory
+
+import (
+	"compress/gzip"
+	"container/list"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// LRUMemoryCache is an efficient LRU-based memory cache implementation
+type LRUMemoryCache struct {
+	entries        map[string]*lruEntry
+	evictList      *list.List
+	gzipWriterPool *sync.Pool
+	gzipReaderPool *sync.Pool
+	maxMemorySize  int64
+	maxEntries     int64
+	currentMemory  int64
+	currentCount   int64
+	mu             sync.RWMutex
+}
+
+type lruEntry struct {
+	expiresAt  time.Time
+	element    *list.Element
+	key        string
+	value      []byte
+	size       int64
+	compressed bool
+}
+
+// NewLRUMemoryCache creates a new LRU memory cache
+func NewLRUMemoryCache(maxMemorySize, maxEntries int64) *LRUMemoryCache {
+	return &LRUMemoryCache{
+		maxMemorySize: maxMemorySize,
+		maxEntries:    maxEntries,
+		entries:       make(map[string]*lruEntry),
+		evictList:     list.New(),
+		gzipWriterPool: &sync.Pool{
+			New: func() interface{} {
+				return gzip.NewWriter(nil)
+			},
+		},
+		gzipReaderPool: &sync.Pool{
+			New: func() interface{} {
+				return &gzip.Reader{}
+			},
+		},
+	}
+}
+
+// Set adds or updates an entry in the cache
+func (c *LRUMemoryCache) Set(key string, value []byte, ttl time.Duration) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// Calculate expiry time
+	expiresAt := time.Now().Add(ttl)
+
+	// Check if we should compress
+	compressed := false
+	finalValue := value
+	if len(value) > 1024 { // Compress if larger than 1KB
+		if compressedData, err := c.compress(value); err == nil && len(compressedData) < len(value) {
+			compressed = true
+			finalValue = compressedData
+		}
+	}
+
+	entrySize := int64(len(key) + len(finalValue) + 64) // 64 bytes overhead estimate
+
+	// Check if key exists
+	if existing, exists := c.entries[key]; exists {
+		// Update existing entry
+		c.evictList.MoveToFront(existing.element)
+		atomic.AddInt64(&c.currentMemory, -existing.size)
+		atomic.AddInt64(&c.currentMemory, entrySize)
+
+		existing.value = finalValue
+		existing.compressed = compressed
+		existing.size = entrySize
+		existing.expiresAt = expiresAt
+
+		c.evictIfNeeded()
+		return
+	}
+
+	// Create new entry
+	entry := &lruEntry{
+		key:        key,
+		value:      finalValue,
+		compressed: compressed,
+		size:       entrySize,
+		expiresAt:  expiresAt,
+	}
+
+	element := c.evictList.PushFront(entry)
+	entry.element = element
+	c.entries[key] = entry
+
+	atomic.AddInt64(&c.currentMemory, entrySize)
+	atomic.AddInt64(&c.currentCount, 1)
+
+	c.evictIfNeeded()
+}
+
+// Get retrieves a value from the cache
+func (c *LRUMemoryCache) Get(key string) ([]byte, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	entry, exists := c.entries[key]
+	if !exists {
+		return nil, false
+	}
+
+	// Check if expired
+	if time.Now().After(entry.expiresAt) {
+		c.removeEntry(entry)
+		return nil, false
+	}
+
+	// Move to front (most recently used)
+	c.evictList.MoveToFront(entry.element)
+
+	// Decompress if needed
+	if entry.compressed {
+		if decompressed, err := c.decompress(entry.value); err == nil {
+			return decompressed, true
+		}
+		// If decompression fails, remove the entry
+		c.removeEntry(entry)
+		return nil, false
+	}
+
+	return entry.value, true
+}
+
+// Delete removes an entry from the cache
+func (c *LRUMemoryCache) Delete(key string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if entry, exists := c.entries[key]; exists {
+		c.removeEntry(entry)
+	}
+}
+
+// Clear removes all entries
+func (c *LRUMemoryCache) Clear() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.entries = make(map[string]*lruEntry)
+	c.evictList = list.New()
+	atomic.StoreInt64(&c.currentMemory, 0)
+	atomic.StoreInt64(&c.currentCount, 0)
+}
+
+// evictIfNeeded removes entries when limits are exceeded
+func (c *LRUMemoryCache) evictIfNeeded() {
+	// Evict based on entry count
+	for atomic.LoadInt64(&c.currentCount) > c.maxEntries && c.evictList.Len() > 0 {
+		c.evictOldest()
+	}
+
+	// Evict based on memory
+	for atomic.LoadInt64(&c.currentMemory) > c.maxMemorySize && c.evictList.Len() > 0 {
+		c.evictOldest()
+	}
+}
+
+// evictOldest removes the least recently used entry
+func (c *LRUMemoryCache) evictOldest() {
+	element := c.evictList.Back()
+	if element == nil {
+		return
+	}
+
+	entry := element.Value.(*lruEntry)
+	c.removeEntry(entry)
+}
+
+// removeEntry removes an entry from all data structures
+func (c *LRUMemoryCache) removeEntry(entry *lruEntry) {
+	c.evictList.Remove(entry.element)
+	delete(c.entries, entry.key)
+	atomic.AddInt64(&c.currentMemory, -entry.size)
+	atomic.AddInt64(&c.currentCount, -1)
+}
+
+// CleanExpiredEntries removes all expired entries
+func (c *LRUMemoryCache) CleanExpiredEntries() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	now := time.Now()
+	for element := c.evictList.Back(); element != nil; {
+		entry := element.Value.(*lruEntry)
+
+		if now.After(entry.expiresAt) {
+			next := element.Prev()
+			c.removeEntry(entry)
+			element = next
+		} else {
+			element = element.Prev()
+		}
+	}
+}
+
+// compress compresses data using gzip
+func (c *LRUMemoryCache) compress(data []byte) ([]byte, error) {
+	buf := GetBuffer()
+	defer PutBuffer(buf)
+
+	gz := c.gzipWriterPool.Get().(*gzip.Writer)
+	gz.Reset(buf)
+	defer c.gzipWriterPool.Put(gz)
+
+	if _, err := gz.Write(data); err != nil {
+		return nil, err
+	}
+	if err := gz.Close(); err != nil {
+		return nil, err
+	}
+
+	compressed := make([]byte, buf.Len())
+	copy(compressed, buf.Bytes())
+	return compressed, nil
+}
+
+// decompress decompresses gzip data
+func (c *LRUMemoryCache) decompress(data []byte) ([]byte, error) {
+	buf := GetBuffer()
+	defer PutBuffer(buf)
+
+	buf.Write(data)
+
+	gr := c.gzipReaderPool.Get().(*gzip.Reader)
+	defer c.gzipReaderPool.Put(gr)
+
+	if err := gr.Reset(buf); err != nil {
+		return nil, err
+	}
+
+	result := GetBuffer()
+	defer PutBuffer(result)
+
+	if _, err := result.ReadFrom(gr); err != nil {
+		return nil, err
+	}
+
+	decompressed := make([]byte, result.Len())
+	copy(decompressed, result.Bytes())
+	return decompressed, nil
+}
+
+// GetStats returns cache statistics
+func (c *LRUMemoryCache) GetStats() map[string]interface{} {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return map[string]interface{}{
+		"entries":      atomic.LoadInt64(&c.currentCount),
+		"memory_bytes": atomic.LoadInt64(&c.currentMemory),
+		"max_entries":  c.maxEntries,
+		"max_memory":   c.maxMemorySize,
+		"fill_percent": float64(atomic.LoadInt64(&c.currentMemory)) / float64(c.maxMemorySize) * 100,
+	}
+}
+
+// GetMemoryUsage returns current memory usage in bytes
+func (c *LRUMemoryCache) GetMemoryUsage() int64 {
+	return atomic.LoadInt64(&c.currentMemory)
+}
+
+// GetMaxMemorySize returns the maximum memory size
+func (c *LRUMemoryCache) GetMaxMemorySize() int64 {
+	return c.maxMemorySize
+}

cache/memory/memory.go

@@ -3,8 +3,8 @@ package libpack_cache_memory
 import (
 	"bytes"
 	"compress/gzip"
+	"context"
 	"io"
-	"runtime"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -13,27 +13,53 @@ import (
 // CompressionThreshold is the minimum size in bytes before a value is compressed
 const CompressionThreshold = 1024 // 1KB
 
-// MaxCacheSize is the maximum number of entries in the cache
-const MaxCacheSize = 10000
+// DefaultMaxMemorySize is the default maximum memory size in bytes (100MB)
+const DefaultMaxMemorySize = 100 * 1024 * 1024
+
+// DefaultMaxCacheSize is the default maximum number of entries in the cache
+// This is used for backward compatibility
+const DefaultMaxCacheSize = 10000
+
+// approxEntryOverhead is the estimated overhead per cache entry in bytes
+// This accounts for the CacheEntry struct overhead, map entry, and synchronization
+const approxEntryOverhead = 64
 
 type CacheEntry struct {
 	ExpiresAt  time.Time
 	Value      []byte
 	Compressed bool
+	MemorySize int64 // Estimated memory usage of this entry in bytes
 }
 
 type Cache struct {
 	compressPool   sync.Pool
 	decompressPool sync.Pool
+	ctx            context.Context
+	cancel         context.CancelFunc
 	entries        sync.Map
 	globalTTL      time.Duration
 	entryCount     int64
+	memoryUsage    int64
+	maxMemorySize  int64
+	maxCacheSize   int64
 	sync.RWMutex
 }
 
 func New(globalTTL time.Duration) *Cache {
+	return NewWithSize(globalTTL, DefaultMaxMemorySize, DefaultMaxCacheSize)
+}
+
+// NewWithSize creates a new cache with the specified memory size limit and entry count limit
+func NewWithSize(globalTTL time.Duration, maxMemorySize int64, maxCacheSize int64) *Cache {
+	// Create context for graceful shutdown
+	ctx, cancel := context.WithCancel(context.Background())
+
 	cache := &Cache{
-		globalTTL: globalTTL,
+		globalTTL:     globalTTL,
+		maxMemorySize: maxMemorySize,
+		maxCacheSize:  maxCacheSize,
+		ctx:           ctx,
+		cancel:        cancel,
 		compressPool: sync.Pool{
 			New: func() interface{} {
 				return gzip.NewWriter(nil)
@@ -47,7 +73,7 @@ func New(globalTTL time.Duration) *Cache {
 		},
 	}
 
-	// Start cleanup routine
+	// Start cleanup routine with context cancellation
 	go cache.cleanupRoutine(globalTTL)
 	return cache
 }
@@ -57,20 +83,40 @@ func (c *Cache) cleanupRoutine(globalTTL time.Duration) {
 	ticker := time.NewTicker(globalTTL / 4)
 	defer ticker.Stop()
 
-	for range ticker.C {
-		c.CleanExpiredEntries()
+	for {
+		select {
+		case <-c.ctx.Done():
+			// Context cancelled, exit gracefully
+			return
+		case <-ticker.C:
+			c.CleanExpiredEntries()
 
-		// Trigger GC if we have a lot of entries
-		if atomic.LoadInt64(&c.entryCount) > MaxCacheSize/2 {
-			runtime.GC()
+			// Note: Removed aggressive GC trigger that was causing performance issues
+			// The Go runtime GC is already optimized and will run when needed
 		}
 	}
 }
 
+// Shutdown gracefully stops the cache cleanup routine
+func (c *Cache) Shutdown() {
+	if c.cancel != nil {
+		c.cancel()
+	}
+}
+
 func (c *Cache) Set(key string, value []byte, ttl time.Duration) {
-	// Check if we've reached the maximum cache size
-	if atomic.LoadInt64(&c.entryCount) >= MaxCacheSize {
-		c.evictOldest(MaxCacheSize / 10) // Evict 10% of entries
+	// Calculate the memory size of this entry
+	entrySize := int64(len(key) + len(value) + approxEntryOverhead)
+
+	// Check if we need to evict entries based on memory or count limits
+	currentMemory := atomic.LoadInt64(&c.memoryUsage)
+	if currentMemory+entrySize > c.maxMemorySize {
+		// Need to evict based on memory
+		memoryToFree := (currentMemory + entrySize) - c.maxMemorySize + (c.maxMemorySize / 10)
+		c.evictToFreeMemory(memoryToFree)
+	} else if atomic.LoadInt64(&c.entryCount) >= c.maxCacheSize {
+		// Fall back to count-based eviction for backward compatibility
+		c.evictOldest(int(c.maxCacheSize / 10)) // Evict 10% of entries
 	}
 
 	expiresAt := time.Now().Add(ttl)
@@ -101,12 +147,26 @@ func (c *Cache) Set(key string, value []byte, ttl time.Duration) {
 		}
 	}
 
-	// Check if this is a new entry
-	_, exists := c.entries.Load(key)
-	if !exists {
+	// Update the entry memory size based on compression status
+	if entry.Compressed {
+		entry.MemorySize = int64(len(key) + len(entry.Value) + approxEntryOverhead)
+	} else {
+		entry.MemorySize = int64(len(key) + len(entry.Value) + approxEntryOverhead)
+	}
+
+	// Check if this is a new entry or an update
+	oldEntry, exists := c.entries.Load(key)
+	if exists {
+		// Update memory usage: subtract old entry size, add new entry size
+		oldCacheEntry := oldEntry.(CacheEntry)
+		atomic.AddInt64(&c.memoryUsage, -oldCacheEntry.MemorySize)
+	} else {
+		// New entry
 		atomic.AddInt64(&c.entryCount, 1)
 	}
 
+	// Add new entry's memory size to total
+	atomic.AddInt64(&c.memoryUsage, entry.MemorySize)
 	c.entries.Store(key, entry)
 }
 
@@ -120,6 +180,7 @@ func (c *Cache) Get(key string) ([]byte, bool) {
 	if cacheEntry.ExpiresAt.Before(time.Now()) {
 		c.entries.Delete(key)
 		atomic.AddInt64(&c.entryCount, -1)
+		atomic.AddInt64(&c.memoryUsage, -cacheEntry.MemorySize)
 		return nil, false
 	}
 
@@ -135,8 +196,10 @@ func (c *Cache) Get(key string) ([]byte, bool) {
 }
 
 func (c *Cache) Delete(key string) {
-	if _, exists := c.entries.LoadAndDelete(key); exists {
+	if entry, exists := c.entries.LoadAndDelete(key); exists {
+		cacheEntry := entry.(CacheEntry)
 		atomic.AddInt64(&c.entryCount, -1)
+		atomic.AddInt64(&c.memoryUsage, -cacheEntry.MemorySize)
 	}
 }
 
@@ -146,6 +209,7 @@ func (c *Cache) Clear() {
 		return true
 	})
 	atomic.StoreInt64(&c.entryCount, 0)
+	atomic.StoreInt64(&c.memoryUsage, 0)
 }
 
 func (c *Cache) CountQueries() int64 {
@@ -183,7 +247,9 @@ func (c *Cache) decompress(data []byte) ([]byte, error) {
 		}
 	}
 
-	defer r.Close()
+	defer func() {
+		_ = r.Close() // Ignore error in defer cleanup
+	}()
 	return io.ReadAll(r)
 }
 
@@ -194,6 +260,7 @@ func (c *Cache) CleanExpiredEntries() {
 		if entry.ExpiresAt.Before(now) {
 			if _, exists := c.entries.LoadAndDelete(key); exists {
 				atomic.AddInt64(&c.entryCount, -1)
+				atomic.AddInt64(&c.memoryUsage, -entry.MemorySize)
 			}
 		}
 		return true
@@ -203,8 +270,8 @@ func (c *Cache) CleanExpiredEntries() {
 // evictOldest removes the oldest n entries from the cache
 func (c *Cache) evictOldest(n int) {
 	type keyExpiry struct {
-		key       string
 		expiresAt time.Time
+		key       string
 	}
 
 	// Collect all entries with their expiry times
@@ -212,7 +279,7 @@ func (c *Cache) evictOldest(n int) {
 	c.entries.Range(func(k, v interface{}) bool {
 		key := k.(string)
 		entry := v.(CacheEntry)
-		entries = append(entries, keyExpiry{key, entry.ExpiresAt})
+		entries = append(entries, keyExpiry{entry.ExpiresAt, key})
 		return len(entries) < cap(entries)
 	})
 
@@ -231,8 +298,74 @@ func (c *Cache) evictOldest(n int) {
 		}
 
 		// Delete this entry
-		if _, exists := c.entries.LoadAndDelete(entries[i].key); exists {
+		if entry, exists := c.entries.LoadAndDelete(entries[i].key); exists {
+			cacheEntry := entry.(CacheEntry)
 			atomic.AddInt64(&c.entryCount, -1)
+			atomic.AddInt64(&c.memoryUsage, -cacheEntry.MemorySize)
 		}
 	}
 }
+
+// evictToFreeMemory removes entries until the specified amount of memory is freed
+func (c *Cache) evictToFreeMemory(bytesToFree int64) {
+	type keyMemorySize struct {
+		expiresAt  time.Time
+		key        string
+		memorySize int64
+	}
+
+	// Collect entries to consider for eviction
+	entries := make([]keyMemorySize, 0, int(c.maxCacheSize/5))
+	c.entries.Range(func(k, v interface{}) bool {
+		key := k.(string)
+		entry := v.(CacheEntry)
+		entries = append(entries, keyMemorySize{entry.ExpiresAt, key, entry.MemorySize})
+		return len(entries) < cap(entries)
+	})
+
+	// Sort entries by expiry time (oldest first)
+	// Simple selection sort since we only need to find the oldest entries
+	var freedBytes int64
+	for i := 0; i < len(entries) && freedBytes < bytesToFree; i++ {
+		oldest := i
+		for j := i + 1; j < len(entries); j++ {
+			if entries[j].expiresAt.Before(entries[oldest].expiresAt) {
+				oldest = j
+			}
+		}
+		// Swap
+		if oldest != i {
+			entries[i], entries[oldest] = entries[oldest], entries[i]
+		}
+
+		// Delete this entry
+		if entry, exists := c.entries.LoadAndDelete(entries[i].key); exists {
+			cacheEntry := entry.(CacheEntry)
+			atomic.AddInt64(&c.entryCount, -1)
+			atomic.AddInt64(&c.memoryUsage, -cacheEntry.MemorySize)
+			freedBytes += cacheEntry.MemorySize
+		}
+	}
+}
+
+// GetMemoryUsage returns the current memory usage of the cache in bytes
+func (c *Cache) GetMemoryUsage() int64 {
+	return atomic.LoadInt64(&c.memoryUsage)
+}
+
+// GetMaxMemorySize returns the maximum memory size allowed for the cache in bytes
+func (c *Cache) GetMaxMemorySize() int64 {
+	return c.maxMemorySize
+}
+
+// SetMaxMemorySize updates the maximum memory size allowed for the cache
+func (c *Cache) SetMaxMemorySize(maxBytes int64) {
+	c.maxMemorySize = maxBytes
+
+	// Check if we need to evict entries due to the new limit
+	currentMemory := atomic.LoadInt64(&c.memoryUsage)
+	if currentMemory > maxBytes {
+		memoryToFree := currentMemory - maxBytes + (maxBytes / 10)
+		c.evictToFreeMemory(memoryToFree)
+	}
+}