chore: refactor block and state api utils (#6963)

* chore: refactor block and state api utils * Update packages/beacon-node/src/api/impl/beacon/blocks/utils.ts Co-authored-by: Nico Flaig <nflaig@protonmail.com> * Update packages/beacon-node/src/api/impl/beacon/state/utils.ts Co-authored-by: Nico Flaig <nflaig@protonmail.com> * chore: address pr comments * chore: tiny cleanup * Use more informative return types * Align block and state utils * Let server handler deserialize state bytes --------- Co-authored-by: Nico Flaig <nflaig@protonmail.com>
ChainSafe · Jul 18, 2024 · fed08fe · fed08fe
1 parent 5f37580
commit fed08fe
Show file tree

Hide file tree

Showing 7 changed files with 110 additions and 80 deletions.
diff --git a/packages/beacon-node/src/api/impl/beacon/blocks/index.ts b/packages/beacon-node/src/api/impl/beacon/blocks/index.ts
@@ -34,7 +34,7 @@ import {ApiModules} from "../../types.js";
 import {validateGossipBlock} from "../../../../chain/validation/block.js";
 import {verifyBlocksInEpoch} from "../../../../chain/blocks/verifyBlock.js";
 import {BeaconChain} from "../../../../chain/chain.js";
-import {resolveBlockId, toBeaconHeaderResponse} from "./utils.js";
+import {getBlockResponse, toBeaconHeaderResponse} from "./utils.js";
 
 type PublishBlockOpts = ImportBlockOpts;
 
@@ -371,15 +371,15 @@ export function getBeaconBlockApi({
     },
 
     async getBlockHeader({blockId}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       return {
         data: toBeaconHeaderResponse(config, block, true),
         meta: {executionOptimistic, finalized},
       };
     },
 
     async getBlockV2({blockId}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       return {
         data: block,
         meta: {
@@ -391,7 +391,7 @@ export function getBeaconBlockApi({
     },
 
     async getBlindedBlock({blockId}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       const fork = config.getForkName(block.message.slot);
       return {
         data: isForkExecution(fork)
@@ -406,7 +406,7 @@ export function getBeaconBlockApi({
     },
 
     async getBlockAttestations({blockId}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       return {
         data: Array.from(block.message.body.attestations),
         meta: {executionOptimistic, finalized},
@@ -445,7 +445,7 @@ export function getBeaconBlockApi({
       }
 
       // Slow path
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       return {
         data: {root: config.getForkTypes(block.message.slot).BeaconBlock.hashTreeRoot(block.message)},
         meta: {executionOptimistic, finalized},
@@ -464,7 +464,7 @@ export function getBeaconBlockApi({
     },
 
     async getBlobSidecars({blockId, indices}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       const blockRoot = config.getForkTypes(block.message.slot).BeaconBlock.hashTreeRoot(block.message);
 
       let {blobSidecars} = (await db.blobSidecars.get(blockRoot)) ?? {};

diff --git a/packages/beacon-node/src/api/impl/beacon/blocks/utils.ts b/packages/beacon-node/src/api/impl/beacon/blocks/utils.ts
@@ -1,7 +1,8 @@
 import {routes} from "@lodestar/api";
 import {blockToHeader} from "@lodestar/state-transition";
 import {ChainForkConfig} from "@lodestar/config";
-import {SignedBeaconBlock} from "@lodestar/types";
+import {RootHex, SignedBeaconBlock, Slot} from "@lodestar/types";
+import {IForkChoice} from "@lodestar/fork-choice";
 import {GENESIS_SLOT} from "../../../../constants/index.js";
 import {ApiError, ValidationError} from "../../errors.js";
 import {IBeaconChain} from "../../../../chain/interface.js";
@@ -22,50 +23,53 @@ export function toBeaconHeaderResponse(
   };
 }
 
-export async function resolveBlockId(
-  chain: IBeaconChain,
-  blockId: routes.beacon.BlockId
-): Promise<{block: SignedBeaconBlock; executionOptimistic: boolean; finalized: boolean}> {
-  const res = await resolveBlockIdOrNull(chain, blockId);
-  if (!res) {
-    throw new ApiError(404, `No block found for id '${blockId}'`);
-  }
-
-  return res;
-}
-
-async function resolveBlockIdOrNull(
-  chain: IBeaconChain,
-  blockId: routes.beacon.BlockId
-): Promise<{block: SignedBeaconBlock; executionOptimistic: boolean; finalized: boolean} | null> {
+export function resolveBlockId(forkChoice: IForkChoice, blockId: routes.beacon.BlockId): RootHex | Slot {
   blockId = String(blockId).toLowerCase();
   if (blockId === "head") {
-    return chain.getBlockByRoot(chain.forkChoice.getHead().blockRoot);
+    return forkChoice.getHead().blockRoot;
   }
 
   if (blockId === "genesis") {
-    return chain.getCanonicalBlockAtSlot(GENESIS_SLOT);
+    return GENESIS_SLOT;
   }
 
   if (blockId === "finalized") {
-    return chain.getCanonicalBlockAtSlot(chain.forkChoice.getFinalizedBlock().slot);
+    return forkChoice.getFinalizedBlock().blockRoot;
   }
 
   if (blockId === "justified") {
-    return chain.getBlockByRoot(chain.forkChoice.getJustifiedBlock().blockRoot);
+    return forkChoice.getJustifiedBlock().blockRoot;
   }
 
   if (blockId.startsWith("0x")) {
     if (!rootHexRegex.test(blockId)) {
       throw new ValidationError(`Invalid block id '${blockId}'`, "blockId");
     }
-    return chain.getBlockByRoot(blockId);
+    return blockId;
   }
 
   // block id must be slot
   const blockSlot = parseInt(blockId, 10);
   if (isNaN(blockSlot) && isNaN(blockSlot - 0)) {
     throw new ValidationError(`Invalid block id '${blockId}'`, "blockId");
   }
-  return chain.getCanonicalBlockAtSlot(blockSlot);
+  return blockSlot;
+}
+
+export async function getBlockResponse(
+  chain: IBeaconChain,
+  blockId: routes.beacon.BlockId
+): Promise<{block: SignedBeaconBlock; executionOptimistic: boolean; finalized: boolean}> {
+  const rootOrSlot = resolveBlockId(chain.forkChoice, blockId);
+
+  const res =
+    typeof rootOrSlot === "string"
+      ? await chain.getBlockByRoot(rootOrSlot)
+      : await chain.getCanonicalBlockAtSlot(rootOrSlot);
+
+  if (!res) {
+    throw new ApiError(404, `No block found for id '${blockId}'`);
+  }
+
+  return res;
 }
diff --git a/packages/beacon-node/src/api/impl/beacon/rewards/index.ts b/packages/beacon-node/src/api/impl/beacon/rewards/index.ts
@@ -1,14 +1,14 @@
 import {routes} from "@lodestar/api";
 import {ApplicationMethods} from "@lodestar/api/server";
 import {ApiModules} from "../../types.js";
-import {resolveBlockId} from "../blocks/utils.js";
+import {getBlockResponse} from "../blocks/utils.js";
 
 export function getBeaconRewardsApi({
   chain,
 }: Pick<ApiModules, "chain">): ApplicationMethods<routes.beacon.rewards.Endpoints> {
   return {
     async getBlockRewards({blockId}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       const data = await chain.getBlockRewards(block.message);
       return {data, meta: {executionOptimistic, finalized}};
     },
@@ -17,7 +17,7 @@ export function getBeaconRewardsApi({
       return {data: rewards, meta: {executionOptimistic, finalized}};
     },
     async getSyncCommitteeRewards({blockId, validatorIds}) {
-      const {block, executionOptimistic, finalized} = await resolveBlockId(chain, blockId);
+      const {block, executionOptimistic, finalized} = await getBlockResponse(chain, blockId);
       const data = await chain.getSyncCommitteeRewards(block.message, validatorIds);
       return {data, meta: {executionOptimistic, finalized}};
     },

diff --git a/packages/beacon-node/src/api/impl/beacon/state/index.ts b/packages/beacon-node/src/api/impl/beacon/state/index.ts
@@ -15,7 +15,7 @@ import {
   filterStateValidatorsByStatus,
   getStateValidatorIndex,
   getValidatorStatus,
-  resolveStateId,
+  getStateResponse,
   toValidatorResponse,
 } from "./utils.js";
 
@@ -26,7 +26,7 @@ export function getBeaconStateApi({
   async function getState(
     stateId: routes.beacon.StateId
   ): Promise<{state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean}> {
-    return resolveStateId(chain, stateId);
+    return getStateResponse(chain, stateId);
   }
 
   return {
@@ -76,7 +76,7 @@ export function getBeaconStateApi({
     },
 
     async getStateValidators({stateId, validatorIds = [], statuses = []}) {
-      const {state, executionOptimistic, finalized} = await resolveStateId(chain, stateId);
+      const {state, executionOptimistic, finalized} = await getStateResponse(chain, stateId);
       const currentEpoch = getCurrentEpoch(state);
       const {validators, balances} = state; // Get the validators sub tree once for all the loop
       const {pubkey2index} = chain.getHeadState().epochCtx;
@@ -131,7 +131,7 @@ export function getBeaconStateApi({
     },
 
     async getStateValidator({stateId, validatorId}) {
-      const {state, executionOptimistic, finalized} = await resolveStateId(chain, stateId);
+      const {state, executionOptimistic, finalized} = await getStateResponse(chain, stateId);
       const {pubkey2index} = chain.getHeadState().epochCtx;
 
       const resp = getStateValidatorIndex(validatorId, state, pubkey2index);
@@ -152,7 +152,7 @@ export function getBeaconStateApi({
     },
 
     async getStateValidatorBalances({stateId, validatorIds = []}) {
-      const {state, executionOptimistic, finalized} = await resolveStateId(chain, stateId);
+      const {state, executionOptimistic, finalized} = await getStateResponse(chain, stateId);
 
       if (validatorIds.length) {
         const headState = chain.getHeadState();
@@ -193,7 +193,7 @@ export function getBeaconStateApi({
     },
 
     async getEpochCommittees({stateId, ...filters}) {
-      const {state, executionOptimistic, finalized} = await resolveStateId(chain, stateId);
+      const {state, executionOptimistic, finalized} = await getStateResponse(chain, stateId);
 
       const stateCached = state as CachedBeaconStateAltair;
       if (stateCached.epochCtx === undefined) {
@@ -235,7 +235,7 @@ export function getBeaconStateApi({
      */
     async getEpochSyncCommittees({stateId, epoch}) {
       // TODO: Should pick a state with the provided epoch too
-      const {state, executionOptimistic, finalized} = await resolveStateId(chain, stateId);
+      const {state, executionOptimistic, finalized} = await getStateResponse(chain, stateId);
 
       // TODO: If possible compute the syncCommittees in advance of the fork and expose them here.
       // So the validators can prepare and potentially attest the first block. Not critical tho, it's very unlikely

diff --git a/packages/beacon-node/src/api/impl/beacon/state/utils.ts b/packages/beacon-node/src/api/impl/beacon/state/utils.ts
@@ -1,53 +1,31 @@
 import {routes} from "@lodestar/api";
 import {FAR_FUTURE_EPOCH, GENESIS_SLOT} from "@lodestar/params";
 import {BeaconStateAllForks, PubkeyIndexMap} from "@lodestar/state-transition";
-import {BLSPubkey, Epoch, phase0, ValidatorIndex} from "@lodestar/types";
+import {BLSPubkey, Epoch, phase0, RootHex, Slot, ValidatorIndex} from "@lodestar/types";
 import {fromHex} from "@lodestar/utils";
-import {IBeaconChain, StateGetOpts} from "../../../../chain/index.js";
+import {IForkChoice} from "@lodestar/fork-choice";
+import {IBeaconChain} from "../../../../chain/index.js";
 import {ApiError, ValidationError} from "../../errors.js";
-import {isOptimisticBlock} from "../../../../util/forkChoice.js";
 
-export async function resolveStateId(
-  chain: IBeaconChain,
-  stateId: routes.beacon.StateId,
-  opts?: StateGetOpts
-): Promise<{state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean}> {
-  const stateRes = await resolveStateIdOrNull(chain, stateId, opts);
-  if (!stateRes) {
-    throw new ApiError(404, `No state found for id '${stateId}'`);
-  }
-
-  return stateRes;
-}
-
-async function resolveStateIdOrNull(
-  chain: IBeaconChain,
-  stateId: routes.beacon.StateId,
-  opts?: StateGetOpts
-): Promise<{state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean} | null> {
+export function resolveStateId(forkChoice: IForkChoice, stateId: routes.beacon.StateId): RootHex | Slot {
   if (stateId === "head") {
-    // TODO: This is not OK, head and headState must be fetched atomically
-    const head = chain.forkChoice.getHead();
-    const headState = chain.getHeadState();
-    return {state: headState, executionOptimistic: isOptimisticBlock(head), finalized: false};
+    return forkChoice.getHead().stateRoot;
   }
 
   if (stateId === "genesis") {
-    return chain.getStateBySlot(GENESIS_SLOT, opts);
+    return GENESIS_SLOT;
   }
 
   if (stateId === "finalized") {
-    const checkpoint = chain.forkChoice.getFinalizedCheckpoint();
-    return chain.getStateByCheckpoint(checkpoint);
+    return forkChoice.getFinalizedBlock().stateRoot;
   }
 
   if (stateId === "justified") {
-    const checkpoint = chain.forkChoice.getJustifiedCheckpoint();
-    return chain.getStateByCheckpoint(checkpoint);
+    return forkChoice.getJustifiedBlock().stateRoot;
   }
 
   if (typeof stateId === "string" && stateId.startsWith("0x")) {
-    return chain.getStateByStateRoot(stateId, opts);
+    return stateId;
   }
 
   // id must be slot
@@ -56,7 +34,45 @@ async function resolveStateIdOrNull(
     throw new ValidationError(`Invalid block id '${stateId}'`, "blockId");
   }
 
-  return chain.getStateBySlot(blockSlot, opts);
+  return blockSlot;
+}
+
+export async function getStateResponse(
+  chain: IBeaconChain,
+  stateId: routes.beacon.StateId
+): Promise<{state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean}> {
+  const rootOrSlot = resolveStateId(chain.forkChoice, stateId);
+
+  const res =
+    typeof rootOrSlot === "string"
+      ? await chain.getStateByStateRoot(rootOrSlot)
+      : await chain.getStateBySlot(rootOrSlot);
+
+  if (!res) {
+    throw new ApiError(404, `No state found for id '${stateId}'`);
+  }
+
+  return res;
+}
+
+export async function getStateResponseWithRegen(
+  chain: IBeaconChain,
+  stateId: routes.beacon.StateId
+): Promise<{state: BeaconStateAllForks | Uint8Array; executionOptimistic: boolean; finalized: boolean}> {
+  const rootOrSlot = resolveStateId(chain.forkChoice, stateId);
+
+  const res =
+    typeof rootOrSlot === "string"
+      ? await chain.getStateByStateRoot(rootOrSlot, {allowRegen: true})
+      : rootOrSlot >= chain.forkChoice.getFinalizedBlock().slot
+        ? await chain.getStateBySlot(rootOrSlot, {allowRegen: true})
+        : null; // TODO implement historical state regen
+
+  if (!res) {
+    throw new ApiError(404, `No state found for id '${stateId}'`);
+  }
+
+  return res;
 }
 
 /**

diff --git a/packages/beacon-node/src/api/impl/debug/index.ts b/packages/beacon-node/src/api/impl/debug/index.ts
@@ -1,8 +1,10 @@
 import {routes} from "@lodestar/api";
 import {ApplicationMethods} from "@lodestar/api/server";
-import {resolveStateId} from "../beacon/state/utils.js";
+import {BeaconState} from "@lodestar/types";
+import {getStateResponseWithRegen} from "../beacon/state/utils.js";
 import {ApiModules} from "../types.js";
 import {isOptimisticBlock} from "../../../util/forkChoice.js";
+import {getStateSlotFromBytes} from "../../../util/multifork.js";
 
 export function getDebugApi({
   chain,
@@ -34,11 +36,19 @@ export function getDebugApi({
     },
 
     async getStateV2({stateId}, context) {
-      const {state, executionOptimistic, finalized} = await resolveStateId(chain, stateId, {allowRegen: true});
+      const {state, executionOptimistic, finalized} = await getStateResponseWithRegen(chain, stateId);
+      let slot: number, data: Uint8Array | BeaconState;
+      if (state instanceof Uint8Array) {
+        slot = getStateSlotFromBytes(state);
+        data = state;
+      } else {
+        slot = state.slot;
+        data = context?.returnBytes ? state.serialize() : state.toValue();
+      }
       return {
-        data: context?.returnBytes ? state.serialize() : state.toValue(),
+        data,
         meta: {
-          version: config.getForkName(state.slot),
+          version: config.getForkName(slot),
           executionOptimistic,
           finalized,
         },