Simplify data size calculation

vacp2p · Aug 2, 2023 · 7ea1c8b · 7ea1c8b
1 parent 9403349
commit 7ea1c8b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 71 deletions.
diff --git a/libp2p/protocols/pubsub/gossipsub.nim b/libp2p/protocols/pubsub/gossipsub.nim
@@ -370,6 +370,9 @@ proc validateAndRelay(g: GossipSub,
   except CatchableError as exc:
     info "validateAndRelay failed", msg=exc.msg
 
+proc dataAndTopicsIdSize(msgs: seq[Message]): int =
+  msgs.mapIt(it.data.len + it.topicIds.mapIt(it.len).foldl(a + b, 0)).foldl(a + b, 0)
+
 proc rateLimit*(g: GossipSub, peer: PubSubPeer, rpcMsgOpt: Opt[RPCMsg], msgSize: int) {.raises:[PeerRateLimitError].} =
   # In this way we count even ignored fields by protobuf
   var rmsg = rpcMsgOpt.valueOr:
@@ -381,12 +384,7 @@ proc rateLimit*(g: GossipSub, peer: PubSubPeer, rpcMsgOpt: Opt[RPCMsg], msgSize:
     if g.verifySignature:
       byteSize(rmsg.messages)
     else:
-      rmsg.messages.mapIt(it.data.len).foldl(a + b, 0) +
-      rmsg.messages
-        .mapIt(
-          it.topicIds.mapIt(it.len).foldl(a + b, 0)
-        )
-        .foldl(a + b, 0)
+      dataAndTopicsIdSize(rmsg.messages)
 
   var uselessAppBytesNum =  msgSize - usefulMsgBytesNum
   rmsg.control.withValue(control):

diff --git a/libp2p/protocols/pubsub/rpc/messages.nim b/libp2p/protocols/pubsub/rpc/messages.nim
@@ -145,34 +145,3 @@ proc byteSize*(iwant: seq[ControlIWant]): int =
     for msgId in item.messageIds:
       total += msgId.len
   return total
-
-proc byteSize*(msg: RPCMsg): int =
-  var total = 0
-
-  for sub in msg.subscriptions:
-    total += sub.topic.len
-    total += sizeof(sub.subscribe)
-
-  for msg in msg.messages:
-    total += byteSize(msg)
-
-  if msg.control.isSome:
-    let ctrl = msg.control.get()
-    total += byteSize(ctrl.ihave)
-
-    total += byteSize(ctrl.iwant)
-
-    for item in ctrl.graft:
-      total += item.topicId.len
-
-    for item in ctrl.prune:
-      total += item.topicId.len
-      total += sizeof(item.backoff)
-      for peer in item.peers:
-        total += peer.peerId.len
-        total += peer.signedPeerRecord.len * sizeof(byte)
-
-  total += msg.ping.len * sizeof(byte)
-  total += msg.pong.len * sizeof(byte)
-
-  return total
diff --git a/tests/pubsub/testmessage.nim b/tests/pubsub/testmessage.nim
@@ -83,37 +83,3 @@ suite "Message":
       key: @[1'u8],  # 1 byte
       topicIds: @["abc", "defgh"]  # 3 + 5 = 8 bytes
     )
-
-    check byteSize(msg) == 3 + 1 + 1 + 8  # Total: 13 bytes
-
-  test "byteSize for RPCMsg":
-    var msg = RPCMsg(
-      subscriptions: @[
-        SubOpts(topic: "abc", subscribe: true),
-        SubOpts(topic: "def", subscribe: false)
-      ],  # 3 + 3 + 2 * sizeof(bool) bytes
-      messages: @[
-        Message(fromPeer: PeerId(data: @[]), data: @[1'u8, 2, 3], seqno: @[1'u8], signature: @[], key: @[1'u8], topicIds: @["abc", "defgh"]),
-        Message(fromPeer: PeerId(data: @[]), data: @[], seqno: @[], signature: @[], key: @[], topicIds: @["abc"])
-      ],  # byteSize: 13 + 3 = 16 bytes
-      control: some(ControlMessage(
-        ihave: @[
-          ControlIHave(topicId: "ghi", messageIds: @[@[1'u8, 2, 3]])
-        ],  # 3 + 3 bytes
-        iwant: @[
-          ControlIWant(messageIds: @[@[1'u8, 2]])
-        ],  # 2 bytes
-        graft: @[
-          ControlGraft(topicId: "jkl")
-        ],  # 3 bytes
-        prune: @[
-          ControlPrune(topicId: "mno", peers: @[PeerInfoMsg(peerId: PeerId(data: @[]), signedPeerRecord: @[])], backoff: 1)
-        ]  # 3 + sizeof(uint64) bytes
-      )),
-      ping: @[],  # Empty seq[byte]
-      pong: @[]  # Empty seq[byte]
-    )
-
-    let boolSize = sizeof(bool)
-    let uint64Size = sizeof(uint64)
-    check byteSize(msg) == (3 + 3 + 2 * boolSize) + 16 + (3 + 3 + 2 + 3 + 3 + uint64Size)