GH-74116: Allow multiple drain waiters for asyncio.StreamWriter (GH-94705) (#96395)

(cherry picked from commit e5b2453e61ba5376831093236d598ef5f9f1de61)

Co-authored-by: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>

Co-authored-by: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>

1 files changed, 19 insertions, 0 deletions

diff --git a/Lib/test/test_asyncio/test_streams.py b/Lib/test/test_asyncio/test_streams.py
index 098a0da..0c49099 100644
--- a/Lib/test/test_asyncio/test_streams.py
+++ b/Lib/test/test_asyncio/test_streams.py
@@ -864,6 +864,25 @@ os.close(fd)
         self.assertEqual(cm.filename, __file__)
         self.assertIs(protocol._loop, self.loop)
 
+    def test_multiple_drain(self):
+        # See https://github.com/python/cpython/issues/74116
+        drained = 0
+
+        async def drainer(stream):
+            nonlocal drained
+            await stream._drain_helper()
+            drained += 1
+
+        async def main():
+            loop = asyncio.get_running_loop()
+            stream = asyncio.streams.FlowControlMixin(loop)
+            stream.pause_writing()
+            loop.call_later(0.1, stream.resume_writing)
+            await asyncio.gather(*[drainer(stream) for _ in range(10)])
+            self.assertEqual(drained, 10)
+
+        self.loop.run_until_complete(main())
+
     def test_drain_raises(self):
         # See http://bugs.python.org/issue25441