Write flow control for asyncio (includes asyncio.streams overhaul).

Lib/asyncio/protocols.py

@@ -29,6 +29,34 @@ class BaseProtocol:
         aborted or closed).
         """
 
+    def pause_writing(self):
+        """Called when the transport's buffer goes over the high-water mark.
+
+        Pause and resume calls are paired -- pause_writing() is called
+        once when the buffer goes strictly over the high-water mark
+        (even if subsequent writes increases the buffer size even
+        more), and eventually resume_writing() is called once when the
+        buffer size reaches the low-water mark.
+
+        Note that if the buffer size equals the high-water mark,
+        pause_writing() is not called -- it must go strictly over.
+        Conversely, resume_writing() is called when the buffer size is
+        equal or lower than the low-water mark.  These end conditions
+        are important to ensure that things go as expected when either
+        mark is zero.
+
+        NOTE: This is the only Protocol callback that is not called
+        through EventLoop.call_soon() -- if it were, it would have no
+        effect when it's most needed (when the app keeps writing
+        without yielding until pause_writing() is called).
+        """
+
+    def resume_writing(self):
+        """Called when the transport's buffer drains below the low-water mark.
+
+        See pause_writing() for details.
+        """
+
 
 class Protocol(BaseProtocol):
     """ABC representing a protocol.

Lib/asyncio/selector_events.py

@@ -346,8 +346,10 @@ class _SelectorTransport(transports.Transport):
         self._buffer = collections.deque()
         self._conn_lost = 0  # Set when call to connection_lost scheduled.
         self._closing = False  # Set when close() called.
-        if server is not None:
-            server.attach(self)
+        self._protocol_paused = False
+        self.set_write_buffer_limits()
+        if self._server is not None:
+            self._server.attach(self)
 
     def abort(self):
         self._force_close(None)
@@ -392,6 +394,40 @@ class _SelectorTransport(transports.Transport):
                 server.detach(self)
                 self._server = None
 
+    def _maybe_pause_protocol(self):
+        size = self.get_write_buffer_size()
+        if size <= self._high_water:
+            return
+        if not self._protocol_paused:
+            self._protocol_paused = True
+            try:
+                self._protocol.pause_writing()
+            except Exception:
+                tulip_log.exception('pause_writing() failed')
+
+    def _maybe_resume_protocol(self):
+        if self._protocol_paused and self.get_write_buffer_size() <= self._low_water:
+            self._protocol_paused = False
+            try:
+                self._protocol.resume_writing()
+            except Exception:
+                tulip_log.exception('resume_writing() failed')
+
+    def set_write_buffer_limits(self, high=None, low=None):
+        if high is None:
+            if low is None:
+                high = 64*1024
+            else:
+                high = 4*low
+        if low is None:
+            low = high // 4
+        assert 0 <= low <= high, repr((low, high))
+        self._high_water = high
+        self._low_water = low
+
+    def get_write_buffer_size(self):
+        return sum(len(data) for data in self._buffer)
+
 
 class _SelectorSocketTransport(_SelectorTransport):
 
@@ -447,7 +483,7 @@ class _SelectorSocketTransport(_SelectorTransport):
             return
 
         if not self._buffer:
-            # Attempt to send it right away first.
+            # Optimization: try to send now.
             try:
                 n = self._sock.send(data)
             except (BlockingIOError, InterruptedError):
@@ -459,34 +495,36 @@ class _SelectorSocketTransport(_SelectorTransport):
                 data = data[n:]
                 if not data:
                     return
-
-            # Start async I/O.
+            # Not all was written; register write handler.
             self._loop.add_writer(self._sock_fd, self._write_ready)
 
+        # Add it to the buffer.
         self._buffer.append(data)
+        self._maybe_pause_protocol()
 
     def _write_ready(self):
         data = b''.join(self._buffer)
         assert data, 'Data should not be empty'
 
-        self._buffer.clear()
+        self._buffer.clear()  # Optimistically; may have to put it back later.
         try:
             n = self._sock.send(data)
         except (BlockingIOError, InterruptedError):
-            self._buffer.append(data)
+            self._buffer.append(data)  # Still need to write this.
         except Exception as exc:
             self._loop.remove_writer(self._sock_fd)
             self._fatal_error(exc)
         else:
             data = data[n:]
