Repository: tornadoweb/tornado

Issue:
Add IOStream.read_into()
For better performance on large reads, it would be useful to expose a `readinto` or `recv_into`-like operation on `IOStream`. The API could read like:
```python
    def read_into(self, buf, callback=None, partial=False):
        """Asynchronously read a number of bytes.

        ``buf`` must be a writable buffer into which data will be read.
        If a callback is given, it will be run with the number of read
        bytes as an argument; if not, this method returns a `.Future`.

        If ``partial`` is true, the callback is run as soon as any bytes
        have been read.  Otherwise, it is run when the ``buf`` has been
        entirely filled with read data.
        """
```

@bdarnell, what do you think? Is this something you would like to see? See performance analysis at https://mail.python.org/pipermail/async-sig/2017-October/000392.html



Discussion:
This API looks reasonable (it looks just like the Go read API, which makes me wonder if some sort of delayed-allocation interface is necessary to address issues like [this](https://eklitzke.org/goroutines-nonblocking-io-and-memory-usage) or if the ordinary non-`into` APIs are sufficient). 

But I wonder whether we should keep adding more complexity to IOStream to support this use case or start from scratch with a performance-focused interface. Aside from blasting bits onto the wire, what other parts of IOStream do you use? Do you need the other read modes, TLS, HTTPServer integration?

Or on the other hand, if we're going to add support for caller-supplied buffers, should this be an option in *all* read modes and not just `read_bytes`? Should `read_into` be a separate function or a `buf` option to `read_bytes` (there are some subtleties there like the need to return the number of bytes read somehow in the partial case, so combining them into the same method may be too much)?
> which makes me wonder if some sort of delayed-allocation interface is necessary to address issues like this or if the ordinary non-into APIs are sufficient

In the general case I don't know. In our case, we have a framing protocol, so we first read 8 bytes of length and would only allocate the buffer afterwards. At that point we are sure that the other end is sending data to us.

> Aside from blasting bits onto the wire, what other parts of IOStream do you use? Do you need the other read modes, TLS, HTTPServer integration?

We do use TLS. We may use `read_bytes` for tiny transfers out of convenience. The rest (`read_until`, `read_regex`, `read_until_close`, `HTTPServer` integration) we don't. However, I know for a fact that [Bokeh](https://bokeh.pydata.org/en/latest/), which uses Tornado, would also like to improve websockets performance with large transfers, and I'm not sure what that entails.

> Or on the other hand, if we're going to add support for caller-supplied buffers, should this be an option in all read modes and not just read_bytes? Should read_into be a separate function or a buf option to read_bytes

I think it mostly makes sense as a separate method. From an implementation point of view, it's also probably much simpler like that.
> However, I know for a fact that Bokeh, which uses Tornado, would also like to improve websockets performance with large transfers, and I'm not sure what that entails.

No benchmarking has been done here.  I would not be surprised if they have some other bottleneck.

Resolving patch:
diff --git a/tornado/iostream.py b/tornado/iostream.py
--- a/tornado/iostream.py
+++ b/tornado/iostream.py
@@ -254,6 +254,8 @@ def __init__(self, max_buffer_size=None,
         self._read_buffer = bytearray()
         self._read_buffer_pos = 0
         self._read_buffer_size = 0
+        self._user_read_buffer = False
+        self._after_user_read_buffer = None
         self._write_buffer = _StreamBuffer()
         self._total_write_index = 0
         self._total_write_done_index = 0
@@ -300,13 +302,18 @@ def write_to_fd(self, data):
         """
         raise NotImplementedError()
 
-    def read_from_fd(self):
+    def read_from_fd(self, buf):
         """Attempts to read from the underlying file.
 
-        Returns ``None`` if there was nothing to read (the socket
-        returned `~errno.EWOULDBLOCK` or equivalent), otherwise
-        returns the data.  When possible, should return no more than
-        ``self.read_chunk_size`` bytes at a time.
+        Reads up to ``len(buf)`` bytes, storing them in the buffer.
+        Returns the number of bytes read. Returns None if there was
+        nothing to read (the socket returned `~errno.EWOULDBLOCK` or
+        equivalent), and zero on EOF.
+
+        .. versionchanged:: 5.0
+
+           Interface redesigned to take a buffer and return a number
+           of bytes instead of a freshly-allocated object.
         """
         raise NotImplementedError()
 
@@ -415,6 +422,50 @@ def read_bytes(self, num_bytes, callback=None, streaming_callback=None,
             raise
         return future
 
+    def read_into(self, buf, callback=None, partial=False):
+        """Asynchronously read a number of bytes.
+
+        ``buf`` must be a writable buffer into which data will be read.
+        If a callback is given, it will be run with the number of read
+        bytes as an argument; if not, this method returns a `.Future`.
+
+        If ``partial`` is true, the callback is run as soon as any bytes
+        have been read.  Otherwise, it is run when the ``buf`` has been
+        entirely filled with read data.
+
+        .. versionadded:: 5.0
+        """
+        future = self._set_read_callback(callback)
+
+        # First copy data already in read buffer
+        available_bytes = self._read_buffer_size
+        n = len(buf)
+        if available_bytes >= n:
+            end = self._read_buffer_pos + n
+            buf[:] = memoryview(self._read_buffer)[self._read_buffer_pos:end]
+            del self._read_buffer[:end]
+            self._after_user_read_buffer = self._read_buffer
+        elif available_bytes > 0:
+            buf[:available_bytes] = memoryview(self._read_buffer)[self._read_buffer_pos:]
+
+        # Set up the supplied buffer as our temporary read buffer.
+        # The original (if it had any data remaining) has been
+        # saved for later.
+        self._user_read_buffer = True
+        self._read_buffer = buf
+        self._read_buffer_pos = 0
+        self._read_buffer_size = available_bytes
+        self._read_bytes = n
+        self._read_partial = partial
+
+        try:
+            self._try_inline_read()
+        except:
+            if future is not None:
+                future.add_done_callback(lambda f: f.exception())
+            raise
+        return future
+
     def read_until_close(self, callback=None, streaming_callback=None):
         """Asynchronously reads all data from the socket until it is closed.
 
@@ -762,6 +813,15 @@ def _set_read_callback(self, callback):
         return self._read_future
 
     def _run_read_callback(self, size, streaming):
+        if self._user_read_buffer:
+            self._read_buffer = self._after_user_read_buffer or bytearray()
+            self._after_user_read_buffer = None
+            self._read_buffer_pos = 0
+            self._read_buffer_size = len(self._read_buffer)
+            self._user_read_buffer = False
+            result = size
+        else:
+            result = self._consume(size)
         if streaming:
             callback = self._streaming_callback
         else:
@@ -771,10 +831,11 @@ def _run_read_callback(self, size, streaming):
                 assert callback is None
                 future = self._read_future
                 self._read_future = None
-                future.set_result(self._consume(size))
+
+                future.set_result(result)
         if callback is not None:
             assert (self._read_future is None) or streaming
-            self._run_callback(callback, self._consume(size))
+            self._run_callback(callback, result)
         else:
             # If we scheduled a callback, we will add the error listener
             # afterwards.  If we didn't, we have to do it now.
@@ -820,31 +881,44 @@ def _read_to_buffer(self):
         to read (i.e. the read returns EWOULDBLOCK or equivalent).  On
         error closes the socket and raises an exception.
         """
-        while True:
-            try:
-                chunk = self.read_from_fd()
-            except (socket.error, IOError, OSError) as e:
-                if errno_from_exception(e) == errno.EINTR:
-                    continue
-                # ssl.SSLError is a subclass of socket.error
-                if self._is_connreset(e):
-                    # Treat ECONNRESET as a connection close rather than
-                    # an error to minimize log spam  (the exception will
-                    # be available on self.error for apps that care).
+        try:
+            while True:
+                try:
+                    if self._user_read_buffer:
+                        buf = memoryview(self._read_buffer)[self._read_buffer_size:]
+                    else:
+                        buf = bytearray(self.read_chunk_size)
+                    bytes_read = self.read_from_fd(buf)
+                except (socket.error, IOError, OSError) as e:
+                    if errno_from_exception(e) == errno.EINTR:
+                        continue
+                    # ssl.SSLError is a subclass of socket.error
+                    if self._is_connreset(e):
+                        # Treat ECONNRESET as a connection close rather than
+                        # an error to minimize log spam  (the exception will
+                        # be available on self.error for apps that care).
+                        self.close(exc_info=e)
+                        return
                     self.close(exc_info=e)
-                    return
-                self.close(exc_info=e)
-                raise
-            break
-        if chunk is None:
-            return 0
-        self._read_buffer += chunk
-        self._read_buffer_size += len(chunk)
+                    raise
+                break
+            if bytes_read is None:
+                return 0
+            elif bytes_read == 0:
+                self.close()
+                return 0
+            if not self._user_read_buffer:
+                self._read_buffer += memoryview(buf)[:bytes_read]
+            self._read_buffer_size += bytes_read
+        finally:
+            # Break the reference to buf so we don't waste a chunk's worth of
+            # memory in case an exception hangs on to our stack frame.
+            buf = None
         if self._read_buffer_size > self.max_buffer_size:
             gen_log.error("Reached maximum read buffer size")
             self.close()
             raise StreamBufferFullError("Reached maximum read buffer size")
-        return len(chunk)
+        return bytes_read
 
     def _run_streaming_callback(self):
         if self._streaming_callback is not None and self._read_buffer_size:
@@ -1106,18 +1180,16 @@ def get_fd_error(self):
                                        socket.SO_ERROR)
         return socket.error(errno, os.strerror(errno))
 
-    def read_from_fd(self):
+    def read_from_fd(self, buf):
         try:
-            chunk = self.socket.recv(self.read_chunk_size)
+            return self.socket.recv_into(buf)
         except socket.error as e:
             if e.args[0] in _ERRNO_WOULDBLOCK:
                 return None
             else:
                 raise
-        if not chunk:
-            self.close()
-            return None
-        return chunk
+        finally:
+            buf = None
 
     def write_to_fd(self, data):
         try:
@@ -1528,35 +1600,29 @@ def write_to_fd(self, data):
             # See https://github.com/tornadoweb/tornado/pull/2008
             del data
 
-    def read_from_fd(self):
-        if self._ssl_accepting:
-            # If the handshake hasn't finished yet, there can't be anything
-            # to read (attempting to read may or may not raise an exception
-            # depending on the SSL version)
-            return None
+    def read_from_fd(self, buf):
         try:
-            # SSLSocket objects have both a read() and recv() method,
-            # while regular sockets only have recv().
-            # The recv() method blocks (at least in python 2.6) if it is
-            # called when there is nothing to read, so we have to use
-            # read() instead.
-            chunk = self.socket.read(self.read_chunk_size)
-        except ssl.SSLError as e:
-            # SSLError is a subclass of socket.error, so this except
-            # block must come first.
-            if e.args[0] == ssl.SSL_ERROR_WANT_READ:
-                return None
-            else:
-                raise
-        except socket.error as e:
-            if e.args[0] in _ERRNO_WOULDBLOCK:
+            if self._ssl_accepting:
+                # If the handshake hasn't finished yet, there can't be anything
+                # to read (attempting to read may or may not raise an exception
+                # depending on the SSL version)
                 return None
-            else:
-                raise
-        if not chunk:
-            self.close()
-            return None
-        return chunk
+            try:
+                return self.socket.recv_into(buf)
+            except ssl.SSLError as e:
+                # SSLError is a subclass of socket.error, so this except
+                # block must come first.
+                if e.args[0] == ssl.SSL_ERROR_WANT_READ:
+                    return None
+                else:
+                    raise
+            except socket.error as e:
+                if e.args[0] in _ERRNO_WOULDBLOCK:
+                    return None
+                else:
+                    raise
+        finally:
+            buf = None
 
     def _is_connreset(self, e):
         if isinstance(e, ssl.SSLError) and e.args[0] == ssl.SSL_ERROR_EOF:
@@ -1592,9 +1658,9 @@ def write_to_fd(self, data):
             # See https://github.com/tornadoweb/tornado/pull/2008
             del data
 
-    def read_from_fd(self):
+    def read_from_fd(self, buf):
         try:
-            chunk = self._fio.read(self.read_chunk_size)
+            return self._fio.readinto(buf)
         except (IOError, OSError) as e:
             if errno_from_exception(e) == errno.EBADF:
                 # If the writing half of a pipe is closed, select will
@@ -1603,13 +1669,8 @@ def read_from_fd(self):
                 return None
             else:
                 raise
-        if chunk is None:
-            # Read would block
-            return None
-        if not chunk:
-            self.close()
-            return None
-        return chunk
+        finally:
+            buf = None
 
 
 def doctests():


Test patch:
diff --git a/tornado/test/iostream_test.py b/tornado/test/iostream_test.py
--- a/tornado/test/iostream_test.py
+++ b/tornado/test/iostream_test.py
@@ -12,6 +12,7 @@
 from tornado.test.util import unittest, skipIfNonUnix, refusing_port, skipPypy3V58
 from tornado.web import RequestHandler, Application
 import errno
+import hashlib
 import logging
 import os
 import platform
@@ -674,6 +675,143 @@ def test_flow_control(self):
             rs.close()
             ws.close()
 
+    def test_read_into(self):
+        rs, ws = self.make_iostream_pair()
+
+        def sleep_some():
+            self.io_loop.run_sync(lambda: gen.sleep(0.05))
+        try:
+            buf = bytearray(10)
+            rs.read_into(buf, callback=self.stop)
+            ws.write(b"hello")
+            sleep_some()
+            self.assertTrue(rs.reading())
+            ws.write(b"world!!")
+            data = self.wait()
+            self.assertFalse(rs.reading())
+            self.assertEqual(data, 10)
+            self.assertEqual(bytes(buf), b"helloworld")
+
+            # Existing buffer is fed into user buffer
+            rs.read_into(buf, callback=self.stop)
+            sleep_some()
+            self.assertTrue(rs.reading())
+            ws.write(b"1234567890")
+            data = self.wait()
+            self.assertFalse(rs.reading())
+            self.assertEqual(data, 10)
+            self.assertEqual(bytes(buf), b"!!12345678")
+
+            # Existing buffer can satisfy read immediately
+            buf = bytearray(4)
+            ws.write(b"abcdefghi")
+            rs.read_into(buf, callback=self.stop)
+            data = self.wait()
+            self.assertEqual(data, 4)
+            self.assertEqual(bytes(buf), b"90ab")
+
+            rs.read_bytes(7, self.stop)
+            data = self.wait()
+            self.assertEqual(data, b"cdefghi")
+        finally:
+            ws.close()
+            rs.close()
+
+    def test_read_into_partial(self):
+        rs, ws = self.make_iostream_pair()
+
+        def sleep_some():
+            self.io_loop.run_sync(lambda: gen.sleep(0.05))
+        try:
+            # Partial read
+            buf = bytearray(10)
+            rs.read_into(buf, callback=self.stop, partial=True)
+            ws.write(b"hello")
+            data = self.wait()
+            self.assertFalse(rs.reading())
+            self.assertEqual(data, 5)
+            self.assertEqual(bytes(buf), b"hello\0\0\0\0\0")
+
+            # Full read despite partial=True
+            ws.write(b"world!1234567890")
+            rs.read_into(buf, callback=self.stop, partial=True)
+            data = self.wait()
+            self.assertEqual(data, 10)
+            self.assertEqual(bytes(buf), b"world!1234")
+
+            # Existing buffer can satisfy read immediately
+            rs.read_into(buf, callback=self.stop, partial=True)
+            data = self.wait()
+            self.assertEqual(data, 6)
+            self.assertEqual(bytes(buf), b"5678901234")
+
+        finally:
+            ws.close()
+            rs.close()
+
+    def test_read_into_zero_bytes(self):
+        rs, ws = self.make_iostream_pair()
+        try:
+            buf = bytearray()
+            fut = rs.read_into(buf)
+            self.assertEqual(fut.result(), 0)
+        finally:
+            ws.close()
+            rs.close()
+
+    def test_many_mixed_reads(self):
+        # Stress buffer handling when going back and forth between
+        # read_bytes() (using an internal buffer) and read_into()
+        # (using a user-allocated buffer).
+        r = random.Random(42)
+        nbytes = 1000000
+        rs, ws = self.make_iostream_pair()
+
+        produce_hash = hashlib.sha1()
+        consume_hash = hashlib.sha1()
+
+        @gen.coroutine
+        def produce():
+            remaining = nbytes
+            while remaining > 0:
+                size = r.randint(1, min(1000, remaining))
+                data = os.urandom(size)
+                produce_hash.update(data)
+                yield ws.write(data)
+                remaining -= size
+            assert remaining == 0
+
+        @gen.coroutine
+        def consume():
+            remaining = nbytes
+            while remaining > 0:
+                if r.random() > 0.5:
+                    # read_bytes()
+                    size = r.randint(1, min(1000, remaining))
+                    data = yield rs.read_bytes(size)
+                    consume_hash.update(data)
+                    remaining -= size
+                else:
+                    # read_into()
+                    size = r.randint(1, min(1000, remaining))
+                    buf = bytearray(size)
+                    n = yield rs.read_into(buf)
+                    assert n == size
+                    consume_hash.update(buf)
+                    remaining -= size
+            assert remaining == 0
+
+        @gen.coroutine
+        def main():
+            yield [produce(), consume()]
+            assert produce_hash.hexdigest() == consume_hash.hexdigest()
+
+        try:
+            self.io_loop.run_sync(main)
+        finally:
+            ws.close()
+            rs.close()
+
 
 class TestIOStreamMixin(TestReadWriteMixin):
     def _make_server_iostream(self, connection, **kwargs):