Add write flow control to unix pipes.

gvanrossum · gvanrossum · commit 0e313f439d8c · 2014-01-28T09:47:27.000-08:00
diff --git a/asyncio/unix_events.py b/asyncio/unix_events.py
@@ -246,7 +246,8 @@ def _call_connection_lost(self, exc):
             self._loop = None
 
 
-class _UnixWritePipeTransport(transports.WriteTransport):
+class _UnixWritePipeTransport(selector_events._FlowControlMixin,
+                              transports.WriteTransport):
 
     def __init__(self, loop, pipe, protocol, waiter=None, extra=None):
         super().__init__(extra)
@@ -277,12 +278,17 @@ def __init__(self, loop, pipe, protocol, waiter=None, extra=None):
         if waiter is not None:
             self._loop.call_soon(waiter.set_result, None)
 
+    def get_write_buffer_size(self):
+        return sum(len(data) for data in self._buffer)
+
     def _read_ready(self):
         # Pipe was closed by peer.
         self._close()
 
     def write(self, data):
-        assert isinstance(data, bytes), repr(data)
+        assert isinstance(data, (bytes, bytearray, memoryview)), repr(data)
+        if isinstance(data, bytearray):
+            data = memoryview(data)
         if not data:
             return
 
@@ -310,6 +316,7 @@ def write(self, data):
             self._loop.add_writer(self._fileno, self._write_ready)
 
         self._buffer.append(data)
+        self._maybe_pause_protocol()
 
     def _write_ready(self):
         data = b''.join(self._buffer)
@@ -329,7 +336,8 @@ def _write_ready(self):
         else:
             if n == len(data):
                 self._loop.remove_writer(self._fileno)
-                if self._closing:
+                self._maybe_resume_protocol()  # May append to buffer.
+                if not self._buffer and self._closing:
                     self._loop.remove_reader(self._fileno)
                     self._call_connection_lost(None)
                 return
