Merge pull request #76 from Lukasa/readahead

Add read-ahead logic.

Author: Lukasa

hyper/http20/bufsocket.py

@@ -11,6 +11,7 @@
 process.
 """
 import select
+from .exceptions import ConnectionResetError
 
 class BufferedSocket(object):
     """
@@ -61,6 +62,20 @@ def _buffer_end(self):
         """
         return self._index + self._bytes_in_buffer
 
+    @property
+    def can_read(self):
+        """
+        Whether or not there is more data to read from the socket.
+        """
+        if self._bytes_in_buffer:
+            return True
+
+        read = select.select([self._sck], [], [], 0)[0]
+        if read:
+            return True
+
+        return False
+
     def recv(self, amt):
         """
         Read some data from the socket.
@@ -105,6 +120,11 @@ def read_all_from_buffer():
         if ((self._remaining_capacity > self._bytes_in_buffer) and
             (should_read)):
             count = self._sck.recv_into(self._buffer_view[self._buffer_end:])
+
+            # The socket just got closed. We should throw an exception if we
+            # were asked for more data than we can return.
+            if not count and amt > self._bytes_in_buffer:
+                raise ConnectionResetError()
             self._bytes_in_buffer += count
 
         # Read out the bytes and update the index.

hyper/http20/connection.py

@@ -15,7 +15,7 @@
 )
 from .response import HTTP20Response, HTTP20Push
 from .window import FlowControlManager
-from .exceptions import ConnectionError
+from .exceptions import ConnectionError, ConnectionResetError
 from .bufsocket import BufferedSocket
 
 import errno
@@ -468,15 +468,13 @@ def _adjust_receive_window(self, frame_len):
 
         return
 
-    def _recv_cb(self):
+    def _consume_single_frame(self):
         """
-        This is the callback used by streams to read data from the connection.
+        Consumes a single frame from the TCP stream.
 
-        It expects to read a single frame, and then to deserialize that frame
-        and pass it to the relevant stream. This is generally called by a
-        stream, not by the connection itself, and it's likely that streams will
-        read a frame that doesn't belong to them. That's ok: streams need to
-        make a decision to spin around again.
+        Right now this method really does a bit too much: it shouldn't be
+        responsible for determining if a frame is valid or to increase the
+        flow control window.
         """
         # Begin by reading 9 bytes from the socket.
         header = self._sock.recv(9)
@@ -522,6 +520,28 @@ def _recv_cb(self):
         else:
             self.receive_frame(frame)
 
+    def _recv_cb(self):
+        """
+        This is the callback used by streams to read data from the connection.
+
+        It expects to read a single frame, and then to deserialize that frame
+        and pass it to the relevant stream. It then attempts to optimistically
+        read further frames (in an attempt to ensure that we see control frames
+        as early as possible).
+
+        This is generally called by a stream, not by the connection itself, and
+        it's likely that streams will read a frame that doesn't belong to them.
+        """
+        self._consume_single_frame()
+        count = 9
+
+        while count and self._sock is not None and self._sock.can_read:
+            # If the connection has been closed, bail out.
+            try:
+                self._consume_single_frame()
+            except ConnectionResetError:
+                break
+
 
     # The following two methods are the implementation of the context manager
     # protocol.

hyper/http20/exceptions.py

@@ -39,3 +39,13 @@ class ProtocolError(HTTP20Error):
     The remote party violated the HTTP/2 protocol.
     """
     pass
+
+
+# Create our own ConnectionResetError.
+try:  # pragma: no cover
+    ConnectionResetError = ConnectionResetError
+except NameError:  # pragma: no cover
+    class ConnectionResetError(Exception):
+        """
+        A HTTP/2 connection was unexpectedly reset.
+        """

test/test_hyper.py

@@ -2029,6 +2029,7 @@ class DummySocket(object):
     def __init__(self):
         self.queue = []
         self.buffer = BytesIO()
+        self.can_read = False
 
     def send(self, data):
         self.queue.append(data)

test/test_socket.py

@@ -63,16 +63,16 @@ def test_receive_multiple_packets_one_at_a_time(self, monkeypatch):
 
         assert d == b'Herebeginsthetestdata'
 
-    def test_receive_empty_packet(self, monkeypatch):
+    def test_receive_small_packets(self, monkeypatch):
         monkeypatch.setattr(
             hyper.http20.bufsocket.select, 'select', dummy_select
         )
         s = DummySocket()
         b = BufferedSocket(s)
-        s.inbound_packets = [b'Here', b'begins', b'', b'the', b'', b'test', b'data']
+        s.inbound_packets = [b'Here', b'begins', b'the', b'test', b'data']
 
         d = b''
-        for _ in range(7):
+        for _ in range(5):
             d += b.recv(100).tobytes()
 
         assert d == b'Herebeginsthetestdata'