added extra tests for identifu push for concurrency cap

Author: lla-dane

libp2p/identity/identify_push/identify_push.py

@@ -176,10 +176,7 @@ async def push_identify_to_peers(
     host: IHost,
     peer_ids: set[ID] | None = None,
     observed_multiaddr: Multiaddr | None = None,
-    counter: dict[str, int] | None = None,
-    lock: trio.Lock | None = None,
-    limit: int = CONCURRENCY_LIMIT,
-) -> int:  # <-- return the max concurrency
+) -> None:
     """
     Push an identify message to multiple peers in parallel.
 
@@ -193,6 +190,4 @@ async def push_identify_to_peers(
     # limiting concurrent connections to 10
     async with trio.open_nursery() as nursery:
         for peer_id in peer_ids:
-            nursery.start_soon(limited_push, peer_id)
-
-    return counter["max"] if counter else 0
+            nursery.start_soon(push_identify_to_peer, host, peer_id, observed_multiaddr)

tests/core/identity/identify_push/test_identify_push.py

@@ -10,7 +10,6 @@
 from libp2p import (
     new_host,
 )
-from libp2p.abc import IHost
 from libp2p.crypto.secp256k1 import (
     create_new_key_pair,
 )
@@ -36,6 +35,8 @@
 )
 from tests.utils.utils import (
     create_mock_connections,
+    run_host_forever,
+    wait_until_listening,
 )
 
 logger = logging.getLogger("libp2p.identity.identify-push-test")
@@ -504,3 +505,91 @@ async def mock_push_identify_to_peer(
     assert state["max_observed"] <= CONCURRENCY_LIMIT, (
         f"Max concurrency observed: {state['max_observed']}"
     )
+
+
+@pytest.mark.trio
+async def test_all_peers_receive_identify_push_with_semaphore(security_protocol):
+    dummy_peers = []
+
+    async with host_pair_factory(security_protocol=security_protocol) as (host_a, _):
+        # Create dummy peers
+        for _ in range(50):
+            key_pair = create_new_key_pair()
+            dummy_host = new_host(key_pair=key_pair)
+            dummy_host.set_stream_handler(
+                ID_PUSH, identify_push_handler_for(dummy_host)
+            )
+            listen_addr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
+            dummy_peers.append((dummy_host, listen_addr))
+
+        async with trio.open_nursery() as nursery:
+            # Start all dummy hosts
+            for host, listen_addr in dummy_peers:
+                nursery.start_soon(run_host_forever, host, listen_addr)
+
+            # Wait for all hosts to finish setting up listeners
+            for host, _ in dummy_peers:
+                await wait_until_listening(host)
+
+            # Now connect host_a → dummy peers
+            for host, _ in dummy_peers:
+                await host_a.connect(info_from_p2p_addr(host.get_addrs()[0]))
+
+            await push_identify_to_peers(
+                host_a,
+            )
+
+            await trio.sleep(0.5)
+
+            peer_id_a = host_a.get_id()
+            for host, _ in dummy_peers:
+                dummy_peerstore = host.get_peerstore()
+                assert peer_id_a in dummy_peerstore.peer_ids()
+
+            nursery.cancel_scope.cancel()
+
+
+@pytest.mark.trio
+async def test_all_peers_receive_identify_push_with_semaphore_under_high_peer_load(
+    security_protocol,
+):
+    dummy_peers = []
+
+    async with host_pair_factory(security_protocol=security_protocol) as (host_a, _):
+        # Create dummy peers
+        # Breaking with more than 500 peers
+        # Trio have a async tasks limit of 1000
+        for _ in range(499):
+            key_pair = create_new_key_pair()
+            dummy_host = new_host(key_pair=key_pair)
+            dummy_host.set_stream_handler(
+                ID_PUSH, identify_push_handler_for(dummy_host)
+            )
+            listen_addr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
+            dummy_peers.append((dummy_host, listen_addr))
+
+        async with trio.open_nursery() as nursery:
+            # Start all dummy hosts
+            for host, listen_addr in dummy_peers:
+                nursery.start_soon(run_host_forever, host, listen_addr)
+
+            # Wait for all hosts to finish setting up listeners
+            for host, _ in dummy_peers:
+                await wait_until_listening(host)
+
+            # Now connect host_a → dummy peers
+            for host, _ in dummy_peers:
+                await host_a.connect(info_from_p2p_addr(host.get_addrs()[0]))
+
+            await push_identify_to_peers(
+                host_a,
+            )
+
+            await trio.sleep(0.5)
+
+            peer_id_a = host_a.get_id()
+            for host, _ in dummy_peers:
+                dummy_peerstore = host.get_peerstore()
+                assert peer_id_a in dummy_peerstore.peer_ids()
+
+            nursery.cancel_scope.cancel()

tests/utils/utils.py

@@ -2,13 +2,30 @@
     MagicMock,
 )
 
+import trio
 
-def create_mock_connections() -> dict:
+from libp2p.abc import IHost
+
+
+def create_mock_connections(count: int = 50) -> dict:
     connections = {}
 
-    for i in range(1, 31):
+    for i in range(1, count):
         peer_id = f"peer-{i}"
         mock_conn = MagicMock(name=f"INetConn-{i}")
         connections[peer_id] = mock_conn
 
     return connections
+
+
+async def run_host_forever(host: IHost, addr):
+    async with host.run([addr]):
+        await trio.sleep_forever()
+
+
+async def wait_until_listening(host, timeout=3):
+    with trio.move_on_after(timeout):
+        while not host.get_addrs():
+            await trio.sleep(0.05)
+        return
+    raise RuntimeError("Timed out waiting for host to get an address")