-            if not data:
+            if data:
+                self._buffer.append(data)  # Still need to write this.
+            self._maybe_resume_protocol()  # May append to buffer.
+            if not self._buffer:
                 self._loop.remove_writer(self._sock_fd)
                 if self._closing:
                     self._call_connection_lost(None)
                 elif self._eof:
                     self._sock.shutdown(socket.SHUT_WR)
-                return
-            self._buffer.append(data)  # Try again later.
 
     def write_eof(self):
         if self._eof:
@@ -546,16 +584,23 @@ class _SelectorSslTransport(_SelectorTransport):
             self._loop.add_writer(self._sock_fd, self._on_handshake)
             return
         except Exception as exc:
+            self._loop.remove_reader(self._sock_fd)
+            self._loop.remove_writer(self._sock_fd)
             self._sock.close()
             if self._waiter is not None:
                 self._waiter.set_exception(exc)
             return
         except BaseException as exc:
+            self._loop.remove_reader(self._sock_fd)
+            self._loop.remove_writer(self._sock_fd)
             self._sock.close()
             if self._waiter is not None:
                 self._waiter.set_exception(exc)
             raise
 
+        self._loop.remove_reader(self._sock_fd)
+        self._loop.remove_writer(self._sock_fd)
+
         # Verify hostname if requested.
         peercert = self._sock.getpeercert()
         if (self._server_hostname is not None and
@@ -574,8 +619,6 @@ class _SelectorSslTransport(_SelectorTransport):
                            compression=self._sock.compression(),
                            )
 
-        self._loop.remove_reader(self._sock_fd)
-        self._loop.remove_writer(self._sock_fd)
         self._loop.add_reader(self._sock_fd, self._on_ready)
         self._loop.add_writer(self._sock_fd, self._on_ready)
         self._loop.call_soon(self._protocol.connection_made, self)
@@ -642,6 +685,8 @@ class _SelectorSslTransport(_SelectorTransport):
             if n < len(data):
                 self._buffer.append(data[n:])
 
+            self._maybe_resume_protocol()  # May append to buffer.
+
         if self._closing and not self._buffer:
             self._loop.remove_writer(self._sock_fd)
             self._call_connection_lost(None)
@@ -657,8 +702,9 @@ class _SelectorSslTransport(_SelectorTransport):
             self._conn_lost += 1
             return
 
-        self._buffer.append(data)
         # We could optimize, but the callback can do this for now.
+        self._buffer.append(data)
+        self._maybe_pause_protocol()
 
     def can_write_eof(self):
         return False
@@ -675,11 +721,13 @@ class _SelectorDatagramTransport(_SelectorTransport):
 
     def __init__(self, loop, sock, protocol, address=None, extra=None):
         super().__init__(loop, sock, protocol, extra)
-
         self._address = address
         self._loop.add_reader(self._sock_fd, self._read_ready)
         self._loop.call_soon(self._protocol.connection_made, self)
 
+    def get_write_buffer_size(self):
+        return sum(len(data) for data, _ in self._buffer)
+
     def _read_ready(self):
         try:
             data, addr = self._sock.recvfrom(self.max_size)
@@ -723,6 +771,7 @@ class _SelectorDatagramTransport(_SelectorTransport):
                 return
 
         self._buffer.append((data, addr))
+        self._maybe_pause_protocol()
 
     def _sendto_ready(self):
         while self._buffer:
@@ -743,6 +792,7 @@ class _SelectorDatagramTransport(_SelectorTransport):
                 self._fatal_error(exc)
                 return
 
+        self._maybe_resume_protocol()  # May append to buffer.
         if not self._buffer:
             self._loop.remove_writer(self._sock_fd)
             if self._closing:

Lib/asyncio/transports.py

@@ -49,6 +49,31 @@ class ReadTransport(BaseTransport):
 class WriteTransport(BaseTransport):
     """ABC for write-only transports."""
 
