feat(cache): add 100MB MaxSize and sets Discard MaxBytes limit to NATS KV cache buckets (#3014)

Signed-off-by: Miguel Martinez Trivino <migmartri@gmail.com>
Signed-off-by: Miguel Martinez Trivino <miguel@chainloop.dev>

Author: migmartri

pkg/cache/attestationbundle/attestationbundle.go

@@ -27,6 +27,7 @@ import (
 
 const (
 	ttl         = 5 * 24 * time.Hour
+	maxBytes    = 100 * 1024 * 1024 // 100 MB
 	bucket      = "chainloop-attestation-bundles"
 	description = "Cache for attestation bundles"
 )
@@ -40,6 +41,7 @@ type Cache struct {
 func New(ctx context.Context, rc *natsconn.ReloadableConnection, logger log.Logger) (*Cache, error) {
 	opts := []cache.Option{
 		cache.WithTTL(ttl),
+		cache.WithMaxBytes(maxBytes),
 		cache.WithDescription(description),
 	}
 

pkg/cache/cache.go

@@ -40,13 +40,9 @@ type Logger interface {
 	Errorw(keyvals ...any)
 }
 
-// defaultMaxSize is a sensible upper bound on in-memory cache entries
-// to prevent unbounded growth. 0 means no LRU eviction (TTL-only).
-const defaultMaxSize = 1000
-
 type config struct {
 	ttl         time.Duration
-	maxSize     int
+	maxBytes    int64
 	logger      Logger
 	natsConn    *nats.Conn
 	bucketName  string
@@ -75,6 +71,12 @@ func WithNATS(conn *nats.Conn, bucketName string) Option {
 	}
 }
 
+// WithMaxBytes sets the maximum total size (in bytes) for the NATS KV bucket.
+// When the limit is reached, NATS discards the oldest entries. Ignored for in-memory backend.
+func WithMaxBytes(n int64) Option {
+	return func(c *config) { c.maxBytes = n }
+}
+
 // WithDescription sets the NATS KV bucket description. Ignored for in-memory backend.
 func WithDescription(desc string) Option {
 	return func(c *config) { c.description = desc }
@@ -108,10 +110,6 @@ func New[T any](opts ...Option) (Cache[T], error) {
 		return newNATSKV[T](cfg)
 	}
 
-	if cfg.maxSize == 0 {
-		cfg.maxSize = defaultMaxSize
-	}
-
 	return newMemoryCache[T](cfg), nil
 }
 

pkg/cache/memory.go

@@ -21,15 +21,19 @@ import (
 	"github.com/hashicorp/golang-lru/v2/expirable"
 )
 
+// defaultMaxSize is a sensible upper bound on in-memory cache entries
+// to prevent unbounded growth. 0 means no LRU eviction (TTL-only).
+const defaultMaxSize = 1000
+
 type memoryCache[T any] struct {
 	lru    *expirable.LRU[string, T]
 	logger Logger
 }
 
 func newMemoryCache[T any](cfg *config) *memoryCache[T] {
-	cfg.logger.Infow("msg", "cache: using in-memory LRU backend", "ttl", cfg.ttl, "maxSize", cfg.maxSize)
+	cfg.logger.Infow("msg", "cache: using in-memory LRU backend", "ttl", cfg.ttl, "maxSize", defaultMaxSize)
 	return &memoryCache[T]{
-		lru:    expirable.NewLRU[string, T](cfg.maxSize, nil, cfg.ttl),
+		lru:    expirable.NewLRU[string, T](defaultMaxSize, nil, cfg.ttl),
 		logger: cfg.logger,
 	}
 }

pkg/cache/natskv.go

@@ -20,6 +20,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"sync"
 
 	"github.com/nats-io/nats.go"
@@ -65,12 +66,29 @@ func (c *natsKVCache[T]) initBucket() error {
 		Bucket:      c.bucket,
 		Description: c.cfg.description,
 		TTL:         c.cfg.ttl,
+		MaxBytes:    c.cfg.maxBytes,
 		Storage:     jetstream.MemoryStorage,
 	})
 	if err != nil {
 		return err
 	}
 
+	// NATS KV hardcodes DiscardNew on the backing stream, which rejects writes
+	// when MaxBytes is reached. For cache use-cases we want DiscardOld so that
+	// the oldest entries are evicted automatically to make room for new ones.
+	if c.cfg.maxBytes > 0 {
+		streamName := fmt.Sprintf("KV_%s", c.bucket)
+		stream, err := js.Stream(context.Background(), streamName)
+		if err != nil {
+			return fmt.Errorf("cache: failed to get backing stream %s: %w", streamName, err)
+		}
+		cfg := stream.CachedInfo().Config
+		cfg.Discard = jetstream.DiscardOld
+		if _, err := js.UpdateStream(context.Background(), cfg); err != nil {
+			return fmt.Errorf("cache: failed to set DiscardOld on stream %s: %w", streamName, err)
+		}
+	}
+
 	c.mu.Lock()
 	c.kv = kv
 	c.mu.Unlock()

pkg/cache/natskv_test.go

@@ -216,6 +216,40 @@ func TestNATSKV_NilKVGracefulDegradation(t *testing.T) {
 	require.NoError(t, nkv.Purge(ctx))
 }
 
+func TestNATSKV_MaxBytesEvictsOldEntries(t *testing.T) {
+	nc := startEmbeddedNATS(t)
+
+	// With MaxBytes set, the backing stream is updated to DiscardOld so that
+	// the oldest entries are evicted when the bucket is full.
+	const maxBytes int64 = 10 * 1024
+	c, err := New[[]byte](
+		WithTTL(time.Minute),
+		WithNATS(nc, "test-maxbytes"),
+		WithMaxBytes(maxBytes),
+	)
+	require.NoError(t, err)
+
+	ctx := context.Background()
+	payload := make([]byte, 1024)
+
+	// Write 20 entries, well beyond the 10 KiB limit
+	for i := range 20 {
+		key := "key-" + strings.Repeat("x", i)
+		require.NoError(t, c.Set(ctx, key, payload), "Set should not fail even when bucket is full")
+	}
+
+	// The latest entry should still be retrievable
+	lastKey := "key-" + strings.Repeat("x", 19)
+	_, ok, err := c.Get(ctx, lastKey)
+	require.NoError(t, err)
+	assert.True(t, ok, "most recent entry should still be in the cache")
+
+	// The earliest entries should have been evicted
+	_, ok, err = c.Get(ctx, "key-")
+	require.NoError(t, err)
+	assert.False(t, ok, "oldest entry should have been evicted")
+}
+
 func TestNew_WithNATSReturnsNATSBackend(t *testing.T) {
 	nc := startEmbeddedNATS(t)
 	c, err := New[string](

pkg/cache/policyevalbundle/policyevalbundle.go

@@ -27,6 +27,7 @@ import (
 
 const (
 	ttl         = 24 * time.Hour
+	maxBytes    = 100 * 1024 * 1024 // 100 MB
 	bucket      = "chainloop-policy-eval-bundles"
 	description = "Cache for policy evaluation bundles from CAS"
 )
@@ -40,6 +41,7 @@ type Cache struct {
 func New(ctx context.Context, rc *natsconn.ReloadableConnection, logger log.Logger) (*Cache, error) {
 	opts := []cache.Option{
 		cache.WithTTL(ttl),
+		cache.WithMaxBytes(maxBytes),
 		cache.WithDescription(description),
 	}