now fetching operatorId dynamically in economicCollector (Layr-Labs#54)

* now fetching operatorId dynamically in economicCollector

* fix economic collect error logic

* fix test

* make economic collector more efficient (less stupid)

* caching operatorId in economic collector

* fix bug

* renamed cacheOperatorIfNotCached -> initOperatorId

chainio/avsregistry/reader.go

@@ -20,6 +20,11 @@ type AvsRegistryReader interface {
 		blockNumber uint32,
 	) ([][]blsoperatorstateretrievar.BLSOperatorStateRetrieverOperator, error)
 
+	GetOperatorStakeInQuorumsOfOperatorAtCurrentBlock(
+		ctx context.Context,
+		operatorId types.OperatorId,
+	) (map[types.QuorumNum]types.StakeAmount, error)
+
 	GetOperatorsStakeInQuorumsOfOperatorAtBlock(
 		ctx context.Context,
 		operatorId types.OperatorId,
@@ -91,7 +96,7 @@ func (r *AvsRegistryChainReader) GetOperatorsStakeInQuorumsOfOperatorAtBlock(
 		operatorId,
 		blockNumber)
 	if err != nil {
-		r.logger.Error("Failed to get operators state", "err", err)
+		r.logger.Error("Failed to get operators state", "err", err, "fn", "AvsRegistryChainReader.GetOperatorsStakeInQuorumsOfOperatorAtBlock")
 		return nil, nil, err
 	}
 
@@ -114,6 +119,30 @@ func (r *AvsRegistryChainReader) GetOperatorsStakeInQuorumsOfOperatorAtCurrentBl
 	return r.GetOperatorsStakeInQuorumsOfOperatorAtBlock(ctx, operatorId, uint32(curBlock))
 }
 
+// GetOperatorStakeInQuorumsOfOperatorAtCurrentBlock could have race conditions
+// it currently makes a bunch of calls to fetch "current block" information,
+// so some of them could actually return information from different blocks
+func (r *AvsRegistryChainReader) GetOperatorStakeInQuorumsOfOperatorAtCurrentBlock(
+	ctx context.Context,
+	operatorId types.OperatorId,
+) (map[types.QuorumNum]types.StakeAmount, error) {
+	quorums, err := r.avsRegistryContractsClient.GetOperatorQuorumsAtCurrentBlock(&bind.CallOpts{}, operatorId)
+	if err != nil {
+		r.logger.Error("Failed to get operator quorums", "err", err)
+		return nil, err
+	}
+	quorumStakes := make(map[types.QuorumNum]types.StakeAmount)
+	for _, quorum := range quorums {
+		stake, err := r.avsRegistryContractsClient.GetCurrentOperatorStakeForQuorum(&bind.CallOpts{}, operatorId, quorum)
+		if err != nil {
+			r.logger.Error("Failed to get operator stake", "err", err)
+			return nil, err
+		}
+		quorumStakes[quorum] = stake
+	}
+	return quorumStakes, nil
+}
+
 func (r *AvsRegistryChainReader) GetCheckSignaturesIndices(
 	ctx context.Context,
 	referenceBlockNumber uint32,

chainio/clients/avs_registry_contracts_client.go

@@ -26,6 +26,17 @@ type AVSRegistryContractsClient interface {
 		operatorAddress gethcommon.Address,
 	) ([32]byte, error)
 
+	GetCurrentOperatorStakeForQuorum(
+		opts *bind.CallOpts,
+		operatorId types.OperatorId,
+		quorumNum types.QuorumNum,
+	) (types.StakeAmount, error)
+
+	GetOperatorQuorumsAtCurrentBlock(
+		opts *bind.CallOpts,
+		operatorId types.OperatorId,
+	) ([]types.QuorumNum, error)
+
 	GetOperatorsStakeInQuorumsAtBlock(
 		opts *bind.CallOpts,
 		quorumNumbers []types.QuorumNum,
@@ -128,6 +139,30 @@ func (a *AvsRegistryContractsChainClient) GetOperatorsStakeInQuorumsAtBlock(
 		blockNumber)
 }
 
+func (a *AvsRegistryContractsChainClient) GetCurrentOperatorStakeForQuorum(
+	opts *bind.CallOpts,
+	operatorId types.OperatorId,
+	quorumNum types.QuorumNum,
+) (types.StakeAmount, error) {
+	return a.avsRegistryBindings.StakeRegistry.GetCurrentOperatorStakeForQuorum(opts, operatorId, quorumNum)
+}
+
+// registry currently only has a fct to retrieve operatorQuorums at current block
+// if we need it at a specific block, we first need to call
+// getQuorumBitmapIndicesByOperatorIdsAtBlockNumber to get the index of the quorum bitmap
+// followed by getQuorumBitmapByOperatorIdAtBlockNumberByIndex
+func (a *AvsRegistryContractsChainClient) GetOperatorQuorumsAtCurrentBlock(
+	opts *bind.CallOpts,
+	operatorId types.OperatorId,
+) ([]types.QuorumNum, error) {
+	quorumBitmap, err := a.avsRegistryBindings.RegistryCoordinator.GetCurrentQuorumBitmapByOperatorId(opts, operatorId)
+	if err != nil {
+		return nil, err
+	}
+	quorums := types.BitmapToQuorumIds(quorumBitmap)
+	return quorums, nil
+}
+
 func (a *AvsRegistryContractsChainClient) GetOperatorsStakeInQuorumsOfOperatorAtBlock(
 	opts *bind.CallOpts,
 	operatorId types.OperatorId,

metrics/collectors/economic/economic.go

@@ -64,17 +64,15 @@ func NewCollector(
 	avsName string, logger logging.Logger,
 	operatorAddr common.Address, quorumNames map[types.QuorumNum]string,
 ) *Collector {
-	operatorId, err := avsRegistryReader.GetOperatorId(context.Background(), operatorAddr)
-	if err != nil {
-		logger.Error("Failed to get operator id", "err", err)
-	}
 	return &Collector{
 		elReader:          elReader,
 		avsRegistryReader: avsRegistryReader,
 		logger:            logger,
 		operatorAddr:      operatorAddr,
-		operatorId:        operatorId,
-		quorumNames:       quorumNames,
+		// we don't fetch operatorId here because operator might not yet be registered (and hence not have an operatorId)
+		// we cache operatorId dynamically in the collect function() inside, which allows constructing collector before registering operator
+		operatorId:  [32]byte{},
+		quorumNames: quorumNames,
 		slashingStatus: prometheus.NewDesc(
 			types.EigenPromNamespace+"_slashing_status",
 			"Whether the operator has been slashed",
@@ -110,6 +108,18 @@ func (ec *Collector) Describe(ch chan<- *prometheus.Desc) {
 	// ch <- ec.delegatedShares
 }
 
+// initialize the operatorId if it hasn't already been initialized
+func (ec *Collector) initOperatorId() error {
+	if ec.operatorId == [32]byte{} {
+		operatorId, err := ec.avsRegistryReader.GetOperatorId(context.Background(), ec.operatorAddr)
+		if err != nil {
+			return err
+		}
+		ec.operatorId = operatorId
+	}
+	return nil
+}
+
 // Collect function for the exported slashingIncurredTotal and balanceTotal metrics
 // see https://github.com/prometheus/client_golang/blob/v1.16.0/prometheus/collector.go#L27
 // collect just makes jsonrpc calls to the slasher and delegationManager and then creates
@@ -132,24 +142,21 @@ func (ec *Collector) Collect(ch chan<- prometheus.Metric) {
 	}
 
 	// collect registeredStake metric
-	// TODO(samlaf): implement this. probably have to call the BLSOperatorStateRetriever contract?
-	// probably should start using the avsregistry service instead of chainio clients so that we can
-	// swap out backend for a subgraph eventually
-	quorumNums, blsOperatorStateRetrieverOperator, err := ec.avsRegistryReader.GetOperatorsStakeInQuorumsOfOperatorAtCurrentBlock(context.Background(), ec.operatorId)
+	err = ec.initOperatorId()
 	if err != nil {
-		ec.logger.Error("Failed to get operator stake", "err", err)
+		ec.logger.Error("Failed to fetch and catch operator id. Skipping collection of registeredStake metric.", "err", err)
 	} else {
-		for quorumIdx, quorumNum := range quorumNums {
-			// TODO: this is stupid.. when AVSs scale to have 5K operators we'll be running through a bunch of operators
-			// we should instead just call registryCoordinator.getQuorumBitmapIndicesByOperatorIdsAtBlockNumber
-			// and stakeRegistry.getStakeForOperatorIdForQuorumAtBlockNumber directly
-			for _, operator := range blsOperatorStateRetrieverOperator[quorumIdx] {
-				if operator.OperatorId == ec.operatorId {
-					stakeFloat64, _ := operator.Stake.Float64()
-					ch <- prometheus.MustNewConstMetric(
-						ec.registeredStake, prometheus.GaugeValue, stakeFloat64, strconv.Itoa(int(quorumNum)), ec.quorumNames[quorumNum],
-					)
-				}
+		// probably should start using the avsregistry service instead of avsRegistryReader so that we can
+		// swap out backend for a subgraph eventually
+		quorumStakeMap, err := ec.avsRegistryReader.GetOperatorStakeInQuorumsOfOperatorAtCurrentBlock(context.Background(), ec.operatorId)
+		if err != nil {
+			ec.logger.Error("Failed to get operator stake", "err", err)
+		} else {
+			for quorumNum, stake := range quorumStakeMap {
+				stakeFloat64, _ := stake.Float64()
+				ch <- prometheus.MustNewConstMetric(
+					ec.registeredStake, prometheus.GaugeValue, stakeFloat64, strconv.Itoa(int(quorumNum)), ec.quorumNames[quorumNum],
+				)
 			}
 		}
 	}

metrics/collectors/economic/economic_test.go

@@ -10,14 +10,13 @@ import (
 	"go.uber.org/mock/gomock"
 
 	"github.com/Layr-Labs/eigensdk-go/chainio/mocks"
-	blsoperatorstateretrievar "github.com/Layr-Labs/eigensdk-go/contracts/bindings/BLSOperatorStateRetriever"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 	"github.com/Layr-Labs/eigensdk-go/types"
 )
 
 func TestEconomicCollector(t *testing.T) {
 	operatorAddr := common.HexToAddress("0x0")
-	operatorId := types.OperatorId{0}
+	operatorId := types.OperatorId{1}
 	quorumNames := map[types.QuorumNum]string{
 		0: "ethQuorum",
 		1: "someOtherTokenQuorum",
@@ -28,29 +27,16 @@ func TestEconomicCollector(t *testing.T) {
 
 	ethReader := mocks.NewMockELReader(mockCtrl)
 	ethReader.EXPECT().OperatorIsFrozen(gomock.Any(), operatorAddr).Return(false, nil)
-	// ethReader.EXPECT().GetOperatorSharesInStrategy(gomock.Any(), operatorAddr, strategyAddrs[0]).Return(big.NewInt(1000), nil)
-	// ethReader.EXPECT().GetOperatorSharesInStrategy(gomock.Any(), operatorAddr, strategyAddrs[1]).Return(big.NewInt(2000), nil)
 
 	avsRegistryReader := mocks.NewMockAvsRegistryReader(mockCtrl)
-	avsRegistryReader.EXPECT().GetOperatorsStakeInQuorumsOfOperatorAtCurrentBlock(gomock.Any(), gomock.Any()).Return(
-		[]types.QuorumNum{0, 1},
-		[][]blsoperatorstateretrievar.BLSOperatorStateRetrieverOperator{
-			{
-				{
-					OperatorId: operatorId,
-					Stake:      big.NewInt(1000),
-				},
-			},
-			{
-				{
-					OperatorId: operatorId,
-					Stake:      big.NewInt(2000),
-				},
-			},
+	avsRegistryReader.EXPECT().GetOperatorId(gomock.Any(), operatorAddr).Return(operatorId, nil)
+	avsRegistryReader.EXPECT().GetOperatorStakeInQuorumsOfOperatorAtCurrentBlock(gomock.Any(), gomock.Any()).Return(
+		map[types.QuorumNum]*big.Int{
+			0: big.NewInt(1000),
+			1: big.NewInt(2000),
 		},
 		nil,
 	)
-	avsRegistryReader.EXPECT().GetOperatorId(gomock.Any(), operatorAddr).Return(operatorId, nil)
 
 	logger := logging.NewNoopLogger()
 	economicCollector := NewCollector(ethReader, avsRegistryReader, "testavs", logger, operatorAddr, quorumNames)

services/avsregistry/avsregistry_chaincaller.go

@@ -74,7 +74,7 @@ func (ar *AvsRegistryServiceChainCaller) GetOperatorsAvsStateAtBlock(ctx context
 func (ar *AvsRegistryServiceChainCaller) GetQuorumsAvsStateAtBlock(ctx context.Context, quorumNumbers []types.QuorumNum, blockNumber types.BlockNum) (map[types.QuorumNum]types.QuorumAvsState, error) {
 	operatorsAvsState, err := ar.GetOperatorsAvsStateAtBlock(ctx, quorumNumbers, blockNumber)
 	if err != nil {
-		ar.logger.Error("Failed to get operator state", "err", err, "service", "AvsRegistryServiceChainCaller")
+		ar.logger.Error("Failed to get quorum state", "err", err, "service", "AvsRegistryServiceChainCaller")
 		return nil, err
 	}
 	quorumsAvsState := make(map[types.QuorumNum]types.QuorumAvsState)