Add test for worker ordering

Author: AlanCoding

dispatcherd/protocols.py

@@ -210,6 +210,8 @@ def __iter__(self) -> Iterator[PoolWorker]: ...
 
     def get_by_id(self, worker_id: int) -> PoolWorker: ...
 
+    def move_to_end(self, worker_id: int) -> None: ...
+
 
 class SharedAsyncObjects:
     exit_event: asyncio.Event

dispatcherd/service/pool.py

@@ -5,6 +5,7 @@
 import signal
 import time
 from typing import Any, Iterator, Literal, Optional
+from collections import OrderedDict
 
 from ..processors.blocker import Blocker
 from ..processors.queuer import Queuer
@@ -173,7 +174,7 @@ def __init__(self) -> None:
 
 class WorkerData(WorkerDataProtocol):
     def __init__(self) -> None:
-        self.workers: dict[int, PoolWorker] = {}
+        self.workers: OrderedDict[int, PoolWorker] = OrderedDict()
         self.management_lock = asyncio.Lock()
 
     def __iter__(self) -> Iterator[PoolWorker]:
@@ -194,6 +195,9 @@ def get_by_id(self, worker_id: int) -> PoolWorker:
     def remove_by_id(self, worker_id: int) -> None:
         del self.workers[worker_id]
 
+    def move_to_end(self, worker_id: int) -> None:
+        self.workers.move_to_end(worker_id)
+
 
 class WorkerPool(WorkerPoolProtocol):
     def __init__(
@@ -559,6 +563,7 @@ async def process_finished(self, worker: PoolWorker, message: dict) -> None:
             else:
                 self.finished_count += 1
             worker.mark_finished_task()
+            self.workers.move_to_end(worker.worker_id)
 
         if not self.queuer.queued_messages and all(worker.current_task is None for worker in self.workers):
             self.events.work_cleared.set()

tests/unit/service/test_worker_order.py

@@ -0,0 +1,69 @@
+import asyncio
+from typing import AsyncIterator
+
+import pytest
+import pytest_asyncio
+
+from dispatcherd.protocols import DispatcherMain
+from dispatcherd.testing.asyncio import adispatcher_service
+
+
+@pytest.fixture(scope='session')
+def order_config():
+    return {
+        "version": 2,
+        "service": {
+            "pool_kwargs": {"min_workers": 2, "max_workers": 2},
+            "main_kwargs": {"node_id": "order-test"},
+        },
+    }
+
+
+@pytest_asyncio.fixture
+async def aorder_dispatcher(order_config) -> AsyncIterator[DispatcherMain]:
+    async with adispatcher_service(order_config) as dispatcher:
+        yield dispatcher
+
+
+@pytest.mark.asyncio
+async def test_workers_reorder_and_dispatch_longest_idle(aorder_dispatcher):
+    pool = aorder_dispatcher.pool
+    assert list(pool.workers.workers.keys()) == [0, 1]
+
+    pool.events.work_cleared.clear()
+    await aorder_dispatcher.process_message({
+        "task": "tests.data.methods.sleep_function",
+        "kwargs": {"seconds": 0.1},
+        "uuid": "t1",
+    })
+    await aorder_dispatcher.process_message({
+        "task": "tests.data.methods.sleep_function",
+        "kwargs": {"seconds": 0.05},
+        "uuid": "t2",
+    })
+    await asyncio.wait_for(pool.events.work_cleared.wait(), timeout=1)
+
+    assert list(pool.workers.workers.keys()) == [1, 0]
+
+    pool.events.work_cleared.clear()
+    await aorder_dispatcher.process_message({
+        "task": "tests.data.methods.sleep_function",
+        "kwargs": {"seconds": 0.01},
+        "uuid": "t3",
+    })
+    await asyncio.sleep(0.01)
+    assert pool.workers.get_by_id(1).current_task["uuid"] == "t3"
+    assert pool.workers.get_by_id(0).current_task is None
+    await asyncio.wait_for(pool.events.work_cleared.wait(), timeout=1)
+
+    pool.events.work_cleared.clear()
+    await aorder_dispatcher.process_message({
+        "task": "tests.data.methods.sleep_function",
+        "kwargs": {"seconds": 0.01},
+        "uuid": "t4",
+    })
+    await asyncio.sleep(0.01)
+    assert pool.workers.get_by_id(0).current_task["uuid"] == "t4"
+    await asyncio.wait_for(pool.events.work_cleared.wait(), timeout=1)
+
+    assert list(pool.workers.workers.keys()) == [1, 0]