feat: distributed retrieval operator socket (#1403)

Author: hopeyen

api/clients/mock/node_client.go

@@ -19,7 +19,7 @@ func NewNodeClient() *MockNodeClient {
 	return &MockNodeClient{}
 }
 
-func (c *MockNodeClient) GetBlobHeader(ctx context.Context, socket string, batchHeaderHash [32]byte, blobIndex uint32) (*core.BlobHeader, *merkletree.Proof, error) {
+func (c *MockNodeClient) GetBlobHeader(ctx context.Context, socket core.OperatorSocket, batchHeaderHash [32]byte, blobIndex uint32) (*core.BlobHeader, *merkletree.Proof, error) {
 	args := c.Called(socket, batchHeaderHash, blobIndex)
 	var hashes [][]byte
 	if args.Get(1) != nil {
@@ -46,7 +46,7 @@ func (c *MockNodeClient) GetBlobHeader(ctx context.Context, socket string, batch
 func (c *MockNodeClient) GetChunks(
 	ctx context.Context,
 	opID core.OperatorID,
-	opInfo *core.IndexedOperatorInfo,
+	opInfo *core.OperatorInfo,
 	batchHeaderHash [32]byte,
 	blobIndex uint32,
 	quorumID core.QuorumID,

api/clients/node_client.go

@@ -20,8 +20,8 @@ type RetrievedChunks struct {
 }
 
 type NodeClient interface {
-	GetBlobHeader(ctx context.Context, socket string, batchHeaderHash [32]byte, blobIndex uint32) (*core.BlobHeader, *merkletree.Proof, error)
-	GetChunks(ctx context.Context, opID core.OperatorID, opInfo *core.IndexedOperatorInfo, batchHeaderHash [32]byte, blobIndex uint32, quorumID core.QuorumID, chunksChan chan RetrievedChunks)
+	GetBlobHeader(ctx context.Context, socket core.OperatorSocket, batchHeaderHash [32]byte, blobIndex uint32) (*core.BlobHeader, *merkletree.Proof, error)
+	GetChunks(ctx context.Context, opID core.OperatorID, opInfo *core.OperatorInfo, batchHeaderHash [32]byte, blobIndex uint32, quorumID core.QuorumID, chunksChan chan RetrievedChunks)
 }
 
 type client struct {
@@ -36,12 +36,12 @@ func NewNodeClient(timeout time.Duration) NodeClient {
 
 func (c client) GetBlobHeader(
 	ctx context.Context,
-	socket string,
+	socket core.OperatorSocket,
 	batchHeaderHash [32]byte,
 	blobIndex uint32,
 ) (*core.BlobHeader, *merkletree.Proof, error) {
 	conn, err := grpc.NewClient(
-		core.OperatorSocket(socket).GetV1RetrievalSocket(),
+		socket.GetV1RetrievalSocket(),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {
@@ -79,7 +79,7 @@ func (c client) GetBlobHeader(
 func (c client) GetChunks(
 	ctx context.Context,
 	opID core.OperatorID,
-	opInfo *core.IndexedOperatorInfo,
+	opInfo *core.OperatorInfo,
 	batchHeaderHash [32]byte,
 	blobIndex uint32,
 	quorumID core.QuorumID,

api/clients/retrieval_client.go

@@ -53,7 +53,7 @@ type BlobChunks struct {
 
 type retrievalClient struct {
 	logger                logging.Logger
-	indexedChainState     core.IndexedChainState
+	chainState            core.ChainState
 	assignmentCoordinator core.AssignmentCoordinator
 	nodeClient            NodeClient
 	verifier              encoding.Verifier
@@ -63,15 +63,15 @@ type retrievalClient struct {
 // NewRetrievalClient creates a new retrieval client.
 func NewRetrievalClient(
 	logger logging.Logger,
-	chainState core.IndexedChainState,
+	chainState core.ChainState,
 	assignmentCoordinator core.AssignmentCoordinator,
 	nodeClient NodeClient,
 	verifier encoding.Verifier,
 	numConnections int) (RetrievalClient, error) {
 
 	return &retrievalClient{
 		logger:                logger.With("component", "RetrievalClient"),
-		indexedChainState:     chainState,
+		chainState:            chainState,
 		assignmentCoordinator: assignmentCoordinator,
 		nodeClient:            nodeClient,
 		verifier:              verifier,
@@ -104,11 +104,12 @@ func (r *retrievalClient) RetrieveBlobChunks(ctx context.Context,
 	batchRoot [32]byte,
 	quorumID core.QuorumID) (*BlobChunks, error) {
 
-	indexedOperatorState, err := r.indexedChainState.GetIndexedOperatorState(ctx, referenceBlockNumber, []core.QuorumID{quorumID})
+	operatorState, err := r.chainState.GetOperatorStateWithSocket(ctx, referenceBlockNumber, []core.QuorumID{quorumID})
 	if err != nil {
+		r.logger.Error("failed to get operator state", "err", err)
 		return nil, err
 	}
-	operators, ok := indexedOperatorState.Operators[quorumID]
+	operators, ok := operatorState.Operators[quorumID]
 	if !ok {
 		return nil, fmt.Errorf("no quorum with ID: %d", quorumID)
 	}
@@ -118,7 +119,7 @@ func (r *retrievalClient) RetrieveBlobChunks(ctx context.Context,
 	var proof *merkletree.Proof
 	var proofVerified bool
 	for opID := range operators {
-		opInfo := indexedOperatorState.IndexedOperators[opID]
+		opInfo := operators[opID]
 		blobHeader, proof, err = r.nodeClient.GetBlobHeader(ctx, opInfo.Socket, batchHeaderHash, blobIndex)
 		if err != nil {
 			// try another operator
@@ -172,7 +173,7 @@ func (r *retrievalClient) RetrieveBlobChunks(ctx context.Context,
 		return nil, err
 	}
 
-	assignments, info, err := r.assignmentCoordinator.GetAssignments(indexedOperatorState.OperatorState, blobHeader.Length, quorumHeader)
+	assignments, info, err := r.assignmentCoordinator.GetAssignments(operatorState, blobHeader.Length, quorumHeader)
 	if err != nil {
 		return nil, errors.New("failed to get assignments")
 	}
@@ -181,8 +182,7 @@ func (r *retrievalClient) RetrieveBlobChunks(ctx context.Context,
 	chunksChan := make(chan RetrievedChunks, len(operators))
 	pool := workerpool.New(r.numConnections)
 	for opID := range operators {
-		opID := opID
-		opInfo := indexedOperatorState.IndexedOperators[opID]
+		opInfo := operators[opID]
 		pool.Submit(func() {
 			r.nodeClient.GetChunks(ctx, opID, opInfo, batchHeaderHash, blobIndex, quorumID, chunksChan)
 		})

api/clients/retrieval_client_test.go

@@ -100,12 +100,7 @@ func setup(t *testing.T) {
 	indexer = &indexermock.MockIndexer{}
 	indexer.On("Index").Return(nil).Once()
 
-	ics, err := coreindexer.NewIndexedChainState(chainState, indexer)
-	if err != nil {
-		panic("failed to create a new indexed chain state")
-	}
-
-	retrievalClient, err = clients.NewRetrievalClient(logger, ics, coordinator, nodeClient, v, 2)
+	retrievalClient, err = clients.NewRetrievalClient(logger, chainState, coordinator, nodeClient, v, 2)
 	if err != nil {
 		panic("failed to create a new retrieval client")
 	}

api/clients/v2/retrieval_client.go

@@ -34,11 +34,11 @@ type RetrievalClient interface {
 }
 
 type retrievalClient struct {
-	logger            logging.Logger
-	ethClient         core.Reader
-	indexedChainState core.IndexedChainState
-	verifier          encoding.Verifier
-	numConnections    int
+	logger         logging.Logger
+	ethClient      core.Reader
+	chainState     core.ChainState
+	verifier       encoding.Verifier
+	numConnections int
 }
 
 var _ RetrievalClient = &retrievalClient{}
@@ -47,16 +47,16 @@ var _ RetrievalClient = &retrievalClient{}
 func NewRetrievalClient(
 	logger logging.Logger,
 	ethClient core.Reader,
-	chainState core.IndexedChainState,
+	chainState core.ChainState,
 	verifier encoding.Verifier,
 	numConnections int,
 ) RetrievalClient {
 	return &retrievalClient{
-		logger:            logger.With("component", "RetrievalClient"),
-		ethClient:         ethClient,
-		indexedChainState: chainState,
-		verifier:          verifier,
-		numConnections:    numConnections,
+		logger:         logger.With("component", "RetrievalClient"),
+		ethClient:      ethClient,
+		chainState:     chainState,
+		verifier:       verifier,
+		numConnections: numConnections,
 	}
 }
 
@@ -75,11 +75,11 @@ func (r *retrievalClient) GetBlob(
 		return nil, err
 	}
 
-	indexedOperatorState, err := r.indexedChainState.GetIndexedOperatorState(ctx, uint(referenceBlockNumber), []core.QuorumID{quorumID})
+	operatorState, err := r.chainState.GetOperatorStateWithSocket(ctx, uint(referenceBlockNumber), []core.QuorumID{quorumID})
 	if err != nil {
 		return nil, err
 	}
-	operators, ok := indexedOperatorState.Operators[quorumID]
+	operators, ok := operatorState.Operators[quorumID]
 	if !ok {
 		return nil, fmt.Errorf("no quorum with ID: %d", quorumID)
 	}
@@ -99,7 +99,7 @@ func (r *retrievalClient) GetBlob(
 		return nil, err
 	}
 
-	assignments, err := corev2.GetAssignments(indexedOperatorState.OperatorState, blobParam, quorumID)
+	assignments, err := corev2.GetAssignments(operatorState, blobParam, quorumID)
 	if err != nil {
 		return nil, errors.New("failed to get assignments")
 	}
@@ -108,8 +108,7 @@ func (r *retrievalClient) GetBlob(
 	chunksChan := make(chan clients.RetrievedChunks, len(operators))
 	pool := workerpool.New(r.numConnections)
 	for opID := range operators {
-		opID := opID
-		opInfo := indexedOperatorState.IndexedOperators[opID]
+		opInfo := operatorState.Operators[quorumID][opID]
 		pool.Submit(func() {
 			r.getChunksFromOperator(ctx, opID, opInfo, blobKey, quorumID, chunksChan)
 		})
@@ -161,7 +160,7 @@ func (r *retrievalClient) GetBlob(
 func (r *retrievalClient) getChunksFromOperator(
 	ctx context.Context,
 	opID core.OperatorID,
-	opInfo *core.IndexedOperatorInfo,
+	opInfo *core.OperatorInfo,
 	blobKey corev2.BlobKey,
 	quorumID core.QuorumID,
 	chunksChan chan clients.RetrievedChunks,

core/chainio.go

@@ -16,6 +16,12 @@ type OperatorStake struct {
 	Stake      *big.Int
 }
 
+type OperatorStakeWithSocket struct {
+	OperatorID OperatorID
+	Stake      *big.Int
+	Socket     OperatorSocket
+}
+
 type OperatorToChurn struct {
 	QuorumId QuorumID
 	Operator gethcommon.Address
@@ -29,6 +35,7 @@ type OperatorSetParam struct {
 }
 
 type OperatorStakes map[QuorumID]map[OperatorIndex]OperatorStake
+type OperatorStakesWithSocket map[QuorumID]map[OperatorIndex]OperatorStakeWithSocket
 
 type Reader interface {
 
@@ -44,6 +51,10 @@ type Reader interface {
 	// The indices of the operators within each quorum are also returned.
 	GetOperatorStakesForQuorums(ctx context.Context, quorums []QuorumID, blockNumber uint32) (OperatorStakes, error)
 
+	// GetOperatorStakesWithSocketForQuorums returns the stakes of all operators within the supplied quorums. The returned stakes are for the block number supplied.
+	// The indices of the operators within each quorum are also returned.
+	GetOperatorStakesWithSocketForQuorums(ctx context.Context, quorums []QuorumID, blockNumber uint32) (OperatorStakesWithSocket, error)
+
 	// GetBlockStaleMeasure returns the BLOCK_STALE_MEASURE defined onchain.
 	GetBlockStaleMeasure(ctx context.Context) (uint32, error)
 

core/eth/reader.go

@@ -483,6 +483,43 @@ func (t *Reader) GetOperatorStakesForQuorums(ctx context.Context, quorums []core
 	return state, nil
 }
 
+// GetOperatorStakesForQuorums returns the stakes of all operators within the supplied quorums. The returned stakes are for the block number supplied.
+// The indices of the operators within each quorum are also returned.
+func (t *Reader) GetOperatorStakesWithSocketForQuorums(ctx context.Context, quorums []core.QuorumID, blockNumber uint32) (core.OperatorStakesWithSocket, error) {
+	quorumBytes := make([]byte, len(quorums))
+	for ind, quorum := range quorums {
+		quorumBytes[ind] = byte(uint8(quorum))
+	}
+
+	// result is a struct{Operators [][]opstateretriever.OperatorStateRetrieverOperator; Sockets [][]string}
+	// Operators is a [][]*opstateretriever.OperatorStake with the same length and order as quorumBytes, and then indexed by operator index
+	// Sockets is a [][]string with the same length and order as quorumBytes, and then indexed by operator index
+	// By contract definition, Operators and Sockets are parallel arrays
+	result, err := t.bindings.OpStateRetriever.GetOperatorStateWithSocket(&bind.CallOpts{
+		Context: ctx,
+	}, t.bindings.RegCoordinatorAddr, quorumBytes, blockNumber)
+	if err != nil {
+		t.logger.Errorf("Failed to fetch operator state: %s", err)
+		return nil, fmt.Errorf("failed to fetch operator state: %w", err)
+	}
+
+	state := make(core.OperatorStakesWithSocket, len(result.Operators))
+	for i := range result.Operators {
+		quorumID := quorums[i]
+		state[quorumID] = make(map[core.OperatorIndex]core.OperatorStakeWithSocket, len(result.Operators[i]))
+		for j, op := range result.Operators[i] {
+			operatorIndex := core.OperatorIndex(j)
+			state[quorumID][operatorIndex] = core.OperatorStakeWithSocket{
+				Stake:      op.Stake,
+				OperatorID: op.OperatorId,
+				Socket:     core.OperatorSocket(result.Sockets[i][j]),
+			}
+		}
+	}
+
+	return state, nil
+}
+
 func (t *Reader) StakeRegistry(ctx context.Context) (gethcommon.Address, error) {
 	return t.bindings.RegistryCoordinator.StakeRegistry(&bind.CallOpts{
 		Context: ctx,

core/eth/state.go

@@ -41,6 +41,15 @@ func (cs *ChainState) GetOperatorState(ctx context.Context, blockNumber uint, qu
 	return getOperatorState(operatorsByQuorum, uint32(blockNumber))
 }
 
+func (cs *ChainState) GetOperatorStateWithSocket(ctx context.Context, blockNumber uint, quorums []core.QuorumID) (*core.OperatorState, error) {
+	operatorsByQuorum, err := cs.Tx.GetOperatorStakesWithSocketForQuorums(ctx, quorums, uint32(blockNumber))
+	if err != nil {
+		return nil, err
+	}
+
+	return getOperatorStateWithSocket(operatorsByQuorum, uint32(blockNumber))
+}
+
 func (cs *ChainState) GetCurrentBlockNumber(ctx context.Context) (uint, error) {
 	header, err := cs.Client.HeaderByNumber(ctx, nil)
 	if err != nil {
@@ -88,3 +97,36 @@ func getOperatorState(operatorsByQuorum core.OperatorStakes, blockNumber uint32)
 
 	return state, nil
 }
+
+func getOperatorStateWithSocket(operatorsByQuorum core.OperatorStakesWithSocket, blockNumber uint32) (*core.OperatorState, error) {
+	operators := make(map[core.QuorumID]map[core.OperatorID]*core.OperatorInfo)
+	totals := make(map[core.QuorumID]*core.OperatorInfo)
+
+	for quorumID, quorum := range operatorsByQuorum {
+		totalStake := big.NewInt(0)
+		operators[quorumID] = make(map[core.OperatorID]*core.OperatorInfo)
+
+		for ind, op := range quorum {
+			operators[quorumID][op.OperatorID] = &core.OperatorInfo{
+				Stake:  op.Stake,
+				Index:  core.OperatorIndex(ind),
+				Socket: core.OperatorSocket(op.Socket),
+			}
+			totalStake.Add(totalStake, op.Stake)
+		}
+
+		totals[quorumID] = &core.OperatorInfo{
+			Stake:  totalStake,
+			Index:  core.OperatorIndex(len(quorum)),
+			Socket: core.OperatorSocket(""),
+		}
+	}
+
+	state := &core.OperatorState{
+		Operators:   operators,
+		Totals:      totals,
+		BlockNumber: uint(blockNumber),
+	}
+
+	return state, nil
+}

core/mock/state.go

@@ -245,6 +245,12 @@ func (d *ChainDataMock) GetOperatorState(ctx context.Context, blockNumber uint,
 	return state.OperatorState, nil
 }
 
+func (d *ChainDataMock) GetOperatorStateWithSocket(ctx context.Context, blockNumber uint, quorums []core.QuorumID) (*core.OperatorState, error) {
+	state := d.GetTotalOperatorStateWithQuorums(ctx, blockNumber, quorums)
+
+	return state.OperatorState, nil
+}
+
 func (d *ChainDataMock) GetOperatorStateByOperator(ctx context.Context, blockNumber uint, operator core.OperatorID) (*core.OperatorState, error) {
 	quorums := make([]core.QuorumID, 0)
 	for quorumID, stake := range d.Stakes {

core/mock/writer.go

@@ -95,6 +95,15 @@ func (t *MockWriter) GetOperatorStakesForQuorums(ctx context.Context, quorums []
 	return result.(core.OperatorStakes), args.Error(1)
 }
 
+func (t *MockWriter) GetOperatorStakesWithSocketForQuorums(ctx context.Context, quorums []core.QuorumID, blockNumber uint32) (core.OperatorStakesWithSocket, error) {
+	args := t.Called()
+	result := args.Get(0)
+	if fn, ok := result.(func([]core.QuorumID, uint32) core.OperatorStakesWithSocket); ok {
+		return fn(quorums, blockNumber), args.Error(1)
+	}
+	return result.(core.OperatorStakesWithSocket), args.Error(1)
+}
+
 func (t *MockWriter) BuildConfirmBatchTxn(ctx context.Context, batchHeader *core.BatchHeader, quorums map[core.QuorumID]*core.QuorumResult, signatureAggregation *core.SignatureAggregation) (*types.Transaction, error) {
 	args := t.Called(ctx, batchHeader, quorums, signatureAggregation)
 	result := args.Get(0)