More efficient ancestor head retrieval for GetAttestationData (#5669)

* Use `BlockRootAtSlot` to look up historical head root

* Update test

* Typo

beacon-chain/rpc/validator/attester.go

@@ -78,27 +78,27 @@ func (vs *Server) GetAttestationData(ctx context.Context, req *ethpb.Attestation
 		return nil, status.Errorf(codes.Internal, "Could not retrieve head root: %v", err)
 	}
 
-	// In the case that we receive an attestation request after a newer state/block has been
-	// processed, we walk up the chain until state.Slot <= req.Slot to prevent producing an
-	// attestation that violates processing constraints.
-	fetchState := vs.BeaconDB.State
-	if featureconfig.Get().NewStateMgmt {
-		fetchState = vs.StateGen.StateByRoot
-	}
-	for headState.Slot() > req.Slot {
-		if ctx.Err() != nil {
-			return nil, status.Errorf(codes.Aborted, ctx.Err().Error())
-		}
-		parent := headState.ParentRoot()
-		headRoot = parent[:]
-		headState, err = fetchState(ctx, parent)
+	// In the case that we receive an attestation request after a newer state/block has been processed.
+	if headState.Slot() > req.Slot {
+		headRoot, err = helpers.BlockRootAtSlot(headState, req.Slot)
 		if err != nil {
-			return nil, status.Error(codes.Internal, err.Error())
+			return nil, status.Errorf(codes.Internal, "Could not get historical head root: %v", err)
 		}
-		if headState == nil {
-			return nil, status.Error(codes.Internal, "Failed to lookup parent state from head.")
+		if featureconfig.Get().NewStateMgmt {
+			headState, err = vs.StateGen.StateByRoot(ctx, bytesutil.ToBytes32(headRoot))
+			if err != nil {
+				return nil, status.Errorf(codes.Internal, "Could not get historical head state: %v", err)
+			}
+		} else {
+			headState, err = vs.BeaconDB.State(ctx, bytesutil.ToBytes32(headRoot))
+			if err != nil {
+				return nil, status.Errorf(codes.Internal, "Could not get historical head state: %v", err)
+			}
 		}
 	}
+	if headState == nil {
+		return nil, status.Error(codes.Internal, "Failed to lookup parent state from head.")
+	}
 
 	if helpers.CurrentEpoch(headState) < helpers.SlotToEpoch(req.Slot) {
 		headState, err = state.ProcessSlots(ctx, headState, helpers.StartSlot(helpers.SlotToEpoch(req.Slot)))

beacon-chain/rpc/validator/attester_test.go

@@ -488,6 +488,7 @@ func TestServer_GetAttestationData_HeadStateSlotGreaterThanRequestSlot(t *testin
 	blockRoots[1] = blockRoot[:]
 	blockRoots[1*params.BeaconConfig().SlotsPerEpoch] = targetRoot[:]
 	blockRoots[2*params.BeaconConfig().SlotsPerEpoch] = justifiedRoot[:]
+	blockRoots[3*params.BeaconConfig().SlotsPerEpoch] = blockRoot2[:]
 	if err := beaconState.SetBlockRoots(blockRoots); err != nil {
 		t.Fatal(err)
 	}