clear the frame fragment cache on finishing the stream

Author: jell-o-fishi

rsocket/frame_fragment_cache.py

@@ -8,7 +8,7 @@ class FrameFragmentCache:
     __slots__ = '_frames_by_stream_id'
 
     def __init__(self):
-        self._frames_by_stream_id: Dict[str, FragmentableFrame] = {}
+        self._frames_by_stream_id: Dict[int, FragmentableFrame] = {}
 
     def append(self, frame: FragmentableFrame) -> Optional[FragmentableFrame]:
         if frame.flags_follows:
@@ -20,6 +20,9 @@ def append(self, frame: FragmentableFrame) -> Optional[FragmentableFrame]:
                 self._frames_by_stream_id.pop(frame.stream_id)
             return frame
 
+    def remove(self, stream_id: int):
+        self._frames_by_stream_id.pop(stream_id, None)
+
     def _frame_fragment_builder(self, next_fragment: FragmentableFrame) -> FragmentableFrame:
 
         current_frame_from_fragments = self._frames_by_stream_id.get(next_fragment.stream_id)

rsocket/rsocket_base.py

@@ -149,6 +149,7 @@ def _get_first_stream_id(self) -> int:
 
     def finish_stream(self, stream_id: int):
         self._stream_control.finish_stream(stream_id)
+        self._frame_fragment_cache.remove(stream_id)
 
     def send_request(self, frame: RequestFrame):
         if self._honor_lease and not self._is_frame_allowed_to_send(frame):