From 70029545ea759c147c98845d92fa272b6eb21625 Mon Sep 17 00:00:00 2001
From: Agnish Ghosh <iamagnix@gmail.com>
Date: Tue, 20 Aug 2024 03:58:27 +0530
Subject: [PATCH] added support for customisable custody subnet count on boot,
 from bn startup flags

---
 beacon_chain/networking/eth2_network.nim | 16 ++++++++++++++++
 beacon_chain/nimbus_beacon_node.nim      | 13 ++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/beacon_chain/networking/eth2_network.nim b/beacon_chain/networking/eth2_network.nim
index d9ba035d7..9a9764780 100644
--- a/beacon_chain/networking/eth2_network.nim
+++ b/beacon_chain/networking/eth2_network.nim
@@ -2598,6 +2598,22 @@ proc updateStabilitySubnetMetadata*(node: Eth2Node, attnets: AttnetBits) =
   else:
     debug "Stability subnets changed; updated ENR attnets", attnets
 
+proc loadCscnetsMetadata*(node: Eth2Node, cscnets: uint64) =
+  if node.metadata.custody_subnet_count == cscnets:
+    return
+
+  node.metadata.custody_subnet_count = cscnets
+
+  let res = node.discovery.updateRecord({
+    enrCustodySubnetCountField: SSZ.encode(node.metadata.custody_subnet_count)})
+
+  if res.isErr:
+    # This should not occur in this scenario as the private key would always
+    # be the correct one and the ENR will not increase in size.
+    warn "Failed to update the ENR syncnets field", error = res.error
+  else:
+    debug "Sync committees changed; updated ENR syncnets", cscnets
+
 proc updateSyncnetsMetadata*(node: Eth2Node, syncnets: SyncnetBits) =
   # https://github.com/ethereum/consensus-specs/blob/v1.5.0-alpha.2/specs/altair/validator.md#sync-committee-subnet-stability
   if node.metadata.syncnets == syncnets:
diff --git a/beacon_chain/nimbus_beacon_node.nim b/beacon_chain/nimbus_beacon_node.nim
index 6f71e1307..58f566bf5 100644
--- a/beacon_chain/nimbus_beacon_node.nim
+++ b/beacon_chain/nimbus_beacon_node.nim
@@ -1153,6 +1153,15 @@ proc addDenebMessageHandlers(
   for topic in dataColumnSidecarTopics(forkDigest, targetSubnets):
     node.network.subscribe(topic, basicParams)
 
+  if node.config.subscribeAllSubnets:
+    node.network.loadCscnetsMetadata(DATA_COLUMN_SIDECAR_SUBNET_COUNT.uint64)
+  elif not node.config.subscribeAllSubnets:
+    let csc = node.config.custodySubnetCount
+    if csc.isSome and csc.get < DATA_COLUMN_SIDECAR_SUBNET_COUNT:
+      node.network.loadCscnetsMetadata(csc.get.uint64)
+    else:
+      node.network.loadCscnetsMetadata(CUSTODY_REQUIREMENT.uint64)
+
 proc addElectraMessageHandlers(
     node: BeaconNode, forkDigest: ForkDigest, slot: Slot) =
   node.addDenebMessageHandlers(forkDigest, slot)
@@ -1544,7 +1553,9 @@ proc reconstructAndSendDataColumns*(node: BeaconNode) {.async.} =
         return
       notice "Data Column Reconstructed and Saved Successfully"
       let dc = data_column_sidecars.get
-      var das_workers = newSeq[Future[SendResult]](len(dc))
+      var
+        worker_count = len(dc)
+        das_workers = newSeq[Future[SendResult]](worker_count)
       for i in 0..<dc.lenu64:
         let subnet_id = compute_subnet_for_data_column_sidecar(i)
         das_workers[i] =