Shared validator pubkey (#5883)

This PR allows sharing the pubkey data between validators by using a
thread-local cache for pubkey data, netting about a 400mb mem usage
reduction on holesky due to us keeping 3 permanent + several ephemeral
state copies in memory at all times and each state copy holding a full
validator.

The PR also introduces a hash cache for the key which gives ~14% speedup
for a full state `hash_tree_root` - the key makes up for a large part of
the `Validator` htr time.

Finally, the time it takes to copy a state goes down as well from ~80m
ms to ~60, for reasons similar to htr.

We use a `ptr` even if a `ref` could in theory have been used - there is
not much practical benefit to a `ref` (given it's mutable) while a `ptr`
is cheaper and easier to copy (when copying temporary states).

We could go further and cache a cooked pubkey but it turns out this is
quite intrusive - in all the relevant places, we're already using a
cooked key from the immutable validator data so there are no immediate
performance gains of doing so while managing the compressed -> cooked
key mapping would become more difficult - something for a future PR
perhaps.

Co-authored-by: Etan Kissling <etan@status.im>

AllTests-mainnet.md

@@ -700,8 +700,9 @@ OK: 1/1 Fail: 0/1 Skip: 0/1
 + KzgCommitment                                                                              OK
 + KzgProof                                                                                   OK
 + RestPublishedSignedBlockContents decoding                                                  OK
++ Validator pubkey hack                                                                      OK
 ```
-OK: 5/5 Fail: 0/5 Skip: 0/5
+OK: 6/6 Fail: 0/6 Skip: 0/6
 ## Remove keystore testing suite
 ```diff
 + Many remotes                                                                               OK
@@ -992,4 +993,4 @@ OK: 2/2 Fail: 0/2 Skip: 0/2
 OK: 9/9 Fail: 0/9 Skip: 0/9
 
 ---TOTAL---
-OK: 669/674 Fail: 0/674 Skip: 5/674
+OK: 670/675 Fail: 0/675 Skip: 5/675

beacon_chain/beacon_chain_db.nim

@@ -1120,8 +1120,9 @@ proc getStateOnlyMutableValidators(
         dstValidator = addr output.validators.data[i]
 
       assign(
-        dstValidator.pubkey,
-        immutableValidators[i].pubkey.toPubKey())
+        dstValidator.pubkeyData,
+        HashedValidatorPubKey.init(
+          immutableValidators[i].pubkey.toPubKey()))
       assign(
         dstValidator.withdrawal_credentials,
         immutableValidators[i].withdrawal_credentials)
@@ -1158,7 +1159,10 @@ proc getStateOnlyMutableValidators(
       # Bypass hash cache invalidation
       let dstValidator = addr output.validators.data[i]
 
-      assign(dstValidator.pubkey, immutableValidators[i].pubkey.toPubKey())
+      assign(
+        dstValidator.pubkeyData,
+         HashedValidatorPubKey.init(
+          immutableValidators[i].pubkey.toPubKey()))
       assign(
         dstValidator.withdrawal_credentials,
         immutableValidators[i].withdrawal_credentials)
@@ -1195,7 +1199,10 @@ proc getStateOnlyMutableValidators(
     for i in prevNumValidators ..< numValidators:
       # Bypass hash cache invalidation
       let dstValidator = addr output.validators.data[i]
-      assign(dstValidator.pubkey, immutableValidators[i].pubkey.toPubKey())
+      assign(
+        dstValidator.pubkeyData,
+        HashedValidatorPubKey.init(
+          immutableValidators[i].pubkey.toPubKey()))
       output.validators.clearCaches(i)
 
     true

beacon_chain/networking/eth2_network.nim

@@ -1276,11 +1276,15 @@ proc toPeerAddr*(r: enr.TypedRecord,
   case proto
   of tcpProtocol:
     if r.ip.isSome and r.tcp.isSome:
-      let ip = ipv4(r.ip.get)
+      let ip = IpAddress(
+        family: IpAddressFamily.IPv4,
+        address_v4: r.ip.get)
       addrs.add MultiAddress.init(ip, tcpProtocol, Port r.tcp.get)
 
     if r.ip6.isSome:
-      let ip = ipv6(r.ip6.get)
+      let ip = IpAddress(
+        family: IpAddressFamily.IPv6,
+        address_v6: r.ip6.get)
       if r.tcp6.isSome:
         addrs.add MultiAddress.init(ip, tcpProtocol, Port r.tcp6.get)
       elif r.tcp.isSome:
@@ -1290,11 +1294,15 @@ proc toPeerAddr*(r: enr.TypedRecord,
 
   of udpProtocol:
     if r.ip.isSome and r.udp.isSome:
-      let ip = ipv4(r.ip.get)
+      let ip = IpAddress(
+        family: IpAddressFamily.IPv4,
+        address_v4: r.ip.get)
       addrs.add MultiAddress.init(ip, udpProtocol, Port r.udp.get)
 
     if r.ip6.isSome:
-      let ip = ipv6(r.ip6.get)
+      let ip = IpAddress(
+        family: IpAddressFamily.IPv6,
+        address_v6: r.ip6.get)
       if r.udp6.isSome:
         addrs.add MultiAddress.init(ip, udpProtocol, Port r.udp6.get)
       elif r.udp.isSome:

beacon_chain/spec/beaconstate.nim

@@ -57,7 +57,7 @@ func get_validator_from_deposit*(deposit: DepositData):
       amount - amount mod EFFECTIVE_BALANCE_INCREMENT, MAX_EFFECTIVE_BALANCE)
 
   Validator(
-    pubkey: deposit.pubkey,
+    pubkeyData: HashedValidatorPubKey.init(deposit.pubkey),
     withdrawal_credentials: deposit.withdrawal_credentials,
     activation_eligibility_epoch: FAR_FUTURE_EPOCH,
     activation_epoch: FAR_FUTURE_EPOCH,

beacon_chain/spec/datatypes/base.nim

@@ -323,9 +323,16 @@ type
     pubkey*: CookedPubKey
     withdrawal_credentials*: Eth2Digest
 
+  HashedValidatorPubKeyItem* = object
+    key*: ValidatorPubKey
+    root*: Eth2Digest
+
+  HashedValidatorPubKey* = object
+    value*: ptr HashedValidatorPubKeyItem
+
   # https://github.com/ethereum/consensus-specs/blob/v1.4.0-beta.6/specs/phase0/beacon-chain.md#validator
   Validator* = object
-    pubkey*: ValidatorPubKey
+    pubkeyData*{.serializedFieldName: "pubkey".}: HashedValidatorPubKey
 
     withdrawal_credentials*: Eth2Digest
       ## Commitment to pubkey for withdrawals and transfers
@@ -441,7 +448,7 @@ type
     # serialized. They're represented in memory to allow in-place SSZ reading
     # and writing compatibly with the full Validator object.
 
-    pubkey* {.dontSerialize.}: ValidatorPubKey
+    pubkeyData* {.dontSerialize.}: HashedValidatorPubKey
 
     withdrawal_credentials* {.dontSerialize.}: Eth2Digest
       ## Commitment to pubkey for withdrawals
@@ -467,7 +474,7 @@ type
     # serialized. They're represented in memory to allow in-place SSZ reading
     # and writing compatibly with the full Validator object.
 
-    pubkey* {.dontSerialize.}: ValidatorPubKey
+    pubkeyData* {.dontSerialize.}: HashedValidatorPubKey
 
     withdrawal_credentials*: Eth2Digest
       ## Commitment to pubkey for withdrawals
@@ -545,6 +552,28 @@ type
 
     flags*: set[RewardFlags]
 
+func pubkey*(v: HashedValidatorPubKey): ValidatorPubKey =
+  if isNil(v.value):
+    # This should never happen but we guard against it in case a
+    # default-initialized Validator instance makes it through the other safety
+    # nets
+    ValidatorPubKey()
+  else:
+    v.value[].key
+
+template pubkey*(v: Validator): ValidatorPubKey =
+  v.pubkeyData.pubkey
+
+func hash_tree_root*(v: HashedValidatorPubKey): Eth2Digest =
+  if isNil(v.value):
+    # Safety net - have to use a constant because the general hash_tree_root
+    # function is not yet declared at this point
+    const zeroPubkeyHash = Eth2Digest.fromHex(
+      "fa324a462bcb0f10c24c9e17c326a4e0ebad204feced523eccaf346c686f06ee")
+    zeroPubkeyHash
+  else:
+    v.value[].root
+
 func getImmutableValidatorData*(validator: Validator): ImmutableValidatorData2 =
   let cookedKey = validator.pubkey.loadValid()  # `Validator` has valid key
   ImmutableValidatorData2(

beacon_chain/spec/eth2_apis/eth2_rest_serialization.nim

@@ -1217,6 +1217,17 @@ proc readValue*(reader: var JsonReader[RestJson], value: var ValidatorPubKey) {.
   else:
     reader.raiseUnexpectedValue($res.error())
 
+proc readValue*(reader: var JsonReader[RestJson], value: var HashedValidatorPubKey) {.
+     raises: [IOError, SerializationError].} =
+  var key: ValidatorPubKey
+  readValue(reader, key)
+
+  value = HashedValidatorPubKey.init(key)
+
+proc writeValue*(
+    writer: var JsonWriter[RestJson], value: HashedValidatorPubKey) {.raises: [IOError].} =
+  writeValue(writer, value.pubkey)
+
 ## BitSeq
 proc readValue*(reader: var JsonReader[RestJson], value: var BitSeq) {.
      raises: [IOError, SerializationError].} =

beacon_chain/spec/eth2_merkleization.nim

@@ -11,9 +11,11 @@
 
 import
   stew/endians2,
+  std/sets,
   ssz_serialization/[merkleization, proofs],
   ./ssz_codec
 
+from ./datatypes/base import HashedValidatorPubKeyItem
 from ./datatypes/phase0 import HashedBeaconState, SignedBeaconBlock
 from ./datatypes/altair import HashedBeaconState, SignedBeaconBlock
 from ./datatypes/bellatrix import HashedBeaconState, SignedBeaconBlock
@@ -66,3 +68,40 @@ func toDepositContractState*(merkleizer: DepositsMerkleizer): DepositContractSta
 
 func getDepositsRoot*(m: var DepositsMerkleizer): Eth2Digest =
   mixInLength(m.getFinalHash, int m.totalChunks)
+
+func hash*(v: ref HashedValidatorPubKeyItem): Hash =
+  if not isNil(v):
+    hash(v[].key)
+  else:
+    default(Hash)
+
+func `==`*(a, b: ref HashedValidatorPubKeyItem): bool =
+  if isNil(a):
+    isNil(b)
+  elif isNil(b):
+    false
+  else:
+    a[].key == b[].key
+
+func init*(T: type HashedValidatorPubKey, key: ValidatorPubKey): HashedValidatorPubKey =
+  {.noSideEffect.}:
+    var keys {.threadvar.}: HashSet[ref HashedValidatorPubKeyItem]
+
+    let
+      tmp = (ref HashedValidatorPubKeyItem)(
+        key: key,
+        root: hash_tree_root(key)
+      )
+      cached =
+        if keys.containsOrIncl(tmp):
+          try:
+            # The interface of HashSet is such that we must construct a full
+            # instance to check if it's in the set - then we can return that
+            # instace and discard the one we just created temporarily
+            keys[tmp]
+          except KeyError:
+            raiseAssert "just checked"
+        else:
+          tmp
+
+  HashedValidatorPubKey(value: addr cached[])

beacon_chain/spec/eth2_ssz_serialization.nim

@@ -71,3 +71,11 @@ func readSszBytes(
   var res: T
   readSszBytes(data, res, updateRoot)
   res
+
+proc fromSszBytes*(
+    T: type HashedValidatorPubKey, bytes: openArray[byte]
+): T {.raises: [SszError].} =
+  let
+    key = ValidatorPubKey.fromSszBytes(bytes)
+
+  HashedValidatorPubKey.init(key)

beacon_chain/spec/ssz_codec.nim

@@ -66,3 +66,5 @@ func fromSszBytes*(
   # TODO https://github.com/nim-lang/Nim/issues/21123
   let tmp = cast[ptr List[ParticipationFlags, Limit VALIDATOR_REGISTRY_LIMIT]](addr result)
   readSszValue(bytes, tmp[])
+
+template toSszType*(v: HashedValidatorPubKey): auto = toRaw(v.pubkey)

beacon_chain/statediff.nim

@@ -31,7 +31,7 @@ func applyValidatorIdentities(
     hl: auto) =
   for item in hl:
     if not validators.add Validator(
-        pubkey: item.pubkey.toPubKey(),
+        pubkeyData: HashedValidatorPubKey.init(item.pubkey.toPubKey()),
         withdrawal_credentials: item.withdrawal_credentials):
       raiseAssert "cannot readd"
 

tests/test_rest_json_serialization.nim

@@ -333,3 +333,24 @@ suite "REST JSON encoding and decoding":
       zeroBlob[] == zeroBlobRoundTrip[]
       nonzeroBlob[] == nonzeroBlobRoundTrip[]
       zeroBlob[] != nonzeroBlob[]
+
+  test "Validator pubkey hack":
+
+    let
+      encoded = """
+      {
+        "pubkey": "0x933ad9491b62059dd065b560d256d8957a8c402cc6e8d8ee7290ae11e8f7329267a8811c397529dac52ae1342ba58c95",
+        "withdrawal_credentials": "0x00f50428677c60f997aadeab24aabf7fceaef491c96a52b463ae91f95611cf71",
+        "effective_balance": "32000000000",
+        "slashed": false,
+        "activation_eligibility_epoch": "0",
+        "activation_epoch": "0",
+        "exit_epoch": "18446744073709551615",
+        "withdrawable_epoch": "18446744073709551615"
+      }"""
+
+    let validator = RestJson.decode(encoded, Validator)
+    check:
+      validator.pubkey == ValidatorPubKey.fromHex(
+        "0x933ad9491b62059dd065b560d256d8957a8c402cc6e8d8ee7290ae11e8f7329267a8811c397529dac52ae1342ba58c95")[]
+      validator.exit_epoch == FAR_FUTURE_EPOCH

vendor/nim-json-serialization

@@ -1 +1 @@
-Subproject commit d9394dc7286064902d825bbc1203d03d7218633a
+Subproject commit c869dae884336e1bca134ccb8ea1a37517d16a29

vendor/nim-ssz-serialization

@@ -1 +1 @@
-Subproject commit 66de36a9ecc67c98ee858faf555b8a8dd2ea2b5f
+Subproject commit 9f9c08b9a748b13942594ca0f075ff9dbaaf9bb9