Ensure stream messages are always ordered

distributed/core.py

@@ -1006,10 +1006,7 @@
                             break
                         handler = self.stream_handlers[op]
                         if iscoroutinefunction(handler):
-                            self._ongoing_background_tasks.call_soon(
-                                handler, **merge(extra, msg)
-                            )
-                            await asyncio.sleep(0)
+                            await handler(**merge(extra, msg))
                         else:
                             handler(**merge(extra, msg))
                     else:
@@ -1521,6 +1518,14 @@
 
         return _().__await__()
 
+    async def __aenter__(self):
+        await self
+        return self
+
+    async def __aexit__(self, *args):
+        await self.close()
+        return
+
     async def start(self) -> None:
         # Invariant: semaphore._value == limit - open - _n_connecting
         self.semaphore = asyncio.Semaphore(self.limit)

distributed/tests/test_core.py

@@ -15,6 +15,7 @@
 
 import dask
 
+from distributed.batched import BatchedSend
 from distributed.comm.core import CommClosedError
 from distributed.comm.registry import backends
 from distributed.comm.tcp import TCPBackend, TCPListener
@@ -1380,3 +1381,78 @@ async def test_async_listener_stop(monkeypatch):
         async with Server({}) as s:
             await s.listen(0)
             assert s.listeners
+
+
+@gen_test()
+async def test_messages_are_ordered_bsend():
+    ledger = []
+
+    async def async_handler(val):
+        await asyncio.sleep(0.01)
+        ledger.append(val)
+
+    def sync_handler(val):
+        ledger.append(val)
+
+    async with Server(
+        {},
+        stream_handlers={
+            "sync_handler": sync_handler,
+            "async_handler": async_handler,
+        },
+    ) as s:
+        await s.listen()
+        comm = await connect(s.address)
+        try:
+            b = BatchedSend(interval=10)
+            try:
+                await comm.write({"op": "connection_stream"})
+                b.start(comm)
+                n = 100
+                for ix in range(n):
+                    if ix % 2:
+                        b.send({"op": "sync_handler", "val": ix})
+                    else:
+                        b.send({"op": "async_handler", "val": ix})
+                while not len(ledger) == n:
+                    await asyncio.sleep(0.01)
+                assert ledger == list(range(n))
+            finally:
+                await b.close()
+        finally:
+            await comm.close()
+
+
+@gen_test()
+async def test_messages_are_ordered_raw():
+    ledger = []
+
+    async def async_handler(val):
+        await asyncio.sleep(0.01)
+        ledger.append(val)
+
+    def sync_handler(val):
+        ledger.append(val)
+
+    async with Server(
+        {},
+        stream_handlers={
+            "sync_handler": sync_handler,
+            "async_handler": async_handler,
+        },
+    ) as s:
+        await s.listen()
+        comm = await connect(s.address)
+        try:
+            await comm.write({"op": "connection_stream"})
+            n = 100
+            for ix in range(n):
+                if ix % 2:
+                    await comm.write({"op": "sync_handler", "val": ix})
+                else:
+                    await comm.write({"op": "async_handler", "val": ix})
+            while not len(ledger) == n:
+                await asyncio.sleep(0.01)
+            assert ledger == list(range(n))
+        finally:
+            await comm.close()