TODO: throttle on async validators (#755)

* fixed todo: throttle on async validators

* added test: validate message respects concurrency limit

* added newsfragment

* added configurable validator semaphore in the PubSub constructor

* added the concurrency-checker in the original test-validate-msg test case

* separate out a _run_async_validator function

* remove redundant run_async_validator

Author: lla-dane

libp2p/pubsub/pubsub.py

@@ -102,13 +102,17 @@ class TopicValidator(NamedTuple):
     is_async: bool
 
 
+MAX_CONCURRENT_VALIDATORS = 10
+
+
 class Pubsub(Service, IPubsub):
     host: IHost
 
     router: IPubsubRouter
 
     peer_receive_channel: trio.MemoryReceiveChannel[ID]
     dead_peer_receive_channel: trio.MemoryReceiveChannel[ID]
+    _validator_semaphore: trio.Semaphore
 
     seen_messages: LastSeenCache
 
@@ -143,6 +147,7 @@ def __init__(
         msg_id_constructor: Callable[
             [rpc_pb2.Message], bytes
         ] = get_peer_and_seqno_msg_id,
+        max_concurrent_validator_count: int = MAX_CONCURRENT_VALIDATORS,
     ) -> None:
         """
         Construct a new Pubsub object, which is responsible for handling all
@@ -168,6 +173,7 @@ def __init__(
         # Therefore, we can only close from the receive side.
         self.peer_receive_channel = peer_receive
         self.dead_peer_receive_channel = dead_peer_receive
+        self._validator_semaphore = trio.Semaphore(max_concurrent_validator_count)
         # Register a notifee
         self.host.get_network().register_notifee(
             PubsubNotifee(peer_send, dead_peer_send)
@@ -657,7 +663,11 @@ async def publish(self, topic_id: str | list[str], data: bytes) -> None:
 
         logger.debug("successfully published message %s", msg)
 
-    async def validate_msg(self, msg_forwarder: ID, msg: rpc_pb2.Message) -> None:
+    async def validate_msg(
+        self,
+        msg_forwarder: ID,
+        msg: rpc_pb2.Message,
+    ) -> None:
         """
         Validate the received message.
 
@@ -680,23 +690,34 @@ async def validate_msg(self, msg_forwarder: ID, msg: rpc_pb2.Message) -> None:
             if not validator(msg_forwarder, msg):
                 raise ValidationError(f"Validation failed for msg={msg}")
 
-        # TODO: Implement throttle on async validators
-
         if len(async_topic_validators) > 0:
             # Appends to lists are thread safe in CPython
-            results = []
-
-            async def run_async_validator(func: AsyncValidatorFn) -> None:
-                result = await func(msg_forwarder, msg)
-                results.append(result)
+            results: list[bool] = []
 
             async with trio.open_nursery() as nursery:
                 for async_validator in async_topic_validators:
-                    nursery.start_soon(run_async_validator, async_validator)
+                    nursery.start_soon(
+                        self._run_async_validator,
+                        async_validator,
+                        msg_forwarder,
+                        msg,
+                        results,
+                    )
 
             if not all(results):
                 raise ValidationError(f"Validation failed for msg={msg}")
 
+    async def _run_async_validator(
+        self,
+        func: AsyncValidatorFn,
+        msg_forwarder: ID,
+        msg: rpc_pb2.Message,
+        results: list[bool],
+    ) -> None:
+        async with self._validator_semaphore:
+            result = await func(msg_forwarder, msg)
+            results.append(result)
+
     async def push_msg(self, msg_forwarder: ID, msg: rpc_pb2.Message) -> None:
         """
         Push a pubsub message to others.

newsfragments/755.performance.rst

@@ -0,0 +1,2 @@
+Added throttling for async topic validators in validate_msg, enforcing a
+concurrency limit to prevent resource exhaustion under heavy load.

tests/core/pubsub/test_pubsub.py

@@ -5,10 +5,12 @@
 from typing import (
     NamedTuple,
 )
+from unittest.mock import patch
 
 import pytest
 import trio
 
+from libp2p.custom_types import AsyncValidatorFn
 from libp2p.exceptions import (
     ValidationError,
 )
@@ -243,7 +245,37 @@ async def async_validator(peer_id: ID, msg: rpc_pb2.Message) -> bool:
     ((False, True), (True, False), (True, True)),
 )
 @pytest.mark.trio
-async def test_validate_msg(is_topic_1_val_passed, is_topic_2_val_passed):
+async def test_validate_msg_with_throttle_condition(
+    is_topic_1_val_passed, is_topic_2_val_passed
+):
+    CONCURRENCY_LIMIT = 10
+
+    state = {
+        "concurrency_counter": 0,
+        "max_observed": 0,
+    }
+    lock = trio.Lock()
+
+    async def mock_run_async_validator(
+        self,
+        func: AsyncValidatorFn,
+        msg_forwarder: ID,
+        msg: rpc_pb2.Message,
+        results: list[bool],
+    ) -> None:
+        async with self._validator_semaphore:
+            async with lock:
+                state["concurrency_counter"] += 1
+                if state["concurrency_counter"] > state["max_observed"]:
+                    state["max_observed"] = state["concurrency_counter"]
+
+            try:
+                result = await func(msg_forwarder, msg)
+                results.append(result)
+            finally:
+                async with lock:
+                    state["concurrency_counter"] -= 1
+
     async with PubsubFactory.create_batch_with_floodsub(1) as pubsubs_fsub:
 
         def passed_sync_validator(peer_id: ID, msg: rpc_pb2.Message) -> bool:
@@ -280,11 +312,19 @@ async def failed_async_validator(peer_id: ID, msg: rpc_pb2.Message) -> bool:
             seqno=b"\x00" * 8,
         )
 
-        if is_topic_1_val_passed and is_topic_2_val_passed:
-            await pubsubs_fsub[0].validate_msg(pubsubs_fsub[0].my_id, msg)
-        else:
-            with pytest.raises(ValidationError):
+        with patch(
+            "libp2p.pubsub.pubsub.Pubsub._run_async_validator",
+            new=mock_run_async_validator,
+        ):
+            if is_topic_1_val_passed and is_topic_2_val_passed:
                 await pubsubs_fsub[0].validate_msg(pubsubs_fsub[0].my_id, msg)
+            else:
+                with pytest.raises(ValidationError):
+                    await pubsubs_fsub[0].validate_msg(pubsubs_fsub[0].my_id, msg)
+
+    assert state["max_observed"] <= CONCURRENCY_LIMIT, (
+        f"Max concurrency observed: {state['max_observed']}"
+    )
 
 
 @pytest.mark.trio