+    def set_write_buffer_limits(self, high=None, low=None):
+        """Set the high- and low-water limits for write flow control.
+
+        These two values control when to call the protocol's
+        pause_writing() and resume_writing() methods.  If specified,
+        the low-water limit must be less than or equal to the
+        high-water limit.  Neither value can be negative.
+
+        The defaults are implementation-specific.  If only the
+        high-water limit is given, the low-water limit defaults to a
+        implementation-specific value less than or equal to the
+        high-water limit.  Setting high to zero forces low to zero as
+        well, and causes pause_writing() to be called whenever the
+        buffer becomes non-empty.  Setting low to zero causes
+        resume_writing() to be called only once the buffer is empty.
+        Use of zero for either limit is generally sub-optimal as it
+        reduces opportunities for doing I/O and computation
+        concurrently.
+        """
+        raise NotImplementedError
+
+    def get_write_buffer_size(self):
+        """Return the current size of the write buffer."""
+        raise NotImplementedError
+
     def write(self, data):
         """Write some data bytes to the transport.
 

Lib/test/test_asyncio/test_streams.py

@@ -32,7 +32,7 @@ class StreamReaderTests(unittest.TestCase):
     @unittest.mock.patch('asyncio.streams.events')
     def test_ctor_global_loop(self, m_events):
         stream = streams.StreamReader()
-        self.assertIs(stream.loop, m_events.get_event_loop.return_value)
+        self.assertIs(stream._loop, m_events.get_event_loop.return_value)
 
     def test_open_connection(self):
         with test_utils.run_test_server() as httpd:
@@ -81,13 +81,13 @@ class StreamReaderTests(unittest.TestCase):
         stream = streams.StreamReader(loop=self.loop)
 
         stream.feed_data(b'')
-        self.assertEqual(0, stream.byte_count)
+        self.assertEqual(0, stream._byte_count)
 
     def test_feed_data_byte_count(self):
         stream = streams.StreamReader(loop=self.loop)
 
         stream.feed_data(self.DATA)
-        self.assertEqual(len(self.DATA), stream.byte_count)
+        self.assertEqual(len(self.DATA), stream._byte_count)
 
     def test_read_zero(self):
         # Read zero bytes.
@@ -96,7 +96,7 @@ class StreamReaderTests(unittest.TestCase):
 
         data = self.loop.run_until_complete(stream.read(0))
         self.assertEqual(b'', data)
-        self.assertEqual(len(self.DATA), stream.byte_count)
+        self.assertEqual(len(self.DATA), stream._byte_count)
 
     def test_read(self):
         # Read bytes.
@@ -109,7 +109,7 @@ class StreamReaderTests(unittest.TestCase):
 
         data = self.loop.run_until_complete(read_task)
         self.assertEqual(self.DATA, data)
-        self.assertFalse(stream.byte_count)
+        self.assertFalse(stream._byte_count)
 
     def test_read_line_breaks(self):
         # Read bytes without line breaks.
@@ -120,7 +120,7 @@ class StreamReaderTests(unittest.TestCase):
         data = self.loop.run_until_complete(stream.read(5))
 
         self.assertEqual(b'line1', data)
-        self.assertEqual(5, stream.byte_count)
+        self.assertEqual(5, stream._byte_count)
 
     def test_read_eof(self):
         # Read bytes, stop at eof.
@@ -133,7 +133,7 @@ class StreamReaderTests(unittest.TestCase):
 
         data = self.loop.run_until_complete(read_task)
         self.assertEqual(b'', data)
-        self.assertFalse(stream.byte_count)
+        self.assertFalse(stream._byte_count)
 
     def test_read_until_eof(self):
         # Read all bytes until eof.
@@ -149,7 +149,7 @@ class StreamReaderTests(unittest.TestCase):
         data = self.loop.run_until_complete(read_task)
 
         self.assertEqual(b'chunk1\nchunk2', data)
-        self.assertFalse(stream.byte_count)
+        self.assertFalse(stream._byte_count)
 
     def test_read_exception(self):
         stream = streams.StreamReader(loop=self.loop)
@@ -176,7 +176,7 @@ class StreamReaderTests(unittest.TestCase):
 
         line = self.loop.run_until_complete(read_task)
         self.assertEqual(b'chunk1 chunk2 chunk3 \n', line)
-        self.assertEqual(len(b'\n chunk4')-1, stream.byte_count)
+        self.assertEqual(len(b'\n chunk4')-1, stream._byte_count)
 
     def test_readline_limit_with_existing_data(self):
         stream = streams.StreamReader(3, loop=self.loop)
@@ -185,7 +185,7 @@ class StreamReaderTests(unittest.TestCase):
 
         self.assertRaises(
             ValueError, self.loop.run_until_complete, stream.readline())
-        self.assertEqual([b'line2\n'], list(stream.buffer))
+        self.assertEqual([b'line2\n'], list(stream._buffer))
 
         stream = streams.StreamReader(3, loop=self.loop)
         stream.feed_data(b'li')
@@ -194,8 +194,8 @@ class StreamReaderTests(unittest.TestCase):
 
         self.assertRaises(
             ValueError, self.loop.run_until_complete, stream.readline())
-        self.assertEqual([b'li'], list(stream.buffer))
-        self.assertEqual(2, stream.byte_count)
+        self.assertEqual([b'li'], list(stream._buffer))
+        self.assertEqual(2, stream._byte_count)
 
     def test_readline_limit(self):
         stream = streams.StreamReader(7, loop=self.loop)
@@ -209,8 +209,8 @@ class StreamReaderTests(unittest.TestCase):
 
         self.assertRaises(
             ValueError, self.loop.run_until_complete, stream.readline())
-        self.assertEqual([b'chunk3\n'], list(stream.buffer))
-        self.assertEqual(7, stream.byte_count)
+        self.assertEqual([b'chunk3\n'], list(stream._buffer))
+        self.assertEqual(7, stream._byte_count)
 
     def test_readline_line_byte_count(self):
         stream = streams.StreamReader(loop=self.loop)
@@ -220,7 +220,7 @@ class StreamReaderTests(unittest.TestCase):
         line = self.loop.run_until_complete(stream.readline())
 
         self.assertEqual(b'line1\n', line)
-        self.assertEqual(len(self.DATA) - len(b'line1\n'), stream.byte_count)
+        self.assertEqual(len(self.DATA) - len(b'line1\n'), stream._byte_count)
 
     def test_readline_eof(self):
         stream = streams.StreamReader(loop=self.loop)
@@ -248,7 +248,7 @@ class StreamReaderTests(unittest.TestCase):
         self.assertEqual(b'line2\nl', data)
         self.assertEqual(
             len(self.DATA) - len(b'line1\n') - len(b'line2\nl'),
-            stream.byte_count)
+            stream._byte_count)
 
     def test_readline_exception(self):
         stream = streams.StreamReader(loop=self.loop)
@@ -268,11 +268,11 @@ class StreamReaderTests(unittest.TestCase):
 
         data = self.loop.run_until_complete(stream.readexactly(0))
         self.assertEqual(b'', data)
-        self.assertEqual(len(self.DATA), stream.byte_count)
+        self.assertEqual(len(self.DATA), stream._byte_count)
 
         data = self.loop.run_until_complete(stream.readexactly(-1))
         self.assertEqual(b'', data)
-        self.assertEqual(len(self.DATA), stream.byte_count)
+        self.assertEqual(len(self.DATA), stream._byte_count)
 
     def test_readexactly(self):
         # Read exact number of bytes.
@@ -289,7 +289,7 @@ class StreamReaderTests(unittest.TestCase):
 
         data = self.loop.run_until_complete(read_task)
         self.assertEqual(self.DATA + self.DATA, data)
-        self.assertEqual(len(self.DATA), stream.byte_count)
+        self.assertEqual(len(self.DATA), stream._byte_count)
 
     def test_readexactly_eof(self):
         # Read exact number of bytes (eof).
@@ -304,7 +304,7 @@ class StreamReaderTests(unittest.TestCase):
 
         data = self.loop.run_until_complete(read_task)
         self.assertEqual(self.DATA, data)
-        self.assertFalse(stream.byte_count)
+        self.assertFalse(stream._byte_count)
 
     def test_readexactly_exception(self):
         stream = streams.StreamReader(loop=self.loop)
@@ -357,7 +357,7 @@ class StreamReaderTests(unittest.TestCase):
         # The following line fails if set_exception() isn't careful.
         stream.set_exception(RuntimeError('message'))
         test_utils.run_briefly(self.loop)
-        self.assertIs(stream.waiter, None)
+        self.assertIs(stream._waiter, None)
 
 
 if __name__ == '__main__':