Provide some functionality for closing streams.

Author: Lukasa

hyper/http20/connection.py

@@ -329,8 +329,8 @@ def _new_stream(self):
         Returns a new stream object for this connection.
         """
         s = Stream(
-            self.next_stream_id, self._send_cb, self._recv_cb, self.encoder,
-            self.decoder
+            self.next_stream_id, self._send_cb, self._recv_cb,
+            self._close_stream, self.encoder, self.decoder
         )
         s._out_flow_control_window = (
             self._settings[SettingsFrame.INITIAL_WINDOW_SIZE]
@@ -339,6 +339,12 @@ def _new_stream(self):
 
         return s
 
+    def _close_stream(self, stream_id):
+        """
+        Called by a stream when it would like to be 'closed'.
+        """
+        del self.streams[stream_id]
+
     def _send_cb(self, frame):
         """
         This is the callback used by streams to send data on the connection.

hyper/http20/response.py

@@ -155,3 +155,11 @@ def fileno(self):
         currently not implemented.
         """
         pass
+
+    def close(self):
+        """
+        Close the response. In effect this closes the backing HTTP/2.0 stream.
+
+        :returns: Nothing.
+        """
+        self._stream.close()

hyper/http20/stream.py

@@ -48,6 +48,7 @@ def __init__(self,
                  stream_id,
                  data_cb,
                  recv_cb,
+                 close_cb,
                  header_encoder,
                  header_decoder):
         self.stream_id = stream_id
@@ -69,6 +70,9 @@ def __init__(self,
         # connection.
         self._recv_cb = recv_cb
 
+        # This is the callback to be called when the stream is closed.
+        self._close_cb = close_cb
+
         # A reference to the header encoder and decoder objects belonging to
         # the parent connection.
         self._encoder = header_encoder
@@ -230,6 +234,17 @@ def getresponse(self):
         # Create the HTTP response.
         return HTTP20Response(headers, self)
 
+    def close(self):
+        """
+        Closes the stream. If the stream is currently open, attempts to close
+        it as gracefully as possible.
+
+        :returns: Nothing.
+        """
+        # Right now let's not bother with grace, let's just call close on the
+        # connection.
+        self._close_cb(self.stream_id)
+
     def _send_chunk(self, data, final):
         """
         Implements most of the sending logic.

test/test_hyper.py

@@ -860,18 +860,35 @@ def test_connection_doesnt_send_window_update_on_zero_length_data_frame(self):
         # No frame should have been sent on the connection.
         assert len(sock.queue) == 0
 
+    def test_streams_are_cleared_from_connections_on_close(self):
+        # Prepare a socket so we can open a stream.
+        sock = DummySocket()
+        c = HTTP20Connection('www.google.com')
+        c._sock = sock
+
+        # Open a request (which creates a stream)
+        c.request('GET', '/')
+
+        # Close the stream.
+        c.streams[1].close()
+
+        # There should be nothing left, but the next stream ID should be
+        # unchanged.
+        assert not c.streams
+        assert c.next_stream_id == 3
+
 
 class TestHyperStream(object):
     def test_streams_have_ids(self):
-        s = Stream(1, None, None, None, None)
+        s = Stream(1, None, None, None, None, None)
         assert s.stream_id == 1
 
     def test_streams_initially_have_no_headers(self):
-        s = Stream(1, None, None, None, None)
+        s = Stream(1, None, None, None, None, None)
         assert s.headers == []
 
     def test_streams_can_have_headers(self):
-        s = Stream(1, None, None, None, None)
+        s = Stream(1, None, None, None, None, None)
         s.add_header("name", "value")
         assert s.headers == [("name", "value")]
 
@@ -881,14 +898,14 @@ def data_callback(frame):
             assert frame.data == 'testkeyTestVal'
             assert frame.flags == set(['END_STREAM', 'END_HEADERS'])
 
-        s = Stream(1, data_callback, None, NullEncoder, None)
+        s = Stream(1, data_callback, None, None, NullEncoder, None)
         s.add_header("TestKey", "TestVal")
         s.open(True)
 
         assert s.state == STATE_HALF_CLOSED_LOCAL
 
     def test_receiving_a_frame_queues_it(self):
-        s = Stream(1, None, None, None, None)
+        s = Stream(1, None, None, None, None, None)
         s.receive_frame(Frame(0))
         assert len(s._queued_frames) == 1
 
@@ -898,7 +915,7 @@ def data_callback(frame):
             assert frame.data == b'Hi there!'
             assert frame.flags == set(['END_STREAM'])
 
-        s = Stream(1, data_callback, None, NullEncoder, None)
+        s = Stream(1, data_callback, None, None, NullEncoder, None)
         s.state = STATE_OPEN
         s.send_data(BytesIO(b'Hi there!'), True)
 
@@ -917,7 +934,7 @@ def data_callback(frame):
 
         data = b'test' * (MAX_CHUNK + 1)
 
-        s = Stream(1, data_callback, None, NullEncoder, None)
+        s = Stream(1, data_callback, None, None, NullEncoder, None)
         s.state = STATE_OPEN
         s.send_data(BytesIO(data), True)
 
