Ensure that conn is only used inside a context so that locking is ensured

Author: fredthomsen

hyper/http20/stream.py

@@ -72,7 +72,7 @@ def __init__(self,
         # one for data being sent to us.
         self._in_window_manager = window_manager
 
-        # Save off a reference to the state machine.
+        # Save off a reference to the state machine wrapped with lock.
         self._conn = connection
 
         # Save off a data callback.
@@ -94,8 +94,9 @@ def send_headers(self, end_stream=False):
         Sends the complete saved header block on the stream.
         """
         headers = self.get_headers()
-        self._conn.send_headers(self.stream_id, headers, end_stream)
-        self._send_cb(self._conn.data_to_send())
+        with self._conn as conn:
+            conn.send_headers(self.stream_id, headers, end_stream)
+            self._send_cb(conn.data_to_send())
 
         if end_stream:
             self.local_closed = True
@@ -186,10 +187,11 @@ def receive_data(self, event):
         self.data.append(event.data)
 
         if increment and not self.remote_closed:
-            self._conn.increment_flow_control_window(
-                increment, stream_id=self.stream_id
-            )
-            self._send_cb(self._conn.data_to_send())
+            with self._conn as conn:
+                conn.increment_flow_control_window(
+                    increment, stream_id=self.stream_id
+                )
+                self._send_cb(conn.data_to_send())
 
     def receive_end_stream(self, event):
         """
@@ -277,16 +279,17 @@ def close(self, error_code=None):
         """
         # FIXME: I think this is overbroad, but for now it's probably ok.
         if not (self.remote_closed and self.local_closed):
-            try:
-                self._conn.reset_stream(self.stream_id, error_code or 0)
-            except h2.exceptions.ProtocolError:
-                # If for any reason we can't reset the stream, just tolerate
-                # it.
-                pass
-            else:
-                self._send_cb(
-                    self._conn.data_to_send(), tolerate_peer_gone=True
-                )
+            with self._conn as conn:
+                try:
+                    conn.reset_stream(self.stream_id, error_code or 0)
+                except h2.exceptions.ProtocolError:
+                    # If for any reason we can't reset the stream, just tolerate
+                    # it.
+                    pass
+                else:
+                    self._send_cb(
+                        conn.data_to_send(), tolerate_peer_gone=True
+                    )
             self.remote_closed = True
             self.local_closed = True
 
@@ -297,7 +300,8 @@ def _out_flow_control_window(self):
         """
         The size of our outbound flow control window.
         """
-        return self._conn.local_flow_control_window(self.stream_id)
+
+        return self._conn._obj.local_flow_control_window(self.stream_id)
 
     def _send_chunk(self, data, final):
         """
@@ -321,10 +325,11 @@ def _send_chunk(self, data, final):
             end_stream = True
 
         # Send the frame and decrement the flow control window.
-        self._conn.send_data(
-            stream_id=self.stream_id, data=data, end_stream=end_stream
-        )
-        self._send_cb(self._conn.data_to_send())
+        with self._conn as conn:
+            conn.send_data(
+                stream_id=self.stream_id, data=data, end_stream=end_stream
+            )
+            self._send_cb(conn.data_to_send())
 
         if end_stream:
             self.local_closed = True