Merge pull request #1165 from hellysmile/limit_closed_streams

Limit amount of _closed_streams

Author: Lukasa

h2/connection.py

@@ -33,7 +33,7 @@
 from .frame_buffer import FrameBuffer
 from .settings import Settings, SettingCodes
 from .stream import H2Stream, StreamClosedBy
-from .utilities import guard_increment_window
+from .utilities import SizeLimitDict, guard_increment_window
 from .windows import WindowManager
 
 
@@ -281,6 +281,9 @@ class H2Connection(object):
     # The initial default value of SETTINGS_MAX_HEADER_LIST_SIZE.
     DEFAULT_MAX_HEADER_LIST_SIZE = 2**16
 
+    # Keep in memory limited amount of results for streams closes
+    MAX_CLOSED_STREAMS = 2**16
+
     def __init__(self, config=None):
         self.state_machine = H2ConnectionStateMachine()
         self.streams = {}
@@ -355,7 +358,9 @@ def __init__(self, config=None):
         # Also used to determine whether we should consider a frame received
         # while a stream is closed as either a stream error or a connection
         # error.
-        self._closed_streams = {}
+        self._closed_streams = SizeLimitDict(
+            size_limit=self.MAX_CLOSED_STREAMS
+        )
 
         # The flow control window manager for the connection.
         self._inbound_flow_control_window_manager = WindowManager(

h2/utilities.py

@@ -617,3 +617,22 @@ def validate_outbound_headers(headers, hdr_validation_flags):
     headers = _check_path_header(headers, hdr_validation_flags)
 
     return headers
+
+
+class SizeLimitDict(collections.OrderedDict):
+
+    def __init__(self, *args, **kwargs):
+        self._size_limit = kwargs.pop("size_limit", None)
+        super(SizeLimitDict, self).__init__(*args, **kwargs)
+
+        self._check_size_limit()
+
+    def __setitem__(self, key, value):
+        super(SizeLimitDict, self).__setitem__(key, value)
+
+        self._check_size_limit()
+
+    def _check_size_limit(self):
+        if self._size_limit is not None:
+            while len(self) > self._size_limit:
+                self.popitem(last=False)

test/test_utility_functions.py

@@ -12,7 +12,7 @@
 import h2.errors
 import h2.events
 import h2.exceptions
-from h2.utilities import extract_method_header
+from h2.utilities import SizeLimitDict, extract_method_header
 
 # These tests require a non-list-returning range function.
 try:
@@ -176,3 +176,51 @@ class TestExtractHeader(object):
     )
     def test_extract_header_method(self, headers):
         assert extract_method_header(headers) == b'GET'
+
+
+def test_size_limit_dict_limit():
+    dct = SizeLimitDict(size_limit=2)
+
+    dct[1] = 1
+    dct[2] = 2
+
+    assert len(dct) == 2
+    assert dct[1] == 1
+    assert dct[2] == 2
+
+    dct[3] = 3
+
+    assert len(dct) == 2
+    assert dct[2] == 2
+    assert dct[3] == 3
+    assert 1 not in dct
+
+
+def test_size_limit_dict_limit_init():
+    initial_dct = {
+        1: 1,
+        2: 2,
+        3: 3,
+    }
+
+    dct = SizeLimitDict(initial_dct, size_limit=2)
+
+    assert len(dct) == 2
+
+
+def test_size_limit_dict_no_limit():
+    dct = SizeLimitDict(size_limit=None)
+
+    dct[1] = 1
+    dct[2] = 2
+
+    assert len(dct) == 2
+    assert dct[1] == 1
+    assert dct[2] == 2
+
+    dct[3] = 3
+
+    assert len(dct) == 3
+    assert dct[1] == 1
+    assert dct[2] == 2
+    assert dct[3] == 3