refactor

Author: aIbrahiim

sdks/python/apache_beam/runners/worker/data_plane.py

@@ -474,8 +474,8 @@ def __init__(self, data_buffer_time_limit_ms=0):
     self._start_send_forwarder()
     self._received = collections.defaultdict(
         lambda: ByteLimitedQueue(
-            maxsize=_DEFAULT_RECEIVE_QUEUE_MAX_ELEMENTS,
-            maxbytes=_DEFAULT_SEND_QUEUE_MAX_BYTES)
+            maxsize=_DEFAULT_RECEIVE_QUEUE_MAX_ELEMENTS, maxbytes=
+            _DEFAULT_SEND_QUEUE_MAX_BYTES)
     )  # type: DefaultDict[str, ByteLimitedQueue[DataOrTimers]]
 
     # Keep a cache of completed instructions. Data for completed instructions
@@ -494,27 +494,29 @@ def close(self):
     self._enqueue_to_send(self._WRITES_FINISHED)
     if self._send_forwarder is not None:
       self._send_forwarder.join()
+    if self._exception:
+      raise self._exception
     self._closed = True
 
   def _start_send_forwarder(self):
     # type: () -> None
     forwarder = threading.Thread(
-        target=self._forward_pending_to_send,
-        name='forward_grpc_outputs')
+        target=self._forward_pending_to_send, name='forward_grpc_outputs')
     forwarder.daemon = True
     forwarder.start()
     self._send_forwarder = forwarder
 
   def _enqueue_to_send(self, elem):
     # type: (DataOrTimers) -> None
-    self._pending_send.put(elem, self._get_element_size_bytes(elem))
+    size = self._get_element_size_bytes(elem)
+    self._pending_send.put((elem, size), size)
 
   def _forward_pending_to_send(self):
     # type: () -> None
     try:
       while True:
-        elem = self._pending_send.get()
-        self._to_send.put(elem, self._get_element_size_bytes(elem))
+        elem, size = self._pending_send.get()
+        self._to_send.put(elem, size)
         if elem is self._WRITES_FINISHED:
           return
     except Exception as e: