Repository: locustio/locust

Issue:
Worker processes should have an id number
Often you will want different workers performing using different data, doing things in different ways etc.

Right now that is kind of hard to implement (the best workaround I'm aware of involves passing different env vars to different worker processes)

But we could build this into locust. Every locust Runner should be given an ID number (runner.id), and MasterRunner should keep track of the current max id, so that new workers that connect during the run will be given an incremented id.

I think it is significantly easier (and maybe even better, at least more predictable) to allow any workers that disconnect during the test to leave "holes" in the sequence of IDs rather than reusing the ID for any new workers that connect.

Discussion:
gdm85: Any idea why the `KeyError` is happening in the CI check?
cyberw: I think the test is bad (fragile) because it assumes a msg with .data is not an ack
cyberw: Looks like a good start! Two questions:

Is the lock really needed? There are no io calls so I dont see how another greenlet could start to run?

Do we really need the old id:s or could this replace it? A number of places would be nicer if we used ordinal IDs instead of uuid:s (specifically logs)
cyberw: This needs some unit tests. 
gdm85: > I think the test is bad (fragile) because it assumes a msg with .data is not an ack

Ok, will try to fix it.



> Is the lock really needed? There are no io calls so I dont see how another greenlet could start to run?

I assumed so; but since `recv_from_client()` only happens serially then I will remove it.



> Do we still need the old id:s or could this replace it? A number of places would be nicer if we used ordinal IDs instead of uuid:s (e.g. the logs)

You mean the `self.client_id` (which is an UUID)? I guess it could be replaced but better in a separate PR? Also, some people downstream might rely on the hostname being there.

> This needs some unit tests.

Any suggestion? If we have one for ACKs then I would modify that one; and perhaps add another one for old masters (I already verified manually that new workers can connect to old masters, but a test would be best).
cyberw: > You mean the self.client_id (which is an UUID)? I guess it could be replaced but better in a separate PR? Also, some people downstream might rely on the hostname being there.

Yes. We could do it later on but I kind of would like to do it now (rather than having two potentially breaking changes)

>> This needs some unit tests.

> Any suggestion? If we have one for ACKs then I would modify that one; and perhaps add another one for old masters (I already verified manually that new workers can connect to old masters, but a test would be best).

no suggestions off hand (I’m on mobile atm). No need to validate backwards/forward compatibility, it looks reasonable. Just validate that the ID is available and incremented correctly. Check the existing tests for inspiration. 
gdm85: > Yes. We could do it later on but I kind of would like to do it now (rather than having two potentially breaking changes).

The changes introduced in this PR are not breaking though.
I can make more changes so that `client_id` is completely replaced and becomes the ordinal `int`; it will become a bigger PR though.

> no suggestions off hand (I’m on mobile atm). No need to validate backwards/forward compatibility, it looks reasonable. Just validate that the ID is available and incremented correctly. Check the existing tests for inspiration.

Ok, I have added one; thanks!
cyberw: >> Yes. We could do it later on but I kind of would like to do it now (rather than having two potentially breaking changes).
> The changes introduced in this PR are not breaking though.

Sorry, I was unclear: what I meant was that the second PR would be a breaking one, as it would remove the `ordinal_id` added here (as it would no longer be needed). 

I’d like to see it investigated at least. Doing it as a separate PR is fine if it becomes too complicated. 
gdm85: I am going to fix the tests and then add another commit which "promotes" the ordinal ID as the `client_uid`; note however that I think it is beneficial that clients have an UUID, so that when they connect back the same ordinal will be assigned to them (at least until master is restarted).

I will try to find the best approach to preserve this functionality.
gdm85: I have been thinking through it; I think removing the UUIDs is a bad idea because right now these are used for a sort of "weak" handshake: clients report their UUIDs and the messaging/events is based on that.

I will make a commit that renames `client_id` into `client_uuid` so that `client_id` can be used for this new ordinal (which is like an index).
cyberw: > I have been thinking through it; I think removing the UUIDs is a bad idea because right now these are used for a sort of "weak" handshake: clients report their UUIDs and the messaging/events is based on that.
> 
> I will make a commit that renames `client_id` into `client_uuid` (removing the socket hostname from it), and also introduce the new field as `worker_index`.

I see! Given that, I have changed my mind. I think it is preferable to keep client_id exactly as it is (including hostname), to promote backwards compatibility.

Can you also add an "integration test" (maybe in test_main.py) that checks that is correct on that side? (maybe by printing self.runner.worker_index from inside a task).

For convenience, how about adding worker_index to LocalRunner and just making it zero?
cyberw: MasterRunner.client_ordinals should probably be renamed worker_indexes instead (just to keep it consistent)
gdm85: > For convenience, how about adding worker_index to LocalRunner and just making it zero?

If there can never be more than 1 `LocalRunner`, then it is fine; I have changed it.

I have applied the other changes as well.

Patch:
diff --git a/docs/api.rst b/docs/api.rst
index 423b1e88..f4044dbc 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -133,7 +133,7 @@ Runner classes
     :members: register_message, send_message
 
 .. autoclass:: locust.runners.WorkerRunner
-    :members: register_message, send_message
+    :members: register_message, send_message, client_id, worker_index
 
 Web UI class
 ============
diff --git a/locust/runners.py b/locust/runners.py
index 9cfef4af..bdcebb44 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -443,6 +443,9 @@ class LocalRunner(Runner):
     Runner for running single process load test
     """
 
+    # always set to 0 for LocalRunner
+    worker_index = 0
+
     def __init__(self, environment) -> None:
         """
         :param environment: Environment instance
@@ -659,6 +662,8 @@ class MasterRunner(DistributedRunner):
         self.master_bind_port = master_bind_port
         self.spawn_rate: float = 0.0
         self.spawning_completed = False
+        self.worker_indexes: dict[str, int] = {}
+        self.worker_index_max = 0
 
         self.clients = WorkerNodes()
         try:
@@ -700,6 +705,19 @@ class MasterRunner(DistributedRunner):
             bool, self.environment.parsed_options.enable_rebalancing
         )
 
+    def get_worker_index(self, client_id):
+        """
+        Get the worker index for the specified client ID;
+        this is a deterministic 0-based ordinal number and guaranteed to not change
+        while Master is alive.
+        """
+        if client_id in self.worker_indexes:
+            return self.worker_indexes[client_id]
+        index = self.worker_index_max
+        self.worker_indexes[client_id] = index
+        self.worker_index_max += 1
+        return index
+
     @property
     def user_count(self) -> int:
         return sum([c.user_count for c in self.clients.values()])
@@ -999,7 +1017,7 @@ class MasterRunner(DistributedRunner):
                         logger.warning(
                             f"A worker ({client_id}) running a different version ({msg.data}) connected, master version is {__version__}"
                         )
-                self.send_message("ack", client_id=client_id)
+                self.send_message("ack", client_id=client_id, data={"index": self.get_worker_index(client_id)})
                 worker_node_id = msg.node_id
                 self.clients[worker_node_id] = WorkerNode(worker_node_id, heartbeat_liveness=HEARTBEAT_LIVENESS)
                 if self._users_dispatcher is not None:
@@ -1121,6 +1139,9 @@ class WorkerRunner(DistributedRunner):
     take the stats generated by the running users and send back to the :class:`MasterRunner`.
     """
 
+    # the worker index is set on ACK, if master provided it (masters <= 2.10.2 do not provide it)
+    worker_index = -1
+
     def __init__(self, environment: "Environment", master_host: str, master_port: int) -> None:
         """
         :param environment: Environment instance
@@ -1245,6 +1266,9 @@ class WorkerRunner(DistributedRunner):
                 logger.error(f"RPCError found when receiving from master: {e}")
                 continue
             if msg.type == "ack":
+                # backward-compatible support of masters that do not send a worker index
+                if msg.data is not None and "index" in msg.data:
+                    self.worker_index = msg.data["index"]
                 self.connection_event.set()
             elif msg.type == "spawn":
                 self.client.send(Message("spawning", None, self.client_id))


Test patch:
diff --git a/locust/test/test_main.py b/locust/test/test_main.py
index 0c47f838..469f1916 100644
--- a/locust/test/test_main.py
+++ b/locust/test/test_main.py
@@ -1325,3 +1325,91 @@ class SecondUser(HttpUser):
 
             self.assertEqual(0, proc.returncode)
             self.assertEqual(0, proc_worker.returncode)
+
+    def test_worker_indexes(self):
+        content = (
+            MOCK_LOCUSTFILE_CONTENT
+            + """
+class AnyUser(HttpUser):
+    host = "http://127.0.0.1:8089"
+    wait_time = between(0, 0.1)
+    @task
+    def my_task(self):
+        print("worker index:", self.environment.runner.worker_index)
+"""
+        )
+        with mock_locustfile(content=content) as mocked:
+            master = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--headless",
+                    "--master",
+                    "--expect-workers",
+                    "2",
+                    "-t",
+                    "5",
+                    "-u",
+                    "2",
+                    "-L",
+                    "DEBUG",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            proc_worker_1 = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--worker",
+                    "-L",
+                    "DEBUG",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            proc_worker_2 = subprocess.Popen(
+                [
+                    "locust",
+                    "-f",
+                    mocked.file_path,
+                    "--worker",
+                    "-L",
+                    "DEBUG",
+                ],
+                stdout=PIPE,
+                stderr=PIPE,
+                text=True,
+            )
+            stdout, stderr = master.communicate()
+            self.assertNotIn("Traceback", stderr)
+            self.assertEqual(0, master.returncode)
+
+            stdout_worker_1, stderr_worker_1 = proc_worker_1.communicate()
+            stdout_worker_2, stderr_worker_2 = proc_worker_2.communicate()
+            self.assertEqual(0, proc_worker_1.returncode)
+            self.assertEqual(0, proc_worker_2.returncode)
+            self.assertNotIn("Traceback", stderr_worker_1)
+            self.assertNotIn("Traceback", stderr_worker_2)
+
+            PREFIX = "worker index: "
+            p1 = stdout_worker_1.find(PREFIX)
+            if p1 == -1:
+                raise Exception(stdout_worker_1 + stderr_worker_1)
+            self.assertNotEqual(-1, p1)
+            p2 = stdout_worker_2.find(PREFIX)
+            if p2 == -1:
+                raise Exception(stdout_worker_2 + stderr_worker_2)
+            self.assertNotEqual(-1, p2)
+            found = [
+                int(stdout_worker_1[p1 + len(PREFIX) :].split("\n")[0]),
+                int(stdout_worker_2[p1 + len(PREFIX) :].split("\n")[0]),
+            ]
+            found.sort()
+            for i in range(2):
+                if found[i] != i:
+                    raise Exception(f"expected index {i} but got", found[i])
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 460cd756..aef75fdd 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -2635,7 +2635,9 @@ class TestMasterRunner(LocustRunnerTestCase):
             master.start(7, 7)
             self.assertEqual(10, len(server.outbox))
 
-            num_users = sum(sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.data)
+            num_users = sum(
+                sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.type != "ack"
+            )
 
             self.assertEqual(7, num_users, "Total number of locusts that would have been spawned is not 7")
 
@@ -2654,10 +2656,43 @@ class TestMasterRunner(LocustRunnerTestCase):
             master.start(2, 2)
             self.assertEqual(10, len(server.outbox))
 
-            num_users = sum(sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.data)
+            num_users = sum(
+                sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.type != "ack"
+            )
 
             self.assertEqual(2, num_users, "Total number of locusts that would have been spawned is not 2")
 
+    def test_spawn_correct_worker_indexes(self):
+        """
+        Tests that workers would receive a monotonic sequence of ordinal IDs.
+        """
+
+        class TestUser(User):
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            master = self.get_runner(user_classes=[TestUser])
+
+            USERS_COUNT = 5
+
+            for i in range(USERS_COUNT):
+                server.mocked_send(Message("client_ready", __version__, "fake_client%i" % i))
+
+            master.start(USERS_COUNT, USERS_COUNT)
+            self.assertEqual(USERS_COUNT * 2, len(server.outbox))
+
+            indexes = []
+            for _, msg in server.outbox:
+                if msg.type == "ack":
+                    indexes.append(msg.data["index"])
+            self.assertEqual(USERS_COUNT, len(indexes), "Total number of locusts/workers is not 5")
+
+            indexes.sort()
+            for i in range(USERS_COUNT):
+                self.assertEqual(indexes[i], i, "Worker index mismatch")
+
     def test_custom_shape_scale_up(self):
         class MyUser(User):
             @task
@@ -2687,14 +2722,18 @@ class TestMasterRunner(LocustRunnerTestCase):
             sleep(0.5)
 
             # Wait for shape_worker to update user_count
-            num_users = sum(sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.data)
+            num_users = sum(
+                sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.type != "ack"
+            )
             self.assertEqual(
                 1, num_users, "Total number of users in first stage of shape test is not 1: %i" % num_users
             )
 
             # Wait for shape_worker to update user_count again
             sleep(2)
-            num_users = sum(sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.data)
+            num_users = sum(
+                sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.type != "ack"
+            )
             self.assertEqual(
                 3, num_users, "Total number of users in second stage of shape test is not 3: %i" % num_users
             )
@@ -2732,7 +2771,9 @@ class TestMasterRunner(LocustRunnerTestCase):
             sleep(0.5)
 
             # Wait for shape_worker to update user_count
-            num_users = sum(sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.data)
+            num_users = sum(
+                sum(msg.data["user_classes_count"].values()) for _, msg in server.outbox if msg.type != "ack"
+            )
             self.assertEqual(
                 5, num_users, "Total number of users in first stage of shape test is not 5: %i" % num_users
             )
@@ -2741,7 +2782,7 @@ class TestMasterRunner(LocustRunnerTestCase):
             sleep(2)
             msgs = defaultdict(dict)
             for _, msg in server.outbox:
-                if not msg.data:
+                if msg.type == "ack":
                     continue
                 msgs[msg.node_id][msg.data["timestamp"]] = sum(msg.data["user_classes_count"].values())
             # Count users for the last received messages
@@ -2977,6 +3018,7 @@ class TestMasterRunner(LocustRunnerTestCase):
             self.assertEqual(1, len(master.clients))
             self.assertEqual("ack", server.outbox[0][1].type)
             self.assertEqual(1, len(server.outbox))
+            self.assertEqual(0, server.outbox[0][1].data["index"])
 
     def test_worker_sends_bad_message_to_master(self):
         """
@@ -3018,7 +3060,7 @@ class TestWorkerRunner(LocustTestCase):
 
     def get_runner(self, client, environment=None, user_classes=None, auto_connect=True):
         if auto_connect:
-            client.mocked_send(Message("ack", {}, "dummy_client_id"))
+            client.mocked_send(Message("ack", {"index": 0}, "dummy_client_id"))
         if environment is None:
             environment = self.environment
         user_classes = user_classes or []