feat: implement flood_publish option in GossipSub and add corresponding tests

Author: yashksaini-coder

libp2p/pubsub/gossipsub.py

@@ -140,8 +140,14 @@ def __init__(
         self.mcache = MessageCache(gossip_window, gossip_history)
 
         # Whether to flood publish to all peers instead of following gossipsub
-        # mesh/fanout logic. Kept as an option primarily for tests and
-        # compatibility with the test factory. Default is False.
+        # mesh/fanout logic when acting as the original publisher.
+        # When enabled, this behaves as a hybrid between FloodSub and GossipSub:
+        # - When this node is the original publisher: Message is sent to ALL peers
+        #   who are subscribed to the topic (flood publishing behavior)
+        # - When this node is forwarding a message: Regular GossipSub behavior is used
+        # This provides better reliability at publication time with a reasonable
+        # bandwidth cost since it only affects the original publisher.
+        # Default is False.
         self.flood_publish = flood_publish
 
         # Create heartbeat timer
@@ -306,43 +312,52 @@ def _get_peers_to_send(
             if topic not in self.pubsub.peer_topics:
                 continue
 
-            # direct peers
-            _direct_peers: set[ID] = {_peer for _peer in self.direct_peers}
-            send_to.update(_direct_peers)
-
-            # floodsub peers
-            floodsub_peers: set[ID] = {
-                peer_id
-                for peer_id in self.pubsub.peer_topics[topic]
-                if peer_id in self.peer_protocol
-                and self.peer_protocol[peer_id] == floodsub.PROTOCOL_ID
-            }
-            send_to.update(floodsub_peers)
-
-            # gossipsub peers
-            gossipsub_peers: set[ID] = set()
-            if topic in self.mesh:
-                gossipsub_peers = self.mesh[topic]
+            # If flood_publish is enabled and we are the original publisher,
+            # send to all peers in the topic (flood publishing behavior)
+            if self.flood_publish and msg_forwarder == self.pubsub.my_id:
+                for peer in self.pubsub.peer_topics[topic]:
+                    # TODO: add score threshold check when peer scoring is implemented
+                    # if direct peer then skip score check
+                    send_to.add(peer)
             else:
-                # When we publish to a topic that we have not subscribe to, we randomly
-                # pick `self.degree` number of peers who have subscribed to the topic
-                # and add them as our `fanout` peers.
-                topic_in_fanout: bool = topic in self.fanout
-                fanout_peers: set[ID] = self.fanout[topic] if topic_in_fanout else set()
-                fanout_size = len(fanout_peers)
-                if not topic_in_fanout or (
-                    topic_in_fanout and fanout_size < self.degree
-                ):
-                    if topic in self.pubsub.peer_topics:
-                        # Combine fanout peers with selected peers
-                        fanout_peers.update(
-                            self._get_in_topic_gossipsub_peers_from_minus(
-                                topic, self.degree - fanout_size, fanout_peers
+                # Regular GossipSub routing logic
+                # direct peers
+                _direct_peers: set[ID] = {_peer for _peer in self.direct_peers}
+                send_to.update(_direct_peers)
+
+                # floodsub peers
+                floodsub_peers: set[ID] = {
+                    peer_id
+                    for peer_id in self.pubsub.peer_topics[topic]
+                    if peer_id in self.peer_protocol
+                    and self.peer_protocol[peer_id] == floodsub.PROTOCOL_ID
+                }
+                send_to.update(floodsub_peers)
+
+                # gossipsub peers
+                gossipsub_peers: set[ID] = set()
+                if topic in self.mesh:
+                    gossipsub_peers = self.mesh[topic]
+                else:
+                    # When we publish to a topic that we have not subscribe to, we randomly
+                    # pick `self.degree` number of peers who have subscribed to the topic
+                    # and add them as our `fanout` peers.
+                    topic_in_fanout: bool = topic in self.fanout
+                    fanout_peers: set[ID] = self.fanout[topic] if topic_in_fanout else set()
+                    fanout_size = len(fanout_peers)
+                    if not topic_in_fanout or (
+                        topic_in_fanout and fanout_size < self.degree
+                    ):
+                        if topic in self.pubsub.peer_topics:
+                            # Combine fanout peers with selected peers
+                            fanout_peers.update(
+                                self._get_in_topic_gossipsub_peers_from_minus(
+                                    topic, self.degree - fanout_size, fanout_peers
+                                )
                             )
-                        )
-                self.fanout[topic] = fanout_peers
-                gossipsub_peers = fanout_peers
-            send_to.update(gossipsub_peers)
+                    self.fanout[topic] = fanout_peers
+                    gossipsub_peers = fanout_peers
+                send_to.update(gossipsub_peers)
         # Excludes `msg_forwarder` and `origin`
         yield from send_to.difference([msg_forwarder, origin])
 

tests/core/pubsub/test_gossipsub.py

@@ -602,6 +602,8 @@ async def test_sparse_connect():
 
 @pytest.mark.trio
 async def test_flood_publish():
+    """Test that with flood_publish disabled, message propagation still works
+    in a fully connected network topology."""
     async with PubsubFactory.create_batch_with_gossipsub(
         6,
         degree=2,
@@ -629,18 +631,78 @@ async def test_flood_publish():
         await pubsubs_gsub[0].publish(topic, msg_content)
 
         # wait for messages to propagate
-        await trio.sleep(0.5)
+        await trio.sleep(1)
 
-        print(routers[0].mesh[topic])
-        if routers[0].pubsub:
-            print(routers[0].pubsub.peer_topics)
+        # Debug info - only log if needed
+        # print(f"Mesh for topic: {routers[0].mesh[topic]}")
+        # if routers[0].pubsub:
+        #     print(f"Peer topics: {routers[0].pubsub.peer_topics}")
 
-        # verify all nodes received the message
-        for queue in queues:
-            msg = await queue.get()
-            assert msg.data == msg_content, (
-                f"node did not receive expected message: {msg.data}"
-            )
+        # verify all nodes received the message with timeout
+        for i, queue in enumerate(queues):
+            try:
+                with trio.fail_after(5):
+                    msg = await queue.get()
+                    assert msg.data == msg_content, f"node {i} received wrong message: {msg.data}"
+            except trio.TooSlowError:
+                pytest.fail(f"Node {i} did not receive the message (timeout)")
+                
+        # Test passed if all nodes received the message
+        print("Basic flood test passed - all nodes received the message")
+
+
+@pytest.mark.trio
+async def test_flood_publish_enabled():
+    """Test that with flood_publish enabled, all nodes receive the message
+    even with a sparse network topology."""
+    # Create a network with flood_publish enabled
+    async with PubsubFactory.create_batch_with_gossipsub(
+        6,
+        degree=2,
+        degree_low=1,
+        degree_high=3,
+        flood_publish=True,  # Enable flood_publish
+    ) as pubsubs_gsub:
+        routers: list[GossipSub] = []
+        for pubsub in pubsubs_gsub:
+            assert isinstance(pubsub.router, GossipSub)
+            routers.append(pubsub.router)
+        hosts = [ps.host for ps in pubsubs_gsub]
+
+        topic = "flood_test_topic"
+        queues = [await pubsub.subscribe(topic) for pubsub in pubsubs_gsub]
+
+        # Create a sparse topology - only connect to a few nodes
+        # We only connect nodes in a chain, which would normally
+        # prevent complete message propagation without flood_publish
+        await connect(hosts[0], hosts[1])
+        await connect(hosts[1], hosts[2])
+        await connect(hosts[2], hosts[3])
+        await connect(hosts[3], hosts[4])
+        await connect(hosts[4], hosts[5])
+
+        # wait for connections to be established
+        await trio.sleep(1)
+
+        # publish a message from the first host
+        msg_content = b"flood_publish_msg"
+        await pubsubs_gsub[0].publish(topic, msg_content)
+
+        # wait for messages to propagate
+        await trio.sleep(2)
+        
+        # verify all nodes received the message with timeout
+        for i, queue in enumerate(queues):
+            try:
+                with trio.fail_after(5):
+                    msg = await queue.get()
+                    assert msg.data == msg_content, f"node {i} received wrong message: {msg.data}"
+                    print(f"Node {i} received message correctly")
+            except trio.TooSlowError:
+                pytest.fail(f"Node {i} did not receive the message (timeout)")
+        
+        # Test passed if all nodes received the message
+        print("Flood publish test passed - all nodes received the message")
 
 
 @pytest.mark.trio