POC: immutablejs for CachedValidatorsBeaconState

packages/lodestar-beacon-state-transition/package.json

@@ -41,7 +41,8 @@
     "@chainsafe/lodestar-utils": "^0.12.0",
     "@chainsafe/ssz": "^0.6.13",
     "bigint-buffer": "^1.1.5",
-    "buffer-xor": "^2.0.2"
+    "buffer-xor": "^2.0.2",
+    "immutable": "4.0.0-rc.12"
   },
   "devDependencies": {
     "@chainsafe/blst": "^0.1.3",

packages/lodestar-beacon-state-transition/src/fast/block/initiateValidatorExit.ts

@@ -25,7 +25,7 @@ export function initiateValidatorExit(
   const exitEpochs = validatorExitEpochs.filter((exitEpoch) => exitEpoch !== FAR_FUTURE_EPOCH);
   exitEpochs.push(computeActivationExitEpoch(config, currentEpoch));
   let exitQueueEpoch = Math.max(...exitEpochs);
-  const exitQueueChurn = validatorExitEpochs.filter((exitEpoch) => exitEpoch === exitQueueEpoch).length;
+  const exitQueueChurn = validatorExitEpochs.filter((exitEpoch) => exitEpoch === exitQueueEpoch).size;
   if (exitQueueChurn >= getChurnLimit(config, epochCtx.currentShuffling.activeIndices.length)) {
     exitQueueEpoch += 1;
   }

packages/lodestar-beacon-state-transition/src/fast/util/epochContext.ts

@@ -143,9 +143,10 @@ export class EpochContext {
     this.previousShuffling = this.currentShuffling;
     this.currentShuffling = this.nextShuffling;
     const nextEpoch = this.currentShuffling.epoch + 1;
-    const indicesBounded: [ValidatorIndex, Epoch, Epoch][] = state
+    const indicesBounded = state
       .flatValidators()
-      .map((v, i) => [i, v.activationEpoch, v.exitEpoch]);
+      .map((v, i) => [i, v.activationEpoch, v.exitEpoch])
+      .toJS();
     this.nextShuffling = computeEpochShuffling(this.config, state, indicesBounded, nextEpoch);
     this._resetProposers(state);
   }

packages/lodestar-beacon-state-transition/src/fast/util/interface.ts

@@ -1,5 +1,6 @@
 import {IReadonlyEpochShuffling} from ".";
 import {ValidatorIndex, Slot, BeaconState, Validator} from "@chainsafe/lodestar-types";
+import {List} from "immutable";
 import {ByteVector, readOnlyForEach} from "@chainsafe/ssz";
 import {createIFlatValidator, IFlatValidator} from "./flatValidator";
 import {config} from "@chainsafe/lodestar-config/lib/presets/mainnet";
@@ -24,7 +25,7 @@ export type ReadonlyEpochContext = {
  */
 // eslint-disable-next-line @typescript-eslint/naming-convention
 export interface CachedValidatorsBeaconState extends BeaconState {
-  flatValidators(): IFlatValidator[];
+  flatValidators(): List<IFlatValidator>;
   setValidator(i: ValidatorIndex, value: Partial<IFlatValidator>): void;
   addValidator(validator: Validator): void;
   getOriginalState(): BeaconState;
@@ -39,18 +40,11 @@ export interface CachedValidatorsBeaconState extends BeaconState {
  */
 export class CachedValidatorsBeaconState {
   public _state: BeaconState;
-  private _cachedValidators: IFlatValidator[];
-  private isSynced: boolean;
+  private _cachedValidators: List<IFlatValidator>;
 
-  constructor(state: BeaconState, cachedValidators?: IFlatValidator[]) {
+  constructor(state: BeaconState, cachedValidators: List<IFlatValidator>) {
     this._state = state;
-    if (cachedValidators && cachedValidators.length > 0) {
-      this._cachedValidators = cachedValidators;
-      this.isSynced = true;
-    } else {
-      this._cachedValidators = [];
-      this.isSynced = false;
-    }
+    this._cachedValidators = cachedValidators;
   }
 
   public createProxy(): CachedValidatorsBeaconState {
@@ -62,8 +56,8 @@ export class CachedValidatorsBeaconState {
    */
   public setValidator(i: ValidatorIndex, value: Partial<IFlatValidator>): void {
     if (this._cachedValidators) {
-      const validator = this._cachedValidators[i];
-      this._cachedValidators[i] = {...validator, ...value};
+      const validator = this._cachedValidators.get(i)!;
+      this._cachedValidators = this._cachedValidators.set(i, {...validator, ...value});
     }
     const validator = this._state.validators[i];
     if (value.activationEligibilityEpoch !== undefined)
@@ -79,45 +73,39 @@ export class CachedValidatorsBeaconState {
    * Add validator to both the cache and BeaconState
    */
   public addValidator(validator: Validator): void {
-    if (this.isSynced) {
-      this._cachedValidators.push(createIFlatValidator(validator));
-    }
+    this._cachedValidators = this._cachedValidators.push(createIFlatValidator(validator));
     this._state.validators.push(validator);
   }
 
   /**
    * Loop through the cached validators, not the TreeBacked validators inside BeaconState.
    */
-  public flatValidators(): IFlatValidator[] {
-    this.sync();
+  public flatValidators(): List<IFlatValidator> {
     return this._cachedValidators;
   }
 
   public clone(): CachedValidatorsBeaconState {
     const clonedState = config.types.BeaconState.clone(this._state);
-    return cloneCachedValidatorsBeaconState(clonedState, [...this._cachedValidators]);
+    return cloneCachedValidatorsBeaconState(clonedState, this._cachedValidators);
   }
 
   public getOriginalState(): BeaconState {
     return this._state;
   }
 
-  private sync(): void {
-    if (this.isSynced) return;
-    readOnlyForEach(this._state.validators, (validator) => {
-      this._cachedValidators!.push(createIFlatValidator(validator));
-    });
-    this.isSynced = true;
-  }
 }
 
 export function createCachedValidatorsBeaconState(state: BeaconState): CachedValidatorsBeaconState {
-  return new CachedValidatorsBeaconState(state).createProxy();
+  const tmpValidators: IFlatValidator[] = [];
+  readOnlyForEach(state.validators, (validator) => {
+    tmpValidators.push(createIFlatValidator(validator));
+  });
+  return new CachedValidatorsBeaconState(state, List(tmpValidators)).createProxy();
 }
 
 function cloneCachedValidatorsBeaconState(
   state: BeaconState,
-  cachedValidators: IFlatValidator[]
+  cachedValidators: List<IFlatValidator>
 ): CachedValidatorsBeaconState {
   return new CachedValidatorsBeaconState(state, cachedValidators).createProxy();
 }

packages/lodestar/src/api/impl/validator/validator.ts

@@ -10,6 +10,7 @@ import {
   AttestationData,
   AttesterDuty,
   BeaconBlock,
+  BeaconState,
   Bytes96,
   CommitteeIndex,
   Epoch,
@@ -20,7 +21,7 @@ import {
   ValidatorIndex,
 } from "@chainsafe/lodestar-types";
 import {assert, ILogger} from "@chainsafe/lodestar-utils";
-import {readOnlyForEach} from "@chainsafe/ssz";
+import {readOnlyForEach, TreeBacked} from "@chainsafe/ssz";
 import {IAttestationJob, IBeaconChain} from "../../../chain";
 import {assembleAttestationData} from "../../../chain/factory/attestation";
 import {assembleBlock} from "../../../chain/factory/block";
@@ -81,7 +82,13 @@ export class ValidatorApi implements IValidatorApi {
       await checkSyncStatus(this.config, this.sync);
       const headRoot = this.chain.forkChoice.getHeadRoot();
       const {state, epochCtx} = await this.chain.regen.getBlockSlotState(headRoot, slot);
-      return await assembleAttestationData(epochCtx.config, state, headRoot, slot, committeeIndex);
+      return await assembleAttestationData(
+        epochCtx.config,
+        state.getOriginalState() as TreeBacked<BeaconState>,
+        headRoot,
+        slot,
+        committeeIndex
+      );
     } catch (e) {
       this.logger.warn("Failed to produce attestation data", e);
       throw e;

packages/lodestar/src/chain/blocks/stateTransition.ts

@@ -1,5 +1,5 @@
-import {byteArrayEquals, TreeBacked} from "@chainsafe/ssz";
-import {BeaconState, Slot} from "@chainsafe/lodestar-types";
+import {byteArrayEquals} from "@chainsafe/ssz";
+import {Slot} from "@chainsafe/lodestar-types";
 import {assert} from "@chainsafe/lodestar-utils";
 import {
   ZERO_HASH,
@@ -23,7 +23,6 @@ import {sleep} from "@chainsafe/lodestar-utils";
 import {IBeaconDb} from "../../db";
 import {BlockError, BlockErrorCode} from "../errors";
 import {verifySignatureSetsBatch} from "../bls";
-import {createCachedValidatorsBeaconState} from "@chainsafe/lodestar-beacon-state-transition/lib/fast/util";
 import {StateTransitionEpochContext} from "@chainsafe/lodestar-beacon-state-transition/lib/fast/util/epochContext";
 
 /**
@@ -74,7 +73,7 @@ export async function processSlotsToNearestCheckpoint(
   const postSlot = slot;
   const preEpoch = computeEpochAtSlot(config, preSlot);
   const postCtx = cloneStateCtx(stateCtx);
-  const stateTransitionState = createCachedValidatorsBeaconState(postCtx.state);
+  const stateTransitionState = postCtx.state;
   const stateTranstionEpochContext = new StateTransitionEpochContext(undefined, postCtx.epochCtx);
   for (
     let nextEpochSlot = computeStartSlotAtEpoch(config, preEpoch + 1);
@@ -180,7 +179,7 @@ export async function runStateTransition(
   // it should always have epochBalances there bc it's a checkpoint state, ie got through processEpoch
   const justifiedBalances = (await db.checkpointStateCache.get(postStateContext.state.currentJustifiedCheckpoint))
     ?.epochCtx.epochBalances;
-  forkChoice.onBlock(job.signedBlock.message, postStateContext.state, justifiedBalances);
+  forkChoice.onBlock(job.signedBlock.message, postStateContext.state.getOriginalState(), justifiedBalances);
 
   if (postSlot % SLOTS_PER_EPOCH === 0) {
     emitCheckpointEvent(emitter, postStateContext);
@@ -201,7 +200,7 @@ export async function runStateTransition(
  */
 function toTreeStateContext(stateCtx: IStateContext): ITreeStateContext {
   const treeStateCtx: ITreeStateContext = {
-    state: stateCtx.state.getOriginalState() as TreeBacked<BeaconState>,
+    state: stateCtx.state,
     epochCtx: new LodestarEpochContext(undefined, stateCtx.epochCtx),
   };
 

packages/lodestar/src/chain/chain.ts

@@ -149,7 +149,7 @@ export class BeaconChain implements IBeaconChain {
   }
   public async getHeadState(): Promise<TreeBacked<BeaconState>> {
     //head state should always have epoch ctx
-    return (await this.getHeadStateContext()).state;
+    return (await this.getHeadStateContext()).state.getOriginalState() as TreeBacked<BeaconState>;
   }
   public async getHeadEpochContext(): Promise<EpochContext> {
     //head should always have epoch ctx

packages/lodestar/src/chain/eventHandlers.ts

@@ -218,7 +218,7 @@ export async function onBlock(
   job: IBlockJob
 ): Promise<void> {
   const blockRoot = this.config.types.BeaconBlock.hashTreeRoot(block.message);
-  this.logger.debug("Block processed", {
+  this.logger.info("Block processed", {
     slot: block.message.slot,
     root: toHexString(blockRoot),
   });
@@ -308,7 +308,7 @@ export async function onErrorBlock(this: BeaconChain, err: BlockError): Promise<
     return;
   }
 
-  this.logger.debug("Block error", {}, err);
+  this.logger.error("Block error", {}, err);
   const blockRoot = this.config.types.BeaconBlock.hashTreeRoot(err.job.signedBlock.message);
 
   switch (err.type.code) {

packages/lodestar/src/chain/factory/block/index.ts

@@ -3,9 +3,9 @@
  */
 
 import {processBlock} from "@chainsafe/lodestar-beacon-state-transition/lib/fast/block";
-import {createCachedValidatorsBeaconState} from "@chainsafe/lodestar-beacon-state-transition/lib/fast/util";
 import {IBeaconConfig} from "@chainsafe/lodestar-config";
-import {BeaconBlock, Bytes96, Root, Slot} from "@chainsafe/lodestar-types";
+import {BeaconBlock, BeaconState, Bytes96, Root, Slot} from "@chainsafe/lodestar-types";
+import {TreeBacked} from "@chainsafe/ssz";
 import {ZERO_HASH} from "../../../constants";
 import {IBeaconDb} from "../../../db/api";
 import {ITreeStateContext} from "../../../db/api/beacon/stateContextCache";
@@ -30,7 +30,14 @@ export async function assembleBlock(
     proposerIndex: stateContext.epochCtx.getBeaconProposer(slot),
     parentRoot: head.blockRoot,
     stateRoot: ZERO_HASH,
-    body: await assembleBody(config, db, eth1, stateContext.state, randaoReveal, graffiti),
+    body: await assembleBody(
+      config,
+      db,
+      eth1,
+      stateContext.state.getOriginalState() as TreeBacked<BeaconState>,
+      randaoReveal,
+      graffiti
+    ),
   };
   block.stateRoot = computeNewStateRoot(config, stateContext, block);
 
@@ -44,7 +51,7 @@ export async function assembleBlock(
  */
 function computeNewStateRoot(config: IBeaconConfig, stateContext: ITreeStateContext, block: BeaconBlock): Root {
   const postState = {
-    state: createCachedValidatorsBeaconState(stateContext.state.clone()),
+    state: stateContext.state.clone(),
     epochCtx: stateContext.epochCtx.copy(),
   };
   processBlock(postState.epochCtx, postState.state, block, true);

packages/lodestar/src/chain/initState.ts

@@ -16,6 +16,7 @@ import {Eth1Provider} from "../eth1";
 import {IBeaconMetrics} from "../metrics";
 import {GenesisBuilder} from "./genesis/genesis";
 import {IGenesisResult} from "./genesis/interface";
+import {createCachedValidatorsBeaconState} from "@chainsafe/lodestar-beacon-state-transition/lib/fast/util";
 
 export async function persistGenesisResult(
   db: IBeaconDb,
@@ -138,7 +139,7 @@ export async function restoreStateCaches(
   const epochCtx = new EpochContext(config);
   epochCtx.loadState(state);
 
-  const stateCtx = {state, epochCtx};
+  const stateCtx = {state: createCachedValidatorsBeaconState(state), epochCtx};
   await Promise.all([db.stateCache.add(stateCtx), db.checkpointStateCache.add(checkpoint, stateCtx)]);
 }
 

packages/lodestar/src/db/api/beacon/stateContextCache.ts

@@ -1,10 +1,11 @@
 import {ByteVector, toHexString, TreeBacked} from "@chainsafe/ssz";
 import {BeaconState, Gwei} from "@chainsafe/lodestar-types";
 import {EpochContext} from "@chainsafe/lodestar-beacon-state-transition";
+import {CachedValidatorsBeaconState} from "@chainsafe/lodestar-beacon-state-transition/lib/fast/util";
 
 // Lodestar specifc state context
 export interface ITreeStateContext {
-  state: TreeBacked<BeaconState>;
+  state: CachedValidatorsBeaconState; // TreeBacked<BeaconState>
   epochCtx: LodestarEpochContext;
 }
 
@@ -40,7 +41,9 @@ export class StateContextCache {
   }
 
   public async add(item: ITreeStateContext): Promise<void> {
-    this.cache[toHexString(item.state.hashTreeRoot())] = this.clone(item);
+    this.cache[toHexString((item.state.getOriginalState() as TreeBacked<BeaconState>).hashTreeRoot())] = this.clone(
+      item
+    );
   }
 
   public async delete(root: ByteVector): Promise<void> {

packages/lodestar/src/tasks/tasks/archiveStates.ts

@@ -6,7 +6,8 @@ import {ITask} from "../interface";
 import {IBeaconDb} from "../../db/api";
 import {IBeaconConfig} from "@chainsafe/lodestar-config";
 import {ILogger} from "@chainsafe/lodestar-utils";
-import {Checkpoint} from "@chainsafe/lodestar-types";
+import {BeaconState, Checkpoint} from "@chainsafe/lodestar-types";
+import {TreeBacked} from "@chainsafe/ssz";
 
 export interface IArchiveStatesModules {
   db: IBeaconDb;
@@ -41,7 +42,7 @@ export class ArchiveStatesTask implements ITask {
       throw Error("No state in cache for finalized checkpoint state epoch #" + this.finalized.epoch);
     }
     const finalizedState = stateCache.state;
-    await this.db.stateArchive.add(finalizedState);
+    await this.db.stateArchive.add(finalizedState.getOriginalState() as TreeBacked<BeaconState>);
     // don't delete states before the finalized state, auto-prune will take care of it
     this.logger.info("Archive states completed", {finalizedEpoch: this.finalized.epoch});
     this.logger.profile("Archive States epoch #" + this.finalized.epoch);

yarn.lock

@@ -6868,6 +6868,11 @@ immediate@^3.2.3, immediate@~3.2.3:
   resolved "https://registry.yarnpkg.com/immediate/-/immediate-3.2.3.tgz#d140fa8f614659bd6541233097ddaac25cdd991c"
   integrity sha1-0UD6j2FGWb1lQSMwl92qwlzdmRw=
 
+[email protected]:
+  version "4.0.0-rc.12"
+  resolved "https://registry.yarnpkg.com/immutable/-/immutable-4.0.0-rc.12.tgz#ca59a7e4c19ae8d9bf74a97bdf0f6e2f2a5d0217"
+  integrity sha512-0M2XxkZLx/mi3t8NVwIm1g8nHoEmM9p9UBl/G9k4+hm0kBgOVdMV/B3CY5dQ8qG8qc80NN4gDV4HQv6FTJ5q7A==
+
 import-fresh@^2.0.0:
   version "2.0.0"
   resolved "https://registry.yarnpkg.com/import-fresh/-/import-fresh-2.0.0.tgz#d81355c15612d386c61f9ddd3922d4304822a546"