Bound the cachedStates cache

consensus/cometbft/service/cache/cache.go

@@ -23,10 +23,12 @@ package cache
 
 import (
 	"errors"
+	"fmt"
 
 	"github.com/berachain/beacon-kit/consensus/cometbft/service/delay"
 	"github.com/berachain/beacon-kit/primitives/math"
 	"github.com/berachain/beacon-kit/primitives/transition"
+	lru "github.com/hashicorp/golang-lru/v2"
 )
 
 const minActivationHeight = 2
@@ -82,32 +84,35 @@ type Element struct {
 }
 
 type candidateStates struct {
-	states         map[string]*Element
+	states         *lru.Cache[string, *Element]
 	finalStateHash *string
 }
 
-func New() States {
+func New(maxSize int) States {
+	c, err := lru.New[string, *Element](maxSize)
+	if err != nil {
+		panic(fmt.Errorf("failed to create candidate states cache: %w", err))
+	}
 	return &candidateStates{
-		states:         make(map[string]*Element),
+		states:         c,
 		finalStateHash: nil,
 	}
 }
 
 func (cs *candidateStates) SetCached(hash string, toCache *Element) {
-	cs.states[hash] = toCache
+	cs.states.Add(hash, toCache)
 }
 
 func (cs *candidateStates) GetCached(hash string) (*Element, error) {
-	cached, found := cs.states[hash]
+	cached, found := cs.states.Get(hash)
 	if !found {
 		return nil, ErrStateNotFound
 	}
 	return cached, nil
 }
 
 func (cs *candidateStates) MarkAsFinal(hash string) error {
-	_, found := cs.states[hash]
-	if !found {
+	if !cs.states.Contains(hash) {
 		return ErrFinalizingUnknownState
 	}
 	cs.finalStateHash = &hash
@@ -118,14 +123,14 @@ func (cs *candidateStates) GetFinal() (string, *State, error) {
 	if cs.finalStateHash == nil {
 		return "", nil, ErrNoFinalState
 	}
-	cached, found := cs.states[*cs.finalStateHash]
+	cached, found := cs.states.Get(*cs.finalStateHash)
 	if !found {
 		return "", nil, ErrFinalStateIsNil
 	}
 	return *cs.finalStateHash, cached.State, nil
 }
 
 func (cs *candidateStates) Reset() {
-	cs.states = make(map[string]*Element)
+	cs.states.Purge()
 	cs.finalStateHash = nil
 }

consensus/cometbft/service/cache/cache_test.go

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: BUSL-1.1
+//
+// Copyright (C) 2025, Berachain Foundation. All rights reserved.
+// Use of this software is governed by the Business Source License included
+// in the LICENSE file of this repository and at www.mariadb.com/bsl11.
+//
+// ANY USE OF THE LICENSED WORK IN VIOLATION OF THIS LICENSE WILL AUTOMATICALLY
+// TERMINATE YOUR RIGHTS UNDER THIS LICENSE FOR THE CURRENT AND ALL OTHER
+// VERSIONS OF THE LICENSED WORK.
+//
+// THIS LICENSE DOES NOT GRANT YOU ANY RIGHT IN ANY TRADEMARK OR LOGO OF
+// LICENSOR OR ITS AFFILIATES (PROVIDED THAT YOU MAY USE A TRADEMARK OR LOGO OF
+// LICENSOR AS EXPRESSLY REQUIRED BY THIS LICENSE).
+//
+// TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON
+// AN "AS IS" BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS,
+// EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF
+// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND
+// TITLE.
+
+package cache_test
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/berachain/beacon-kit/consensus/cometbft/service/cache"
+)
+
+const maxCachedStates = 10
+
+func TestCandidateStates_BasicOperations(t *testing.T) {
+	t.Parallel()
+	c := cache.New(maxCachedStates)
+
+	// Test SetCached and GetCached
+	elem := &cache.Element{State: &cache.State{}}
+	c.SetCached("test-hash", elem)
+
+	got, err := c.GetCached("test-hash")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if got != elem {
+		t.Error("retrieved element does not match stored element")
+	}
+
+	// Test GetCached for non-existent key
+	_, err = c.GetCached("non-existent")
+	if !errors.Is(err, cache.ErrStateNotFound) {
+		t.Errorf("expected ErrStateNotFound, got: %v", err)
+	}
+
+	// Test MarkAsFinal
+	err = c.MarkAsFinal("test-hash")
+	if err != nil {
+		t.Fatalf("unexpected error marking as final: %v", err)
+	}
+
+	// Test MarkAsFinal for non-existent key
+	err = c.MarkAsFinal("non-existent")
+	if !errors.Is(err, cache.ErrFinalizingUnknownState) {
+		t.Errorf("expected ErrFinalizingUnknownState, got: %v", err)
+	}
+
+	// Test GetFinal
+	hash, state, err := c.GetFinal()
+	if err != nil {
+		t.Fatalf("unexpected error getting final: %v", err)
+	}
+	if hash != "test-hash" {
+		t.Errorf("expected hash 'test-hash', got: %s", hash)
+	}
+	if state != elem.State {
+		t.Error("final state does not match expected state")
+	}
+
+	// Test Reset
+	c.Reset()
+
+	_, err = c.GetCached("test-hash")
+	if !errors.Is(err, cache.ErrStateNotFound) {
+		t.Error("expected cache to be empty after reset")
+	}
+
+	_, _, err = c.GetFinal()
+	if !errors.Is(err, cache.ErrNoFinalState) {
+		t.Errorf("expected ErrNoFinalState after reset, got: %v", err)
+	}
+}
+
+func TestCandidateStates_BoundedSize(t *testing.T) {
+	t.Parallel()
+	c := cache.New(maxCachedStates)
+
+	// Add more entries than maxCachedStates
+	for i := range maxCachedStates + 5 {
+		hash := fmt.Sprintf("hash-%d", i)
+		c.SetCached(hash, &cache.Element{})
+	}
+
+	// Verify oldest entries were evicted (LRU behavior)
+	for i := range 5 {
+		hash := fmt.Sprintf("hash-%d", i)
+		_, err := c.GetCached(hash)
+		if !errors.Is(err, cache.ErrStateNotFound) {
+			t.Errorf("expected hash-%d to be evicted, but it was found", i)
+		}
+	}
+
+	// Verify newest entries are still present
+	for i := 5; i < maxCachedStates+5; i++ {
+		hash := fmt.Sprintf("hash-%d", i)
+		_, err := c.GetCached(hash)
+		if err != nil {
+			t.Errorf("expected hash-%d to be present, got error: %v", i, err)
+		}
+	}
+}

consensus/cometbft/service/service.go

@@ -52,6 +52,7 @@ import (
 const (
 	initialAppVersion uint64 = 0
 	AppName           string = "beacond"
+	maxCachedStates          = 10
 )
 
 type Service struct {
@@ -139,7 +140,7 @@ func NewService(
 		cmtConsensusParams: cmtConsensusParams,
 		cmtCfg:             cmtCfg,
 		telemetrySink:      telemetrySink,
-		cachedStates:       cache.New(),
+		cachedStates:       cache.New(maxCachedStates),
 	}
 
 	s.MountStore(storage.StoreKey, storetypes.StoreTypeIAVL)