From 351472a2d49f190066e9ba04842b38a2d99476ee Mon Sep 17 00:00:00 2001
From: tersec <tersec@users.noreply.github.com>
Date: Wed, 20 Sep 2023 19:00:37 +0200
Subject: [PATCH] per payload-builder payload builder validator registration
 (#5431)

---
 beacon_chain/beacon_node.nim                  | 17 ++++----
 beacon_chain/validators/beacon_validators.nim | 41 ++++++++++++++-----
 2 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/beacon_chain/beacon_node.nim b/beacon_chain/beacon_node.nim
index ac3a873006..1334461927 100644
--- a/beacon_chain/beacon_node.nim
+++ b/beacon_chain/beacon_node.nim
@@ -122,23 +122,26 @@ func getPayloadBuilderAddress*(config: BeaconNodeConf): Opt[string] =
   else:
     Opt.none(string)
 
+proc getPayloadBuilderAddress*(
+    node: BeaconNode, pubkey: ValidatorPubKey): Opt[string] =
+  let defaultPayloadBuilderAddress = node.config.getPayloadBuilderAddress
+  if node.keymanagerHost.isNil:
+    defaultPayloadBuilderAddress
+  else:
+    node.keymanagerHost[].getBuilderConfig(pubkey).valueOr:
+      defaultPayloadBuilderAddress
+
 proc getPayloadBuilderClient*(
     node: BeaconNode, validator_index: uint64): RestResult[RestClientRef] =
   if not node.config.payloadBuilderEnable:
     return err "Payload builder globally disabled"
 
   let
-    defaultPayloadBuilderAddress = node.config.getPayloadBuilderAddress
     pubkey = withState(node.dag.headState):
       if validator_index >= forkyState.data.validators.lenu64:
         return err "Validator index too high"
       forkyState.data.validators.item(validator_index).pubkey
-    payloadBuilderAddress =
-      if node.keymanagerHost.isNil:
-        defaultPayloadBuilderAddress
-      else:
-        node.keymanagerHost[].getBuilderConfig(pubkey).valueOr:
-          defaultPayloadBuilderAddress
+    payloadBuilderAddress = node.getPayloadBuilderAddress(pubkey)
 
   if payloadBuilderAddress.isNone:
     return err "Payload builder disabled"
diff --git a/beacon_chain/validators/beacon_validators.nim b/beacon_chain/validators/beacon_validators.nim
index ea6fa5db5d..7f6c935827 100644
--- a/beacon_chain/validators/beacon_validators.nim
+++ b/beacon_chain/validators/beacon_validators.nim
@@ -1556,17 +1556,24 @@ proc getValidatorRegistration(
 
   return ok validatorRegistration
 
-proc registerValidators*(node: BeaconNode, epoch: Epoch) {.async.} =
-  try:
-    if not node.config.payloadBuilderEnable: return
+proc registerValidatorsPerBuilder(
+    node: BeaconNode, payloadBuilderAddress: string, epoch: Epoch,
+    attachedValidatorPubkeys: seq[ValidatorPubKey]) {.async.} =
+  const HttpOk = 200
 
-    const HttpOk = 200
-
-    let payloadBuilderClient = node.getPayloadBuilderClient(0).valueOr:
+  try:
+    let payloadBuilderClient =
+        RestClientRef.new(payloadBuilderAddress).valueOr:
       debug "Unable to initialize payload builder client while registering validators",
+        payloadBuilderAddress, epoch,
         err = error
       return
 
+    if payloadBuilderClient.isNil:
+      debug "registerValidatorsPerBuilder: got nil payload builder REST client reference",
+        payloadBuilderAddress, epoch
+      return
+
     let restBuilderStatus = awaitWithTimeout(payloadBuilderClient.checkBuilderStatus(),
                                              BUILDER_STATUS_DELAY_TOLERANCE):
       debug "Timeout when obtaining builder status"
@@ -1578,11 +1585,6 @@ proc registerValidators*(node: BeaconNode, epoch: Epoch) {.async.} =
         builderStatus = restBuilderStatus
       return
 
-    # The async aspect of signing the registrations can cause the attached
-    # validators to change during the loop.
-    let attachedValidatorPubkeys =
-      toSeq(node.attachedValidators[].validators.keys)
-
     const emptyNestedSeq = @[newSeq[SignedValidatorRegistrationV1](0)]
     # https://github.com/ethereum/builder-specs/blob/v0.3.0/specs/bellatrix/validator.md#validator-registration
     # Seed with single empty inner list to avoid special cases
@@ -1683,6 +1685,23 @@ proc registerValidators*(node: BeaconNode, epoch: Epoch) {.async.} =
     warn "registerValidators: exception",
       error = exc.msg
 
+proc registerValidators*(node: BeaconNode, epoch: Epoch) {.async.} =
+  if not node.config.payloadBuilderEnable: return
+
+  var builderKeys: Table[string, seq[ValidatorPubKey]]
+  for pubkey in node.attachedValidators[].validators.keys:
+    let payloadBuilderAddress = node.getPayloadBuilderAddress(pubkey).valueOr:
+      continue
+
+    if payloadBuilderAddress in builderKeys:
+      builderKeys[payloadBuilderAddress].add pubkey
+    else:
+      builderKeys[payloadBuilderAddress] = @[pubkey]
+
+  for payloadBuilderAddress in builderKeys.keys:
+    await node.registerValidatorsPerBuilder(
+      payloadBuilderAddress, epoch, builderKeys[payloadBuilderAddress])
+
 proc updateValidators(
     node: BeaconNode, validators: openArray[Validator]) =
   # Since validator indicies are stable, we only check the "updated" range -