Repository: locustio/locust

Issue:
Allow master node to supply data to worker nodes directly
<!-- 
If you have a general question about how to use Locust, please check Stack Overflow first https://stackoverflow.com/questions/tagged/locust

You can also ask new questions on SO, https://stackoverflow.com/questions/ask just remember to tag your question with "locust".
-->

### Is your feature request related to a problem? Please describe.
<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] -->
Pretty much a duplicate of #1506, but it would be extremely beneficial for locust to be able to communicate data between nodes easily. An example that you're probably familiar with is needing to retrieve user accounts from a database for login, and trying to coordinate that between several worker nodes so that the same users are not logged in across multiple workers. 

### Describe the solution you'd like
<!-- A clear and concise description of what you want to happen -->
Adding hooks to allow communication between nodes would be ideal. I'm not sure exactly what the best solution would be but allowing the user to define custom messages and handlers for the greenlets seems like it'd be easy enough to implement and massively improve locust's ability to coordinate itself in distributed mode. 

### Describe alternatives you've considered
<!-- A clear and concise description of any alternative solutions or features you've considered -->
The currently accepted workarounds (e.g. mongodb or a REST api) are much less than ideal, especially for something like a docker environment where each additional component must be installed and set up regularly.

### Additional context
<!-- Add any other context or screenshots about the feature request here -->
I'm available to give this a shot, but I wanted to make sure A) it has value and B) the best implementation from those who are more experienced on this project. 


Discussion:
Absolutely, go for it! Try to build upon the existing zeromq communication. 

Resolving patch:
diff --git a/examples/custom_messages.py b/examples/custom_messages.py
new file mode 100644
--- /dev/null
+++ b/examples/custom_messages.py
@@ -0,0 +1,60 @@
+from locust import HttpUser, task, events, between
+from locust.runners import MasterRunner, WorkerRunner
+
+usernames = []
+
+
+def setup_test_users(environment, msg, **kwargs):
+    # Fired when the worker recieves a message of type 'test_users'
+    usernames.extend(map(lambda u: u["name"], msg.data))
+    environment.runner.send_message("acknowledge_users", f"Thanks for the {len(msg.data)} users!")
+
+
+def on_acknowledge(msg, **kwargs):
+    # Fired when the master recieves a message of type 'acknowledge_users'
+    print(msg.data)
+
+
+@events.init.add_listener
+def on_locust_init(environment, **_kwargs):
+    if not isinstance(environment.runner, MasterRunner):
+        environment.runner.register_message("test_users", setup_test_users)
+    if not isinstance(environment.runner, WorkerRunner):
+        environment.runner.register_message("acknowledge_users", on_acknowledge)
+
+
+@events.test_start.add_listener
+def on_test_start(environment, **_kwargs):
+    # When the test is started, evenly divides list between
+    # worker nodes to ensure unique data across threads
+    if not isinstance(environment.runner, WorkerRunner):
+        users = []
+        for i in range(environment.runner.target_user_count):
+            users.append({"name": f"User{i}"})
+
+        worker_count = environment.runner.worker_count
+        chunk_size = int(len(users) / worker_count)
+
+        for i, worker in enumerate(environment.runner.clients):
+            start_index = i * chunk_size
+
+            if i + 1 < worker_count:
+                end_index = start_index + chunk_size
+            else:
+                end_index = len(users)
+
+            data = users[start_index:end_index]
+            environment.runner.send_message("test_users", data, worker)
+
+
+class WebsiteUser(HttpUser):
+    host = "http://127.0.0.1:8089"
+    wait_time = between(2, 5)
+
+    def __init__(self, parent):
+        self.username = usernames.pop()
+        super(WebsiteUser, self).__init__(parent)
+
+    @task
+    def task(self):
+        print(self.username)
diff --git a/locust/runners.py b/locust/runners.py
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -69,6 +69,7 @@ def __init__(self, environment):
         self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
         self.exceptions = {}
         self.target_user_count = None
+        self.custom_messages = {}
 
         # set up event listeners for recording requests
         def on_request_success(request_type, name, response_time, response_length, **_kwargs):
@@ -381,6 +382,15 @@ def log_exception(self, node_id, msg, formatted_tb):
         row["nodes"].add(node_id)
         self.exceptions[key] = row
 
+    def register_message(self, msg_type, listener):
+        """
+        Register a listener for a custom message from another node
+
+        :param msg_type: The type of the message to listen for
+        :param listener: The function to execute when the message is received
+        """
+        self.custom_messages[msg_type] = listener
+
 
 class LocalRunner(Runner):
     """
@@ -425,6 +435,21 @@ def stop(self):
         super().stop()
         self.environment.events.test_stop.fire(environment=self.environment)
 
+    def send_message(self, msg_type, data=None):
+        """
+        Emulates internodal messaging by calling registered listeners
+
+        :param msg_type: The type of the message to emulate sending
+        :param data: Optional data to include
+        """
+        logger.debug(f"Running locally: sending {msg_type} message to self")
+        if msg_type in self.custom_messages:
+            listener = self.custom_messages[msg_type]
+            msg = Message(msg_type, data, "local")
+            listener(environment=self.environment, msg=msg)
+        else:
+            logger.warning(f"Unknown message type recieved: {msg_type}")
+
 
 class DistributedRunner(Runner):
     def __init__(self, *args, **kwargs):
@@ -709,6 +734,11 @@ def client_listener(self):
                             self.quit()
             elif msg.type == "exception":
                 self.log_exception(msg.node_id, msg.data["msg"], msg.data["traceback"])
+            elif msg.type in self.custom_messages:
+                logger.debug(f"Recieved {msg.type} message from worker {msg.node_id}")
+                self.custom_messages[msg.type](environment=self.environment, msg=msg)
+            else:
+                logger.warning(f"Unknown message type recieved from worker {msg.node_id}: {msg.type}")
 
             self.check_stopped()
 
@@ -716,6 +746,23 @@ def client_listener(self):
     def worker_count(self):
         return len(self.clients.ready) + len(self.clients.spawning) + len(self.clients.running)
 
+    def send_message(self, msg_type, data=None, client_id=None):
+        """
+        Sends a message to attached worker node(s)
+
+        :param msg_type: The type of the message to send
+        :param data: Optional data to send
+        :param client_id: Optional id of the target worker node.
+                            If None, will send to all attached workers
+        """
+        if client_id:
+            logger.debug(f"Sending {msg_type} message to client {client_id}")
+            self.server.send_to_client(Message(msg_type, data, client_id))
+        else:
+            for client in self.clients.all:
+                logger.debug(f"Sending {msg_type} message to client {client.id}")
+                self.server.send_to_client(Message(msg_type, data, client.id))
+
 
 class WorkerRunner(DistributedRunner):
     """
@@ -828,6 +875,11 @@ def worker(self):
                 self.stop()
                 self._send_stats()  # send a final report, in case there were any samples not yet reported
                 self.greenlet.kill(block=True)
+            elif msg.type in self.custom_messages:
+                logger.debug(f"Recieved {msg.type} message from master")
+                self.custom_messages[msg.type](environment=self.environment, msg=msg)
+            else:
+                logger.warning(f"Unknown message type recieved: {msg.type}")
 
     def stats_reporter(self):
         while True:
@@ -837,6 +889,16 @@ def stats_reporter(self):
                 logger.error("Temporary connection lost to master server: %s, will retry later." % (e))
             gevent.sleep(WORKER_REPORT_INTERVAL)
 
+    def send_message(self, msg_type, data=None):
+        """
+        Sends a message to master node
+
+        :param msg_type: The type of the message to send
+        :param data: Optional data to send
+        """
+        logger.debug(f"Sending {msg_type} message to master")
+        self.client.send(Message(msg_type, data, self.client_id))
+
     def _send_stats(self):
         data = {}
         self.environment.events.report_to_master.fire(client_id=self.client_id, data=data)


Test patch:
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -438,6 +438,54 @@ def my_task(self):
         user_count = len(runner.user_greenlets)
         self.assertTrue(user_count == 2, "User count has not decreased correctly to 2, it is : %i" % user_count)
 
