chore: add Cancelled exception to GetConn

vacp2p · Feb 19, 2025 · 63c2533 · 63c2533
1 parent a978660
commit 63c2533
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 4 deletions.
diff --git a/libp2p/protocols/pubsub/pubsub.nim b/libp2p/protocols/pubsub/pubsub.nim
@@ -359,9 +359,11 @@ method getOrCreatePeer*(
       peer[].codec = protoNegotiated
     return peer[]
 
-  proc getConn(): Future[Connection] {.async: (raises: [GetConnDialError]).} =
+  proc getConn(): Future[Connection] {.async: (raises: [CancelledError, GetConnDialError]).} =
     try:
       return await p.switch.dial(peerId, protosToDial)
+    except CancelledError as exc:
+      raise exc
     except CatchableError as e:
       raise (ref GetConnDialError)(parent: e)
 

diff --git a/libp2p/protocols/pubsub/pubsubpeer.nim b/libp2p/protocols/pubsub/pubsubpeer.nim
@@ -81,7 +81,7 @@ type
   PubSubPeerEvent* = object
     kind*: PubSubPeerEventKind
 
-  GetConn* = proc(): Future[Connection] {.async: (raises: [GetConnDialError]).}
+  GetConn* = proc(): Future[Connection] {.async: (raises: [CancelledError, GetConnDialError]).}
   DropConn* = proc(peer: PubSubPeer) {.gcsafe, raises: [].}
     # have to pass peer as it's unknown during init
   OnEvent* = proc(peer: PubSubPeer, event: PubSubPeerEvent) {.gcsafe, raises: [].}

diff --git a/tests/pubsub/utils.nim b/tests/pubsub/utils.nim
@@ -27,9 +27,11 @@ randomize()
 type TestGossipSub* = ref object of GossipSub
 
 proc getPubSubPeer*(p: TestGossipSub, peerId: PeerId): PubSubPeer =
-  proc getConn(): Future[Connection] {.async: (raises: [GetConnDialError]).} =
-    try:
+  proc getConn(): Future[Connection] {.async: (raises: [CancelledError, GetConnDialError]).} =
+    try: 
       return await p.switch.dial(peerId, GossipSubCodec_12)
+    except CancelledError as exc:
+      raise exc
     except CatchableError as e:
       raise (ref GetConnDialError)(parent: e)