beacon api: optimize val assignment lookup (#15558)

beacon-chain/rpc/prysm/v1alpha1/validator/duties_v2.go

@@ -17,6 +17,13 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+const (
+	// validatorLookupThreshold determines when to use full assignment map vs cached linear search.
+	// For requests with fewer validators, we use cached linear search to avoid the overhead
+	// of building a complete assignment map for all validators in the epoch.
+	validatorLookupThreshold = 3000
+)
+
 // GetDutiesV2 returns the duties assigned to a list of validators specified
 // in the request object.
 //
@@ -53,8 +60,7 @@ func (vs *Server) dutiesv2(ctx context.Context, req *ethpb.DutiesRequest) (*ethp
 	span.SetAttributes(trace.Int64Attribute("num_pubkeys", int64(len(req.PublicKeys))))
 	defer span.End()
 
-	// Load committee and proposer metadata
-	meta, err := loadDutiesMetadata(ctx, s, req.Epoch)
+	meta, err := loadDutiesMetadata(ctx, s, req.Epoch, len(req.PublicKeys))
 	if err != nil {
 		return nil, err
 	}
@@ -68,24 +74,22 @@ func (vs *Server) dutiesv2(ctx context.Context, req *ethpb.DutiesRequest) (*ethp
 			return nil, status.Errorf(codes.Aborted, "Could not continue fetching assignments: %v", ctx.Err())
 		}
 
-		idx, ok := s.ValidatorIndexByPubkey(bytesutil.ToBytes48(pubKey))
+		validatorIndex, ok := s.ValidatorIndexByPubkey(bytesutil.ToBytes48(pubKey))
 		if !ok {
-			// Unknown validator: still append placeholder duty with UNKNOWN_STATUS
-			validatorAssignments = append(validatorAssignments, &ethpb.DutiesV2Response_Duty{
+			unknownDuty := &ethpb.DutiesV2Response_Duty{
 				PublicKey: pubKey,
 				Status:    ethpb.ValidatorStatus_UNKNOWN_STATUS,
-			})
-			nextValidatorAssignments = append(nextValidatorAssignments, &ethpb.DutiesV2Response_Duty{
-				PublicKey: pubKey,
-				Status:    ethpb.ValidatorStatus_UNKNOWN_STATUS,
-			})
+			}
+			validatorAssignments = append(validatorAssignments, unknownDuty)
+			nextValidatorAssignments = append(nextValidatorAssignments, unknownDuty)
 			continue
 		}
 
-		meta.current.liteAssignment = helpers.AssignmentForValidator(meta.current.committeesBySlot, meta.current.startSlot, idx)
-		meta.next.liteAssignment = helpers.AssignmentForValidator(meta.next.committeesBySlot, meta.next.startSlot, idx)
+		meta.current.liteAssignment = vs.getValidatorAssignment(meta.current, validatorIndex)
 
-		assignment, nextAssignment, err := vs.buildValidatorDuty(pubKey, idx, s, req.Epoch, meta)
+		meta.next.liteAssignment = vs.getValidatorAssignment(meta.next, validatorIndex)
+
+		assignment, nextAssignment, err := vs.buildValidatorDuty(pubKey, validatorIndex, s, req.Epoch, meta)
 		if err != nil {
 			return nil, err
 		}
@@ -143,17 +147,18 @@ type dutiesMetadata struct {
 }
 
 type metadata struct {
-	committeesAtSlot uint64
-	proposalSlots    map[primitives.ValidatorIndex][]primitives.Slot
-	startSlot        primitives.Slot
-	committeesBySlot [][][]primitives.ValidatorIndex
-	liteAssignment   *helpers.LiteAssignment
+	committeesAtSlot       uint64
+	proposalSlots          map[primitives.ValidatorIndex][]primitives.Slot
+	startSlot              primitives.Slot
+	committeesBySlot       [][][]primitives.ValidatorIndex
+	validatorAssignmentMap map[primitives.ValidatorIndex]*helpers.LiteAssignment
+	liteAssignment         *helpers.LiteAssignment
 }
 
-func loadDutiesMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.Epoch) (*dutiesMetadata, error) {
+func loadDutiesMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.Epoch, numValidators int) (*dutiesMetadata, error) {
 	meta := &dutiesMetadata{}
 	var err error
-	meta.current, err = loadMetadata(ctx, s, reqEpoch)
+	meta.current, err = loadMetadata(ctx, s, reqEpoch, numValidators)
 	if err != nil {
 		return nil, err
 	}
@@ -163,14 +168,14 @@ func loadDutiesMetadata(ctx context.Context, s state.BeaconState, reqEpoch primi
 		return nil, status.Errorf(codes.Internal, "Could not compute proposer slots: %v", err)
 	}
 
-	meta.next, err = loadMetadata(ctx, s, reqEpoch+1)
+	meta.next, err = loadMetadata(ctx, s, reqEpoch+1, numValidators)
 	if err != nil {
 		return nil, err
 	}
 	return meta, nil
 }
 
-func loadMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.Epoch) (*metadata, error) {
+func loadMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.Epoch, numValidators int) (*metadata, error) {
 	meta := &metadata{}
 
 	if err := helpers.VerifyAssignmentEpoch(reqEpoch, s); err != nil {
@@ -193,9 +198,48 @@ func loadMetadata(ctx context.Context, s state.BeaconState, reqEpoch primitives.
 		return nil, err
 	}
 
+	if numValidators >= validatorLookupThreshold {
+		meta.validatorAssignmentMap = buildValidatorAssignmentMap(meta.committeesBySlot, meta.startSlot)
+	}
+
 	return meta, nil
 }
 
+// buildValidatorAssignmentMap creates a map from validator index to assignment for O(1) lookup.
+func buildValidatorAssignmentMap(
+	bySlot [][][]primitives.ValidatorIndex,
+	startSlot primitives.Slot,
+) map[primitives.ValidatorIndex]*helpers.LiteAssignment {
+	validatorToAssignment := make(map[primitives.ValidatorIndex]*helpers.LiteAssignment)
+
+	for relativeSlot, committees := range bySlot {
+		for cIdx, committee := range committees {
+			for pos, vIdx := range committee {
+				validatorToAssignment[vIdx] = &helpers.LiteAssignment{
+					AttesterSlot:            startSlot + primitives.Slot(relativeSlot),
+					CommitteeIndex:          primitives.CommitteeIndex(cIdx),
+					CommitteeLength:         uint64(len(committee)),
+					ValidatorCommitteeIndex: uint64(pos),
+				}
+			}
+		}
+	}
+	return validatorToAssignment
+}
+
+// getValidatorAssignment retrieves the assignment for a validator using either
+// the pre-built assignment map (for large requests) or linear search (for small requests).
+func (vs *Server) getValidatorAssignment(meta *metadata, validatorIndex primitives.ValidatorIndex) *helpers.LiteAssignment {
+	if meta.validatorAssignmentMap != nil {
+		if assignment, exists := meta.validatorAssignmentMap[validatorIndex]; exists {
+			return assignment
+		}
+		return &helpers.LiteAssignment{}
+	}
+
+	return helpers.AssignmentForValidator(meta.committeesBySlot, meta.startSlot, validatorIndex)
+}
+
 // buildValidatorDuty builds both current‑epoch and next‑epoch V2 duty objects
 // for a single validator index.
 func (vs *Server) buildValidatorDuty(

beacon-chain/rpc/prysm/v1alpha1/validator/duties_v2_test.go

@@ -559,3 +559,170 @@ func TestGetDutiesV2_SyncNotReady(t *testing.T) {
 	_, err := vs.GetDutiesV2(t.Context(), &ethpb.DutiesRequest{})
 	assert.ErrorContains(t, "Syncing to latest head", err)
 }
+
+func TestBuildValidatorAssignmentMap(t *testing.T) {
+	start := primitives.Slot(200)
+	bySlot := [][][]primitives.ValidatorIndex{
+		{{1, 2, 3}},        // slot 200, committee 0
+		{{7, 8, 9}},        // slot 201, committee 0  
+		{{4, 5}, {10, 11}}, // slot 202, committee 0 & 1
+	}
+
+	assignmentMap := buildValidatorAssignmentMap(bySlot, start)
+
+	// Test validator 8 assignment (slot 201, committee 0, position 1)
+	vIdx := primitives.ValidatorIndex(8)
+	got, exists := assignmentMap[vIdx]
+	assert.Equal(t, true, exists)
+	require.NotNil(t, got)
+	assert.Equal(t, start+1, got.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(0), got.CommitteeIndex)
+	assert.Equal(t, uint64(3), got.CommitteeLength)
+	assert.Equal(t, uint64(1), got.ValidatorCommitteeIndex)
+
+	// Test validator 1 assignment (slot 200, committee 0, position 0)
+	vIdx1 := primitives.ValidatorIndex(1)
+	got1, exists1 := assignmentMap[vIdx1]
+	assert.Equal(t, true, exists1)
+	require.NotNil(t, got1)
+	assert.Equal(t, start, got1.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(0), got1.CommitteeIndex)
+	assert.Equal(t, uint64(3), got1.CommitteeLength)
+	assert.Equal(t, uint64(0), got1.ValidatorCommitteeIndex)
+
+	// Test validator 10 assignment (slot 202, committee 1, position 0)
+	vIdx10 := primitives.ValidatorIndex(10)
+	got10, exists10 := assignmentMap[vIdx10]
+	assert.Equal(t, true, exists10)
+	require.NotNil(t, got10)
+	assert.Equal(t, start+2, got10.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(1), got10.CommitteeIndex)
+	assert.Equal(t, uint64(2), got10.CommitteeLength)
+	assert.Equal(t, uint64(0), got10.ValidatorCommitteeIndex)
+
+	// Test non-existent validator
+	_, exists99 := assignmentMap[primitives.ValidatorIndex(99)]
+	assert.Equal(t, false, exists99)
+
+	// Verify that we get the same results as the linear search
+	for _, committees := range bySlot {
+		for _, committee := range committees {
+			for _, validatorIdx := range committee {
+				linearResult := helpers.AssignmentForValidator(bySlot, start, validatorIdx)
+				mapResult, mapExists := assignmentMap[validatorIdx]
+				assert.Equal(t, true, mapExists)
+				require.DeepEqual(t, linearResult, mapResult)
+			}
+		}
+	}
+}
+
+func TestGetValidatorAssignment_WithAssignmentMap(t *testing.T) {
+	start := primitives.Slot(100)
+	bySlot := [][][]primitives.ValidatorIndex{
+		{{1, 2, 3}},
+		{{4, 5, 6}},
+	}
+
+	// Test with pre-built assignment map (large request scenario)
+	meta := &metadata{
+		startSlot:              start,
+		committeesBySlot:       bySlot,
+		validatorAssignmentMap: buildValidatorAssignmentMap(bySlot, start),
+	}
+
+	vs := &Server{}
+	
+	// Test existing validator (validator 2 is at position 1 in the committee, not position 2)
+	assignment := vs.getValidatorAssignment(meta, primitives.ValidatorIndex(2))
+	require.NotNil(t, assignment)
+	assert.Equal(t, start, assignment.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(0), assignment.CommitteeIndex)
+	assert.Equal(t, uint64(1), assignment.ValidatorCommitteeIndex)
+
+	// Test non-existent validator should return empty assignment
+	assignment = vs.getValidatorAssignment(meta, primitives.ValidatorIndex(99))
+	require.NotNil(t, assignment)
+	assert.Equal(t, primitives.Slot(0), assignment.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(0), assignment.CommitteeIndex)
+}
+
+func TestGetValidatorAssignment_WithoutAssignmentMap(t *testing.T) {
+	start := primitives.Slot(100)
+	bySlot := [][][]primitives.ValidatorIndex{
+		{{1, 2, 3}},
+		{{4, 5, 6}},
+	}
+
+	// Test without assignment map (small request scenario)
+	meta := &metadata{
+		startSlot:              start,
+		committeesBySlot:       bySlot,
+		validatorAssignmentMap: nil, // No map - should use linear search
+	}
+
+	vs := &Server{}
+	
+	// Test existing validator
+	assignment := vs.getValidatorAssignment(meta, primitives.ValidatorIndex(5))
+	require.NotNil(t, assignment)
+	assert.Equal(t, start+1, assignment.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(0), assignment.CommitteeIndex)
+	assert.Equal(t, uint64(1), assignment.ValidatorCommitteeIndex)
+
+	// Test non-existent validator should return empty assignment
+	assignment = vs.getValidatorAssignment(meta, primitives.ValidatorIndex(99))
+	require.NotNil(t, assignment)
+	assert.Equal(t, primitives.Slot(0), assignment.AttesterSlot)
+	assert.Equal(t, primitives.CommitteeIndex(0), assignment.CommitteeIndex)
+}
+
+func TestLoadMetadata_ThresholdBehavior(t *testing.T) {
+	state, _ := util.DeterministicGenesisState(t, 128)
+	epoch := primitives.Epoch(0)
+
+	tests := []struct {
+		name                    string
+		numValidators          int
+		expectAssignmentMap    bool
+	}{
+		{
+			name:                    "Small request - below threshold",
+			numValidators:          100,
+			expectAssignmentMap:    false,
+		},
+		{
+			name:                    "Large request - at threshold",
+			numValidators:          validatorLookupThreshold,
+			expectAssignmentMap:    true,
+		},
+		{
+			name:                    "Large request - above threshold", 
+			numValidators:          validatorLookupThreshold + 1000,
+			expectAssignmentMap:    true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			meta, err := loadMetadata(t.Context(), state, epoch, tt.numValidators)
+			require.NoError(t, err)
+			require.NotNil(t, meta)
+
+			if tt.expectAssignmentMap {
+				require.NotNil(t, meta.validatorAssignmentMap, "Expected assignment map to be built for large requests")
+				assert.Equal(t, true, len(meta.validatorAssignmentMap) > 0, "Assignment map should not be empty")
+			} else {
+				// For small requests, the map should be nil (not initialized)
+				if meta.validatorAssignmentMap != nil {
+					t.Errorf("Expected no assignment map for small requests, got: %v", meta.validatorAssignmentMap)
+				}
+			}
+
+			// Common fields should always be set
+			assert.Equal(t, true, meta.committeesAtSlot > 0)
+			require.NotNil(t, meta.committeesBySlot)
+			assert.Equal(t, true, len(meta.committeesBySlot) > 0)
+		})
+	}
+}

changelog/tt_opt-val-lookup.md

@@ -0,0 +1,3 @@
+### Changed
+
+- Beacon api optimize validator lookup for large batch request size.