@@ -932,7 +949,7 @@ def data_callback(frame):
             assert frame.data == b'Hi there!'
             assert frame.flags == set(['END_STREAM'])
 
-        s = Stream(1, data_callback, None, NullEncoder, None)
+        s = Stream(1, data_callback, None, None, NullEncoder, None)
         s.state = STATE_OPEN
         s.send_data(b'Hi there!', True)
 
@@ -951,7 +968,7 @@ def data_callback(frame):
 
         data = b'test' * (MAX_CHUNK + 1)
 
-        s = Stream(1, data_callback, None, NullEncoder, None)
+        s = Stream(1, data_callback, None, None, NullEncoder, None)
         s.state = STATE_OPEN
         s.send_data(data, True)
 
@@ -961,7 +978,7 @@ def data_callback(frame):
         assert s._out_flow_control_window == 65535 - len(data)
 
     def test_windowupdate_frames_update_windows(self):
-        s = Stream(1, None, None, None, None)
+        s = Stream(1, None, None, None, None, None)
         f = WindowUpdateFrame(1)
         f.window_increment = 1000
         s.receive_frame(f)
@@ -980,7 +997,7 @@ def inner():
                 s.receive_frame(in_frames.pop(0))
             return inner
 
-        s = Stream(1, send_cb, None, None, None)
+        s = Stream(1, send_cb, None, None, None, None)
         s._recv_cb = recv_cb(s)
         s.state = STATE_HALF_CLOSED_LOCAL
 
@@ -1006,7 +1023,7 @@ def inner():
                 s.receive_frame(in_frames.pop(0))
             return inner
 
-        s = Stream(1, send_cb, None, None, None)
+        s = Stream(1, send_cb, None, None, None, None)
         s._recv_cb = recv_cb(s)
         s.state = STATE_HALF_CLOSED_LOCAL
 
@@ -1038,7 +1055,7 @@ def inner():
                 s.receive_frame(in_frames.pop(0))
             return inner
 
-        s = Stream(1, send_cb, None, None, None)
+        s = Stream(1, send_cb, None, None, None, None)
         s._recv_cb = recv_cb(s)
         s.state = STATE_HALF_CLOSED_LOCAL
 
@@ -1099,6 +1116,13 @@ def test_response_transparently_decrypts_wrong_deflate(self):
 
         assert resp.read() == b'this is test data'
 
+    def test_response_calls_stream_close(self):
+        stream = DummyStream('')
+        resp = HTTP20Response({':status': '200'}, stream)
+        resp.close()
+
+        assert stream.closed
+
 
 class TestHTTP20Adapter(object):
     def test_adapter_reuses_connections(self):
@@ -1131,6 +1155,13 @@ def close(self):
 class DummyStream(object):
     def __init__(self, data):
         self.data = data
+        self.closed = False
 
     def _read(self, *args, **kwargs):
         return self.data
+
+    def close(self):
+        if not self.closed:
+            self.closed = True
+        else:
+            assert False

test/test_integration.py

@@ -15,6 +15,11 @@
 from hyper.http20.frame import (
     Frame, SettingsFrame, WindowUpdateFrame, DataFrame, HeadersFrame
 )
+from hyper.http20.hpack import Encoder
+from hyper.http20.huffman import HuffmanEncoder
+from hyper.http20.huffman_constants import (
+    RESPONSE_CODES, RESPONSE_CODES_LENGTH
+)
 from server import SocketLevelTest
 
 # Turn off certificate verification for the tests.
@@ -27,6 +32,16 @@ def decode_frame(frame_data):
     assert 8 + length == len(frame_data)
     return f
 
+
+def build_headers_frame(headers):
+    f = HeadersFrame(1)
+    e = Encoder()
+    e.huffman_coder = HuffmanEncoder(RESPONSE_CODES, RESPONSE_CODES_LENGTH)
+    f.data = e.encode(headers)
+    f.flags.add('END_HEADERS')
+    return f
+
+
 class TestHyperIntegration(SocketLevelTest):
     def test_connection_string(self):
         self.set_up()
@@ -200,6 +215,42 @@ def socket_handler(listener):
 
         self.tear_down()
 
+    def test_closed_responses_remove_their_streams_from_conn(self):
+        self.set_up()
+
+        recv_event = threading.Event()
+
+        def socket_handler(listener):
+            sock = listener.accept()[0]
+
+            # We're going to get the two messages for the connection open, then
+            # a headers frame.
+            sock.recv(65535)
+            sock.recv(65535)
+            sock.send(SettingsFrame(0).serialize())
+            sock.recv(65535)
+
+            # Now, send the headers for the response.
+            f = build_headers_frame([(':status', '200')])
+            f.stream_id = 1
+            sock.send(f.serialize())
+
+            # Wait for the message from the main thread.
+            recv_event.wait()
+            sock.close()
+
+        self._start_server(socket_handler)
+        conn = HTTP20Connection(self.host, self.port)
+        conn.request('GET', '/')
+        resp = conn.getresponse()
+
+        # Close the response.
+        resp.close()
+
+        recv_event.set()
+
+        assert not conn.streams
+
 
 class TestRequestsAdapter(SocketLevelTest):
     def test_adapter_received_values(self):