From: Tom Christie <tom@tomchristie.com>
Date: Tue, 31 Dec 2019 13:18:00 +0000 (+0000)
Subject: Lock around stream.write and stream.close operations (#699)
X-Git-Tag: 0.11.0~30
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=56afde1f9f84311b9582b413b6e70bb9973e5921;p=thirdparty%2Fhttpx.git

Lock around stream.write and stream.close operations (#699)
---

diff --git a/httpx/backends/asyncio.py b/httpx/backends/asyncio.py
index 1859e96a..854618e1 100644
--- a/httpx/backends/asyncio.py
+++ b/httpx/backends/asyncio.py
@@ -77,6 +77,7 @@ class SocketStream(BaseSocketStream):
         self.stream_reader = stream_reader
         self.stream_writer = stream_writer
         self.read_lock = asyncio.Lock()
+        self.write_lock = asyncio.Lock()
 
         self._inner: typing.Optional[SocketStream] = None
 
@@ -135,11 +136,12 @@ class SocketStream(BaseSocketStream):
         if not data:
             return
 
-        self.stream_writer.write(data)
         try:
-            return await asyncio.wait_for(
-                self.stream_writer.drain(), timeout.write_timeout
-            )
+            async with self.write_lock:
+                self.stream_writer.write(data)
+                return await asyncio.wait_for(
+                    self.stream_writer.drain(), timeout.write_timeout
+                )
         except asyncio.TimeoutError:
             raise WriteTimeout() from None
 
@@ -168,7 +170,8 @@ class SocketStream(BaseSocketStream):
         # This is fine, though, because '.close()' schedules the actual closing of the
         # stream, meaning that at best it will happen during the next event loop
         # iteration, and at worst asyncio will take care of it on program exit.
-        self.stream_writer.close()
+        async with self.write_lock:
+            self.stream_writer.close()
 
 
 class AsyncioBackend(ConcurrencyBackend):
diff --git a/httpx/backends/trio.py b/httpx/backends/trio.py
index 7c26dae3..6e79dc8b 100644
--- a/httpx/backends/trio.py
+++ b/httpx/backends/trio.py
@@ -79,7 +79,8 @@ class SocketStream(BaseSocketStream):
         return stream.socket.is_readable()
 
     async def close(self) -> None:
-        await self.stream.aclose()
+        async with self.write_lock:
+            await self.stream.aclose()
 
 
 class TrioBackend(ConcurrencyBackend):