Add DiscoveryConfig to tune routing table ip limits and bitPerHops (#471

)
status-im · Feb 2, 2022 · 41edd4a · 41edd4a
1 parent cf95b68
commit 41edd4a
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 55 deletions.
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
@@ -116,6 +116,10 @@ const
   ## call
 
 type
+  DiscoveryConfig* = object
+    tableIpLimits*: TableIpLimits
+    bitsPerHop*: int
+
   Protocol* = ref object
     transp: DatagramTransport
     localNode*: Node
@@ -148,6 +152,11 @@ type
 
   DiscResult*[T] = Result[T, cstring]
 
+const
+  defaultDiscoveryConfig* = DiscoveryConfig(
+    tableIpLimits: DefaultTableIpLimits,
+    bitsPerHop: DefaultBitsPerHop)
+
 proc addNode*(d: Protocol, node: Node): bool =
   ## Add `Node` to discovery routing table.
   ##
@@ -880,18 +889,32 @@ proc ipMajorityLoop(d: Protocol) {.async.} =
   except CancelledError:
     trace "ipMajorityLoop canceled"
 
-proc newProtocol*(privKey: PrivateKey,
-                  enrIp: Option[ValidIpAddress],
-                  enrTcpPort, enrUdpPort: Option[Port],
-                  localEnrFields: openArray[(string, seq[byte])] = [],
-                  bootstrapRecords: openArray[Record] = [],
-                  previousRecord = none[enr.Record](),
-                  bindPort: Port,
-                  bindIp = IPv4_any(),
-                  enrAutoUpdate = false,
-                  tableIpLimits = DefaultTableIpLimits,
-                  rng = newRng()):
-                  Protocol =
+func init*(
+    T: type DiscoveryConfig,
+    tableIpLimit: uint,
+    bucketIpLimit: uint,
+    bitsPerHop: int): T =
+
+  DiscoveryConfig(
+    tableIpLimits: TableIpLimits(
+      tableIpLimit: tableIpLimit,
+      bucketIpLimit: bucketIpLimit),
+    bitsPerHop: bitsPerHop
+  )
+
+proc newProtocol*(
+    privKey: PrivateKey,
+    enrIp: Option[ValidIpAddress],
+    enrTcpPort, enrUdpPort: Option[Port],
+    localEnrFields: openArray[(string, seq[byte])] = [],
+    bootstrapRecords: openArray[Record] = [],
+    previousRecord = none[enr.Record](),
+    bindPort: Port,
+    bindIp = IPv4_any(),
+    enrAutoUpdate = false,
+    config = defaultDiscoveryConfig,
+    rng = newRng()):
+    Protocol =
   # TODO: Tried adding bindPort = udpPort as parameter but that gave
   # "Error: internal error: environment misses: udpPort" in nim-beacon-chain.
   # Anyhow, nim-beacon-chain would also require some changes to support port
@@ -933,7 +956,8 @@ proc newProtocol*(privKey: PrivateKey,
     bootstrapRecords: @bootstrapRecords,
     ipVote: IpVote.init(),
     enrAutoUpdate: enrAutoUpdate,
-    routingTable: RoutingTable.init(node, DefaultBitsPerHop, tableIpLimits, rng),
+    routingTable: RoutingTable.init(
+      node, config.bitsPerHop, config.tableIpLimits, rng),
     rng: rng)
 
 template listeningAddress*(p: Protocol): Address =

diff --git a/tests/p2p/discv5_test_helper.nim b/tests/p2p/discv5_test_helper.nim
@@ -9,26 +9,31 @@ export net
 proc localAddress*(port: int): Address =
   Address(ip: ValidIpAddress.init("127.0.0.1"), port: Port(port))
 
-proc initDiscoveryNode*(rng: ref BrHmacDrbgContext, privKey: PrivateKey,
-                        address: Address,
-                        bootstrapRecords: openArray[Record] = [],
-                        localEnrFields: openArray[(string, seq[byte])] = [],
-                        previousRecord = none[enr.Record]()):
-                        discv5_protocol.Protocol =
+proc initDiscoveryNode*(
+    rng: ref BrHmacDrbgContext,
+    privKey: PrivateKey,
+    address: Address,
+    bootstrapRecords: openArray[Record] = [],
+    localEnrFields: openArray[(string, seq[byte])] = [],
+    previousRecord = none[enr.Record]()):
+    discv5_protocol.Protocol =
   # set bucketIpLimit to allow bucket split
-  let tableIpLimits = TableIpLimits(tableIpLimit: 1000,  bucketIpLimit: 24)
+  let config = DiscoveryConfig.init(1000, 24, 5)
 
-  result = newProtocol(privKey,
-                       some(address.ip),
-                       some(address.port), some(address.port),
-                       bindPort = address.port,
-                       bootstrapRecords = bootstrapRecords,
-                       localEnrFields = localEnrFields,
-                       previousRecord = previousRecord,
-                       tableIpLimits = tableIpLimits,
-                       rng = rng)
+  let protocol = newProtocol(
+    privKey,
+    some(address.ip),
+    some(address.port), some(address.port),
+    bindPort = address.port,
+    bootstrapRecords = bootstrapRecords,
+    localEnrFields = localEnrFields,
+    previousRecord = previousRecord,
+    config = config,
+    rng = rng)
 
-  result.open()
+  protocol.open()
+
+  protocol
 
 proc nodeIdInNodes*(id: NodeId, nodes: openArray[Node]): bool =
   for n in nodes:

diff --git a/tests/utp/test_discv5_protocol.nim b/tests/utp/test_discv5_protocol.nim
@@ -14,31 +14,8 @@ import
   ../../eth/p2p/discoveryv5/[enr, node, routing_table],
   ../../eth/p2p/discoveryv5/protocol as discv5_protocol,
   ../../eth/utp/utp_discv5_protocol,
-  ../../eth/keys
-
-proc localAddress*(port: int): Address =
-  Address(ip: ValidIpAddress.init("127.0.0.1"), port: Port(port))
-
-proc initDiscoveryNode*(rng: ref BrHmacDrbgContext,
-    privKey: PrivateKey,
-    address: Address,
-    bootstrapRecords: openArray[Record] = [],
-    localEnrFields: openArray[(string, seq[byte])] = [],
-    previousRecord = none[enr.Record]()): discv5_protocol.Protocol =
-  # set bucketIpLimit to allow bucket split
-  let tableIpLimits = TableIpLimits(tableIpLimit: 1000,  bucketIpLimit: 24)
-
-  result = newProtocol(privKey,
-    some(address.ip),
-    some(address.port), some(address.port),
-    bindPort = address.port,
-    bootstrapRecords = bootstrapRecords,
-    localEnrFields = localEnrFields,
-    previousRecord = previousRecord,
-    tableIpLimits = tableIpLimits,
-    rng = rng)
-
-  result.open()
+  ../../eth/keys,
+  ../p2p/discv5_test_helper
 
 proc generateByteArray(rng: var BrHmacDrbgContext, length: int): seq[byte] =
   var bytes = newSeq[byte](length)