fix: improve state serialization

packages/beacon-node/src/chain/stateCache/persistentCheckpointsCache.ts

@@ -640,7 +640,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
             this.metrics?.statePersistSecFromSlot.observe(this.clock?.secFromSlot(this.clock?.currentSlot ?? 0) ?? 0);
             const cpPersist = {epoch: epoch, root: epochBoundaryRoot};
             {
-              const timer = this.metrics?.statePersistDuration.startTimer();
+              const timer = this.metrics?.stateSerializeDuration.startTimer();
               // automatically free the buffer pool after this scope
               using stateBytesWithKey = this.serializeState(state);
               let stateBytes = stateBytesWithKey?.buffer;
@@ -649,8 +649,8 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
                 this.metrics?.persistedStateAllocCount.inc();
                 stateBytes = state.serialize();
               }
-              persistedKey = await this.datastore.write(cpPersist, stateBytes);
               timer?.();
+              persistedKey = await this.datastore.write(cpPersist, stateBytes);
             }
             persistCount++;
             this.logger.verbose("Pruned checkpoint state from memory and persisted to disk", {
@@ -723,7 +723,7 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
       if (bufferWithKey) {
         const stateBytes = bufferWithKey.buffer;
         const dataView = new DataView(stateBytes.buffer, stateBytes.byteOffset, stateBytes.byteLength);
-        state.type.tree_serializeToBytes({uint8Array: stateBytes, dataView}, 0, state.node);
+        state.serializeToBytes({uint8Array: stateBytes, dataView}, 0);
         return bufferWithKey;
       }
     }
@@ -736,18 +736,16 @@ export class PersistentCheckpointStateCache implements CheckpointStateCache {
    *   - As monitored on holesky as of Jan 2024, it helps save ~500ms state reload time (4.3s vs 3.8s)
    *   - Also `serializeState.test.ts` perf test shows a lot of differences allocating validators bytes once vs every time,
    * This is 2x - 3x faster than allocating memory every time.
-   * TODO: consider serializing validators manually like in `serializeState.test.ts` perf test, this could be 3x faster than this
    */
   private serializeStateValidators(state: CachedBeaconStateAllForks): BufferWithKey | null {
-    // const validatorsSszTimer = this.metrics?.stateReloadValidatorsSszDuration.startTimer();
     const type = state.type.fields.validators;
     const size = type.tree_serializedSize(state.validators.node);
     if (this.bufferPool) {
       const bufferWithKey = this.bufferPool.alloc(size);
       if (bufferWithKey) {
         const validatorsBytes = bufferWithKey.buffer;
         const dataView = new DataView(validatorsBytes.buffer, validatorsBytes.byteOffset, validatorsBytes.byteLength);
-        type.tree_serializeToBytes({uint8Array: validatorsBytes, dataView}, 0, state.validators.node);
+        state.validators.serializeToBytes({uint8Array: validatorsBytes, dataView}, 0);
         return bufferWithKey;
       }
     }

packages/beacon-node/src/metrics/metrics/lodestar.ts

@@ -1195,9 +1195,9 @@ export function createLodestarMetrics(
         help: "Histogram of cloned count per state every time state.clone() is called",
         buckets: [1, 2, 5, 10, 50, 250],
       }),
-      statePersistDuration: register.histogram({
-        name: "lodestar_cp_state_cache_state_persist_seconds",
-        help: "Histogram of time to persist state to db",
+      stateSerializeDuration: register.histogram({
+        name: "lodestar_cp_state_cache_state_serialize_seconds",
+        help: "Histogram of time to serialize state to db",
         buckets: [0.1, 0.5, 1, 2, 3, 4],
       }),
       statePruneFromMemoryCount: register.gauge({

packages/beacon-node/test/e2e/chain/stateCache/nHistoricalStates.test.ts

@@ -399,10 +399,10 @@ describe(
           )?.value
         ).toEqual(reloadCount);
 
-        const persistMetricValues = await (followupBn.metrics?.cpStateCache.statePersistDuration as Histogram).get();
+        const stateSszMetricValues = await (followupBn.metrics?.cpStateCache.stateSerializeDuration as Histogram).get();
         expect(
-          persistMetricValues?.values.find(
-            (value) => value.metricName === "lodestar_cp_state_cache_state_persist_seconds_count"
+          stateSszMetricValues?.values.find(
+            (value) => value.metricName === "lodestar_cp_state_cache_state_serialize_seconds_count"
           )?.value
         ).toEqual(persistCount);
 

packages/state-transition/test/perf/util/serializeState.test.ts

@@ -1,6 +1,6 @@
 import {itBench, setBenchOpts} from "@dapplion/benchmark";
 import {ssz} from "@lodestar/types";
-import {generatePerfTestCachedStateAltair} from "../util.js";
+import {cachedStateAltairPopulateCaches, generatePerfTestCachedStateAltair} from "../util.js";
 
 /**
  * This shows different statistics between allocating memory once vs every time.
@@ -13,8 +13,9 @@ describe.skip("serialize state and validators", function () {
     // increasing this may have different statistics due to gc time
     minMs: 60_000,
   });
-  const valicatorCount = 1_500_000;
-  const seedState = generatePerfTestCachedStateAltair({vc: 1_500_000, goBackOneSlot: false});
+  // change to 1_700_000 for holesky size
+  const valicatorCount = 20_000;
+  const seedState = generatePerfTestCachedStateAltair({vc: valicatorCount, goBackOneSlot: false});
 
   /**
    * Allocate memory every time, on a Mac M1:
@@ -35,6 +36,17 @@ describe.skip("serialize state and validators", function () {
     },
   });
 
+  // now cache nodes
+  cachedStateAltairPopulateCaches(seedState);
+
+  itBench({
+    id: `serialize state ${valicatorCount} validators using new serializeToBytes() method`,
+    fn: () => {
+      stateBytes.fill(0);
+      seedState.serializeToBytes({uint8Array: stateBytes, dataView: stateDataView}, 0);
+    },
+  });
+
   itBench({
     id: `serialize altair state ${valicatorCount} validators`,
     fn: () => {
@@ -80,7 +92,7 @@ describe.skip("serialize state and validators", function () {
    * this is 3x faster than the previous approach.
    */
   const NUMBER_2_POW_32 = 2 ** 32;
-  const output = new Uint8Array(121 * 1_500_000);
+  const output = new Uint8Array(121 * valicatorCount);
   const dataView = new DataView(output.buffer, output.byteOffset, output.byteLength);
   // this caches validators nodes which is what happen after we run a state transition
   const validators = seedState.validators.getAllReadonlyValues();
@@ -119,4 +131,11 @@ describe.skip("serialize state and validators", function () {
       }
     },
   });
+
+  itBench({
+    id: `serialize ${valicatorCount} validators from state `,
+    fn: () => {
+      seedState.validators.serializeToBytes({uint8Array: output, dataView}, 0);
+    },
+  });
 });

packages/types/src/phase0/sszTypes.ts

@@ -2,7 +2,6 @@ import {
   BitListType,
   BitVectorType,
   ContainerType,
-  ContainerNodeStructType,
   ListBasicType,
   ListCompositeType,
   VectorBasicType,
@@ -29,15 +28,14 @@ import {
   VALIDATOR_REGISTRY_LIMIT,
 } from "@lodestar/params";
 import * as primitiveSsz from "../primitive/sszTypes.js";
+import {ValidatorNodeStruct} from "./validator.js";
 
 const {
-  Boolean,
   Bytes32,
   UintNum64,
   UintBn64,
   Slot,
   Epoch,
-  EpochInf,
   CommitteeIndex,
   ValidatorIndex,
   Root,
@@ -226,21 +224,6 @@ export const HistoricalBatchRoots = new ContainerType(
   {typeName: "HistoricalBatchRoots", jsonCase: "eth2"}
 );
 
-export const ValidatorContainer = new ContainerType(
-  {
-    pubkey: BLSPubkey,
-    withdrawalCredentials: Bytes32,
-    effectiveBalance: UintNum64,
-    slashed: Boolean,
-    activationEligibilityEpoch: EpochInf,
-    activationEpoch: EpochInf,
-    exitEpoch: EpochInf,
-    withdrawableEpoch: EpochInf,
-  },
-  {typeName: "Validator", jsonCase: "eth2"}
-);
-
-export const ValidatorNodeStruct = new ContainerNodeStructType(ValidatorContainer.fields, ValidatorContainer.opts);
 // The main Validator type is the 'ContainerNodeStructType' version
 export const Validator = ValidatorNodeStruct;
 

packages/types/src/phase0/validator.ts

@@ -0,0 +1,76 @@
+import {ByteViews, ContainerNodeStructType, ValueOfFields} from "@chainsafe/ssz";
+import * as primitiveSsz from "../primitive/sszTypes.js";
+
+const {Boolean, Bytes32, UintNum64, BLSPubkey, EpochInf} = primitiveSsz;
+
+// this is to work with uint32, see https://github.com/ChainSafe/ssz/blob/ssz-v0.15.1/packages/ssz/src/type/uint.ts
+const NUMBER_2_POW_32 = 2 ** 32;
+
+/*
+ * Below constants are respective to their ssz type in `ValidatorType`.
+ */
+const UINT32_SIZE = 4;
+const PUBKEY_SIZE = 48;
+const WITHDRAWAL_CREDENTIALS_SIZE = 32;
+const SLASHED_SIZE = 1;
+
+export const ValidatorType = {
+  pubkey: BLSPubkey,
+  withdrawalCredentials: Bytes32,
+  effectiveBalance: UintNum64,
+  slashed: Boolean,
+  activationEligibilityEpoch: EpochInf,
+  activationEpoch: EpochInf,
+  exitEpoch: EpochInf,
+  withdrawableEpoch: EpochInf,
+};
+
+/**
+ * Improve serialization performance for state.validators.serialize();
+ */
+export class ValidatorNodeStructType extends ContainerNodeStructType<typeof ValidatorType> {
+  constructor() {
+    super(ValidatorType, {typeName: "Validator", jsonCase: "eth2"});
+  }
+
+  value_serializeToBytes(
+    {uint8Array: output, dataView}: ByteViews,
+    offset: number,
+    validator: ValueOfFields<typeof ValidatorType>
+  ): number {
+    output.set(validator.pubkey, offset);
+    offset += PUBKEY_SIZE;
+    output.set(validator.withdrawalCredentials, offset);
+    offset += WITHDRAWAL_CREDENTIALS_SIZE;
+    const {effectiveBalance, activationEligibilityEpoch, activationEpoch, exitEpoch, withdrawableEpoch} = validator;
+    // effectiveBalance is UintNum64
+    dataView.setUint32(offset, effectiveBalance & 0xffffffff, true);
+    offset += UINT32_SIZE;
+    dataView.setUint32(offset, (effectiveBalance / NUMBER_2_POW_32) & 0xffffffff, true);
+    offset += UINT32_SIZE;
+    output[offset] = validator.slashed ? 1 : 0;
+    offset += SLASHED_SIZE;
+    offset = writeEpochInf(dataView, offset, activationEligibilityEpoch);
+    offset = writeEpochInf(dataView, offset, activationEpoch);
+    offset = writeEpochInf(dataView, offset, exitEpoch);
+    offset = writeEpochInf(dataView, offset, withdrawableEpoch);
+
+    return offset;
+  }
+}
+
+function writeEpochInf(dataView: DataView, offset: number, value: number): number {
+  if (value === Infinity) {
+    dataView.setUint32(offset, 0xffffffff, true);
+    offset += UINT32_SIZE;
+    dataView.setUint32(offset, 0xffffffff, true);
+    offset += UINT32_SIZE;
+  } else {
+    dataView.setUint32(offset, value & 0xffffffff, true);
+    offset += UINT32_SIZE;
+    dataView.setUint32(offset, (value / NUMBER_2_POW_32) & 0xffffffff, true);
+    offset += UINT32_SIZE;
+  }
+  return offset;
+}
+export const ValidatorNodeStruct = new ValidatorNodeStructType();

packages/types/test/unit/phase0/sszTypes.test.ts

@@ -0,0 +1,31 @@
+import {ContainerType} from "@chainsafe/ssz";
+import {describe, it, expect} from "vitest";
+import {ssz} from "../../../src/index.js";
+import {ValidatorType} from "../../../src/phase0/validator.js";
+
+const ValidatorContainer = new ContainerType(ValidatorType, {typeName: "Validator", jsonCase: "eth2"});
+
+describe("Validator ssz types", function () {
+  it("should serialize to the same value", () => {
+    const seedValidator = {
+      activationEligibilityEpoch: 10,
+      activationEpoch: 11,
+      exitEpoch: Infinity,
+      slashed: false,
+      withdrawableEpoch: 13,
+      pubkey: Buffer.alloc(48, 100),
+      withdrawalCredentials: Buffer.alloc(32, 100),
+    };
+
+    const validators = [
+      {...seedValidator, effectiveBalance: 31000000000, slashed: false},
+      {...seedValidator, effectiveBalance: 32000000000, slashed: true},
+    ];
+
+    for (const validator of validators) {
+      const serialized = ValidatorContainer.serialize(validator);
+      const serialized2 = ssz.phase0.Validator.serialize(validator);
+      expect(serialized).toEqual(serialized2);
+    }
+  });
+});