Modify _drain_helper() to handle parallel calls without race-condition (#6028)

Traktormaster · asvetlov · web-flow · commit 3eea7b8dbd90 · 2021-10-27T15:59:19.000Z
* Modify _drain_helper() to handle parallel calls of _send_frame() without race-condition.

* Update CHANGES/2934.bugfix

* Update 2934.bugfix

Co-authored-by: Andrew Svetlov &lt;andrew.svetlov@gmail.com&gt;
diff --git a/CHANGES/2934.bugfix b/CHANGES/2934.bugfix
@@ -0,0 +1 @@
+Modify _drain_helper() to handle concurrent `await resp.write(...)` or `ws.send_json(...)` calls without race-condition.
diff --git a/CONTRIBUTORS.txt b/CONTRIBUTORS.txt
@@ -230,6 +230,7 @@ Navid Sheikhol
 Nicolas Braem
 Nikolay Kim
 Nikolay Novik
+Nándor Mátravölgyi
 Oisin Aylward
 Olaf Conradi
 Oleg Höfling
diff --git a/aiohttp/base_protocol.py b/aiohttp/base_protocol.py
@@ -81,7 +81,7 @@ async def _drain_helper(self) -> None:
         if not self._paused:
             return
         waiter = self._drain_waiter
-        assert waiter is None or waiter.cancelled()
-        waiter = self._loop.create_future()
-        self._drain_waiter = waiter
-        await waiter
+        if waiter is None:
+            waiter = self._loop.create_future()
+            self._drain_waiter = waiter
+        await asyncio.shield(waiter)
diff --git a/tests/test_base_protocol.py b/tests/test_base_protocol.py
@@ -180,3 +180,21 @@ async def wait() -> None:
     with suppress(asyncio.CancelledError):
         await t
     assert pr._drain_waiter is None
+
+
+async def test_parallel_drain_race_condition() -> None:
+    loop = asyncio.get_event_loop()
+    pr = BaseProtocol(loop=loop)
+    tr = mock.Mock()
+    pr.connection_made(tr)
+    pr.pause_writing()
+
+    ts = [loop.create_task(pr._drain_helper()) for _ in range(5)]
+    assert not (await asyncio.wait(ts, timeout=0.5))[
+        0
+    ], "All draining tasks must be pending"
+
+    assert pr._drain_waiter is not None
+    pr.resume_writing()
+    await asyncio.gather(*ts)
+    assert pr._drain_waiter is None

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Modify _drain_helper() to handle concurrent `await resp.write(...)` or `ws.send_json(...)` calls without race-condition.