Add operator state cache to IndexedChainState

api/clients/retrieval_client_test.go

@@ -100,7 +100,7 @@ func setup(t *testing.T) {
 	indexer = &indexermock.MockIndexer{}
 	indexer.On("Index").Return(nil).Once()
 
-	ics, err := coreindexer.NewIndexedChainState(chainState, indexer)
+	ics, err := coreindexer.NewIndexedChainState(chainState, indexer, 0)
 	if err != nil {
 		panic("failed to create a new indexed chain state")
 	}

core/indexer/state.go

@@ -2,28 +2,43 @@ package indexer
 
 import (
 	"context"
+	"encoding/binary"
 	"errors"
 
 	"github.com/Layr-Labs/eigenda/core"
 	"github.com/Layr-Labs/eigenda/indexer"
+	lru "github.com/hashicorp/golang-lru/v2"
 )
 
 type IndexedChainState struct {
 	core.ChainState
 
-	Indexer indexer.Indexer
+	Indexer            indexer.Indexer
+	operatorStateCache *lru.Cache[string, *core.IndexedOperatorState]
 }
 
 var _ core.IndexedChainState = (*IndexedChainState)(nil)
 
 func NewIndexedChainState(
 	chainState core.ChainState,
 	indexer indexer.Indexer,
+	cacheSize int,
 ) (*IndexedChainState, error) {
+	var operatorStateCache *lru.Cache[string, *core.IndexedOperatorState]
+	var err error
+
+	if cacheSize > 0 {
+		operatorStateCache, err = lru.New[string, *core.IndexedOperatorState](cacheSize)
+		if err != nil {
+			return nil, err
+		}
+	}
 
 	return &IndexedChainState{
 		ChainState: chainState,
 		Indexer:    indexer,
+
+		operatorStateCache: operatorStateCache,
 	}, nil
 }
 
@@ -32,6 +47,13 @@ func (ics *IndexedChainState) Start(ctx context.Context) error {
 }
 
 func (ics *IndexedChainState) GetIndexedOperatorState(ctx context.Context, blockNumber uint, quorums []core.QuorumID) (*core.IndexedOperatorState, error) {
+	// Check if the indexed operator state has been cached
+	cacheKey := computeCacheKey(blockNumber, quorums)
+	if ics.operatorStateCache != nil {
+		if val, ok := ics.operatorStateCache.Get(cacheKey); ok {
+			return val, nil
+		}
+	}
 
 	pubkeys, sockets, err := ics.getObjects(blockNumber)
 	if err != nil {
@@ -73,11 +95,14 @@ func (ics *IndexedChainState) GetIndexedOperatorState(ctx context.Context, block
 		AggKeys:          aggKeys,
 	}
 
+	if ics.operatorStateCache != nil {
+		ics.operatorStateCache.Add(cacheKey, state)
+	}
+
 	return state, nil
 }
 
 func (ics *IndexedChainState) GetIndexedOperators(ctx context.Context, blockNumber uint) (map[core.OperatorID]*core.IndexedOperatorInfo, error) {
-
 	pubkeys, sockets, err := ics.getObjects(blockNumber)
 	if err != nil {
 		return nil, err
@@ -138,3 +163,13 @@ func (ics *IndexedChainState) getObjects(blockNumber uint) (*OperatorPubKeys, Op
 	return pubkeys, sockets, nil
 
 }
+
+// Computes a cache key for the operator state cache. The cache key is a
+// combination of the block number and the quorum IDs. Note: the order of the
+// quorum IDs matters.
+func computeCacheKey(blockNumber uint, quorumIDs []uint8) string {
+	bytes := make([]byte, 8+len(quorumIDs))
+	binary.LittleEndian.PutUint64(bytes, uint64(blockNumber))
+	copy(bytes[8:], quorumIDs)
+	return string(bytes)
+}

core/indexer/state_mock_test.go

@@ -0,0 +1,82 @@
+package indexer_test
+
+import (
+	"context"
+	"math/big"
+	"testing"
+
+	"github.com/Layr-Labs/eigenda/core"
+	"github.com/Layr-Labs/eigenda/core/indexer"
+	coremock "github.com/Layr-Labs/eigenda/core/mock"
+	indexermock "github.com/Layr-Labs/eigenda/indexer/mock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+type testComponents struct {
+	ChainState        *coremock.ChainDataMock
+	Indexer           *indexermock.MockIndexer
+	IndexedChainState *indexer.IndexedChainState
+}
+
+func TestIndexedOperatorStateCache(t *testing.T) {
+	c := createTestComponents(t)
+	pubKeys := &indexer.OperatorPubKeys{}
+	c.Indexer.On("GetObject", mock.Anything, 0).Return(pubKeys, nil)
+	sockets := indexer.OperatorSockets{
+		core.OperatorID{0, 1}: "socket1",
+	}
+	c.Indexer.On("GetObject", mock.Anything, 1).Return(sockets, nil)
+
+	operatorState := &core.OperatorState{
+		Operators: map[core.QuorumID]map[core.OperatorID]*core.OperatorInfo{
+			0: {
+				core.OperatorID{0}: {
+					Stake: big.NewInt(100),
+					Index: 0,
+				},
+			},
+		},
+	}
+	c.ChainState.On("GetOperatorState", mock.Anything, uint(100), []core.QuorumID{0}).Return(operatorState, nil)
+	c.ChainState.On("GetOperatorState", mock.Anything, uint(100), []core.QuorumID{1}).Return(operatorState, nil)
+	c.ChainState.On("GetOperatorState", mock.Anything, uint(101), []core.QuorumID{0, 1}).Return(operatorState, nil)
+
+	ctx := context.Background()
+	// Get the operator state for block 100 and quorum 0
+	_, err := c.IndexedChainState.GetIndexedOperatorState(ctx, uint(100), []core.QuorumID{0})
+	assert.NoError(t, err)
+	c.ChainState.AssertNumberOfCalls(t, "GetOperatorState", 1)
+
+	// Get the operator state for block 100 and quorum 0 again
+	_, err = c.IndexedChainState.GetIndexedOperatorState(ctx, uint(100), []core.QuorumID{0})
+	assert.NoError(t, err)
+	c.ChainState.AssertNumberOfCalls(t, "GetOperatorState", 1)
+
+	// Get the operator state for block 100 and quorum 1
+	_, err = c.IndexedChainState.GetIndexedOperatorState(ctx, uint(100), []core.QuorumID{1})
+	assert.NoError(t, err)
+	c.ChainState.AssertNumberOfCalls(t, "GetOperatorState", 2)
+
+	// Get the operator state for block 101 and quorum 0 & 1
+	_, err = c.IndexedChainState.GetIndexedOperatorState(ctx, uint(101), []core.QuorumID{0, 1})
+	assert.NoError(t, err)
+	c.ChainState.AssertNumberOfCalls(t, "GetOperatorState", 3)
+
+	// Get the operator state for block 101 and quorum 0 & 1 again
+	_, err = c.IndexedChainState.GetIndexedOperatorState(ctx, uint(101), []core.QuorumID{0, 1})
+	assert.NoError(t, err)
+	c.ChainState.AssertNumberOfCalls(t, "GetOperatorState", 3)
+}
+
+func createTestComponents(t *testing.T) *testComponents {
+	chainState := &coremock.ChainDataMock{}
+	idx := &indexermock.MockIndexer{}
+	ics, err := indexer.NewIndexedChainState(chainState, idx, 1)
+	assert.NoError(t, err)
+	return &testComponents{
+		ChainState:        chainState,
+		Indexer:           idx,
+		IndexedChainState: ics,
+	}
+}

core/indexer/state_test.go

@@ -137,7 +137,7 @@ func mustMakeChainState(env *deploy.Config, store indexer.HeaderStore, logger lo
 	)
 	Expect(err).ToNot(HaveOccurred())
 
-	chainState, err := indexedstate.NewIndexedChainState(cs, indexer)
+	chainState, err := indexedstate.NewIndexedChainState(cs, indexer, 0)
 	if err != nil {
 		panic(err)
 	}

core/mock/state.go

@@ -225,6 +225,7 @@ func (d *ChainDataMock) GetTotalOperatorStateWithQuorums(ctx context.Context, bl
 }
 
 func (d *ChainDataMock) GetOperatorState(ctx context.Context, blockNumber uint, quorums []core.QuorumID) (*core.OperatorState, error) {
+	_ = d.Called(ctx, blockNumber, quorums)
 	state := d.GetTotalOperatorStateWithQuorums(ctx, blockNumber, quorums)
 
 	return state.OperatorState, nil

core/thegraph/config.go

@@ -8,15 +8,17 @@ import (
 )
 
 const (
-	EndpointFlagName   = "thegraph.endpoint"
-	BackoffFlagName    = "thegraph.backoff"
-	MaxRetriesFlagName = "thegraph.max_retries"
+	EndpointFlagName       = "thegraph.endpoint"
+	BackoffFlagName        = "thegraph.backoff"
+	MaxRetriesFlagName     = "thegraph.max_retries"
+	OperatorStateCacheSize = "thegraph.operator_state_cache_size"
 )
 
 type Config struct {
-	Endpoint     string        // The Graph endpoint
-	PullInterval time.Duration // The interval to pull data from The Graph
-	MaxRetries   int           // The maximum number of retries to pull data from The Graph
+	Endpoint               string        // The Graph endpoint
+	PullInterval           time.Duration // The interval to pull data from The Graph
+	MaxRetries             int           // The maximum number of retries to pull data from The Graph
+	OperatorStateCacheSize int           // The size of the cache
 }
 
 func CLIFlags(envPrefix string) []cli.Flag {
@@ -39,15 +41,22 @@ func CLIFlags(envPrefix string) []cli.Flag {
 			Value:  5,
 			EnvVar: common.PrefixEnvVar(envPrefix, "GRAPH_MAX_RETRIES"),
 		},
+		cli.IntFlag{
+			Name:   OperatorStateCacheSize,
+			Usage:  "The size of the operator state cache in elements (0 to disable)",
+			Value:  0,
+			EnvVar: common.PrefixEnvVar(envPrefix, "GRAPH_OPERATOR_STATE_CACHE_SIZE"),
+		},
 	}
 }
 
 func ReadCLIConfig(ctx *cli.Context) Config {
 
 	return Config{
-		Endpoint:     ctx.String(EndpointFlagName),
-		PullInterval: ctx.Duration(BackoffFlagName),
-		MaxRetries:   ctx.Int(MaxRetriesFlagName),
+		Endpoint:               ctx.String(EndpointFlagName),
+		PullInterval:           ctx.Duration(BackoffFlagName),
+		MaxRetries:             ctx.Int(MaxRetriesFlagName),
+		OperatorStateCacheSize: ctx.Int(OperatorStateCacheSize),
 	}
 
 }

core/thegraph/state.go

@@ -2,6 +2,7 @@ package thegraph
 
 import (
 	"context"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"math"
@@ -10,6 +11,7 @@ import (
 	"github.com/Layr-Labs/eigenda/core"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 	"github.com/consensys/gnark-crypto/ecc/bn254"
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/shurcooL/graphql"
 )
 
@@ -79,29 +81,41 @@ type (
 		core.ChainState
 		querier GraphQLQuerier
 
-		logger logging.Logger
+		logger             logging.Logger
+		operatorStateCache *lru.Cache[string, *core.IndexedOperatorState]
 	}
 )
 
 var _ IndexedChainState = (*indexedChainState)(nil)
 
-func MakeIndexedChainState(config Config, cs core.ChainState, logger logging.Logger) *indexedChainState {
-
+func MakeIndexedChainState(config Config, cs core.ChainState, logger logging.Logger) (*indexedChainState, error) {
 	logger.Info("Using graph node")
 	querier := graphql.NewClient(config.Endpoint, nil)
 
 	// RetryQuerier is a wrapper around the GraphQLQuerier that retries queries on failure
 	retryQuerier := NewRetryQuerier(querier, config.PullInterval, config.MaxRetries)
 
-	return NewIndexedChainState(cs, retryQuerier, logger)
+	return NewIndexedChainState(cs, retryQuerier, logger, config.OperatorStateCacheSize)
 }
 
-func NewIndexedChainState(cs core.ChainState, querier GraphQLQuerier, logger logging.Logger) *indexedChainState {
+func NewIndexedChainState(cs core.ChainState, querier GraphQLQuerier, logger logging.Logger, cacheSize int) (*indexedChainState, error) {
+	var operatorStateCache *lru.Cache[string, *core.IndexedOperatorState]
+	var err error
+
+	if cacheSize > 0 {
+		operatorStateCache, err = lru.New[string, *core.IndexedOperatorState](cacheSize)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return &indexedChainState{
 		ChainState: cs,
-		querier:    querier,
-		logger:     logger.With("component", "IndexedChainState"),
-	}
+
+		querier:            querier,
+		logger:             logger.With("component", "IndexedChainState"),
+		operatorStateCache: operatorStateCache,
+	}, nil
 }
 
 func (ics *indexedChainState) Start(ctx context.Context) error {
@@ -124,6 +138,14 @@ func (ics *indexedChainState) Start(ctx context.Context) error {
 }
 
 func (ics *indexedChainState) GetIndexedOperatorState(ctx context.Context, blockNumber uint, quorums []core.QuorumID) (*core.IndexedOperatorState, error) {
+	// Check if the indexed operator state has been cached
+	cacheKey := computeCacheKey(blockNumber, quorums)
+	if ics.operatorStateCache != nil {
+		if val, ok := ics.operatorStateCache.Get(cacheKey); ok {
+			return val, nil
+		}
+	}
+
 	operatorState, err := ics.ChainState.GetOperatorState(ctx, blockNumber, quorums)
 	if err != nil {
 		return nil, err
@@ -173,6 +195,11 @@ func (ics *indexedChainState) GetIndexedOperatorState(ctx context.Context, block
 		IndexedOperators: indexedOperators,
 		AggKeys:          aggKeys,
 	}
+
+	if ics.operatorStateCache != nil {
+		ics.operatorStateCache.Add(cacheKey, state)
+	}
+
 	return state, nil
 }
 
@@ -365,3 +392,13 @@ func convertIndexedOperatorInfoGqlToIndexedOperatorInfo(operator *IndexedOperato
 		Socket:   string(operator.SocketUpdates[0].Socket),
 	}, nil
 }
+
+// Computes a cache key for the operator state cache. The cache key is a
+// combination of the block number and the quorum IDs. Note: the order of the
+// quorum IDs matters.
+func computeCacheKey(blockNumber uint, quorumIDs []uint8) string {
+	bytes := make([]byte, 8+len(quorumIDs))
+	binary.LittleEndian.PutUint64(bytes, uint64(blockNumber))
+	copy(bytes[8:], quorumIDs)
+	return string(bytes)
+}

core/thegraph/state_integration_test.go

@@ -86,7 +86,8 @@ func TestIndexerIntegration(t *testing.T) {
 	tx, err := eth.NewWriter(logger, client, testConfig.EigenDA.OperatorStateRetreiver, testConfig.EigenDA.ServiceManager)
 	assert.NoError(t, err)
 
-	cs := thegraph.NewIndexedChainState(eth.NewChainState(tx, client), graphql.NewClient(graphUrl, nil), logger)
+	cs, err := thegraph.NewIndexedChainState(eth.NewChainState(tx, client), graphql.NewClient(graphUrl, nil), logger, 0)
+	assert.NoError(t, err)
 	time.Sleep(5 * time.Second)
 
 	err = cs.Start(context.Background())