+    def test_custom_message(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        test_custom_msg = [False]
+        test_custom_msg_data = [{}]
+
+        def on_custom_msg(msg, **kw):
+            test_custom_msg[0] = True
+            test_custom_msg_data[0] = msg.data
+
+        environment = Environment(user_classes=[MyUser])
+        runner = LocalRunner(environment)
+
+        runner.register_message("test_custom_msg", on_custom_msg)
+        runner.send_message("test_custom_msg", {"test_data": 123})
+
+        self.assertTrue(test_custom_msg[0])
+        self.assertEqual(123, test_custom_msg_data[0]["test_data"])
+
+    def test_undefined_custom_message(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        test_custom_msg = [False]
+
+        def on_custom_msg(msg, **kw):
+            test_custom_msg[0] = True
+
+        environment = Environment(user_classes=[MyUser])
+        runner = LocalRunner(environment)
+
+        runner.register_message("test_custom_msg", on_custom_msg)
+        runner.send_message("test_different_custom_msg")
+
+        self.assertFalse(test_custom_msg[0])
+        self.assertEqual(1, len(self.mocked_log.warning))
+        msg = self.mocked_log.warning[0]
+        self.assertIn("Unknown message type recieved", msg)
+
 
 class TestMasterWorkerRunners(LocustTestCase):
     def test_distributed_integration_run(self):
@@ -1321,6 +1369,73 @@ def test_master_reset_connection(self):
                 self.assertEqual(1, len(master.clients))
                 master.quit()
 
+    def test_custom_message_send(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            master = self.get_runner()
+            for i in range(5):
+                master.clients[i] = WorkerNode(i)
+            master.send_message("test_custom_msg", {"test_data": 123})
+
+            self.assertEqual(5, len(server.outbox))
+            for _, msg in server.outbox:
+                self.assertEqual("test_custom_msg", msg.type)
+                self.assertEqual(123, msg.data["test_data"])
+
+    def test_custom_message_receive(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            test_custom_msg = [False]
+            test_custom_msg_data = [{}]
+
+            def on_custom_msg(msg, **kw):
+                test_custom_msg[0] = True
+                test_custom_msg_data[0] = msg.data
+
+            master = self.get_runner()
+            master.register_message("test_custom_msg", on_custom_msg)
+
+            server.mocked_send(Message("test_custom_msg", {"test_data": 123}, "dummy_id"))
+
+            self.assertTrue(test_custom_msg[0])
+            self.assertEqual(123, test_custom_msg_data[0]["test_data"])
+
+    def test_undefined_custom_message_receive(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            test_custom_msg = [False]
+
+            def on_custom_msg(msg, **kw):
+                test_custom_msg[0] = True
+
+            master = self.get_runner()
+            master.register_message("test_custom_msg", on_custom_msg)
+
+            server.mocked_send(Message("unregistered_custom_msg", {}, "dummy_id"))
+
+            self.assertFalse(test_custom_msg[0])
+            self.assertEqual(1, len(self.mocked_log.warning))
+            msg = self.mocked_log.warning[0]
+            self.assertIn("Unknown message type recieved from worker", msg)
+
 
 class TestWorkerRunner(LocustTestCase):
     def setUp(self):
@@ -1470,6 +1585,75 @@ def my_task(self):
             self.assertEqual(9, len(worker.user_greenlets))
             worker.quit()
 
+    def test_custom_message_send(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
+            environment = Environment()
+            worker = self.get_runner(environment=environment, user_classes=[MyUser])
+            client.outbox.clear()
+            worker.send_message("test_custom_msg", {"test_data": 123})
+            self.assertEqual("test_custom_msg", client.outbox[0].type)
+            self.assertEqual(123, client.outbox[0].data["test_data"])
+            worker.quit()
+
+    def test_custom_message_receive(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
+            environment = Environment()
+            test_custom_msg = [False]
+            test_custom_msg_data = [{}]
+
+            def on_custom_msg(msg, **kw):
+                test_custom_msg[0] = True
+                test_custom_msg_data[0] = msg.data
+
+            worker = self.get_runner(environment=environment, user_classes=[MyUser])
+            worker.register_message("test_custom_msg", on_custom_msg)
+
+            client.mocked_send(Message("test_custom_msg", {"test_data": 123}, "dummy_client_id"))
+
+            self.assertTrue(test_custom_msg[0])
+            self.assertEqual(123, test_custom_msg_data[0]["test_data"])
+            worker.quit()
+
+    def test_undefined_custom_message_receive(self):
+        class MyUser(User):
+            wait_time = constant(1)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
+            environment = Environment()
+
+            test_custom_msg = [False]
+
+            def on_custom_msg(msg, **kw):
+                test_custom_msg[0] = True
+
+            worker = self.get_runner(environment=environment, user_classes=[MyUser])
+            worker.register_message("test_custom_msg", on_custom_msg)
+
+            client.mocked_send(Message("unregistered_custom_msg", {}, "dummy_id"))
+
+            self.assertFalse(test_custom_msg[0])
+            self.assertEqual(1, len(self.mocked_log.warning))
+            msg = self.mocked_log.warning[0]
+            self.assertIn("Unknown message type recieved", msg)
+
 
 class TestMessageSerializing(unittest.TestCase):
     def test_message_serialize(self):