[GitHub] [beam] rohdesamuel commented on a change in pull request #12701: [BEAM-10603] Gracefully shutdown the channel reader in the test_stream_impl

[GitBox <gi...@apache.org>]

rohdesamuel commented on a change in pull request #12701:
URL: https://github.com/apache/beam/pull/12701#discussion_r481335695



##########
File path: sdks/python/apache_beam/runners/direct/test_stream_impl.py
##########
@@ -276,11 +276,20 @@ def _stream_events_from_rpc(endpoint, output_tags, coder, channel, is_alive):
         output_ids=[str(tag) for tag in output_tags])
 
     event_stream = stub.Events(event_request)
-    for e in event_stream:
-      channel.put(_TestStream.test_stream_payload_to_events(e, coder))
-      if not is_alive():
+    try:
+      for e in event_stream:
+        channel.put(_TestStream.test_stream_payload_to_events(e, coder))
+        if not is_alive():
+          return
+    except grpc.RpcError as e:
+      # This happens when the Python interpreter shuts down or whn in a
+      # notebook environment when the kernel is interrupted.
+      if e.code() == grpc.StatusCode.UNAVAILABLE:

Review comment:
       So far I've only seen UNAVAILABLE as a non-error code. But looking at https://developers.google.com/maps-booking/reference/grpc-api/status_codes I think we can also add CANCELLED to not forward the exception.
   
   At L290, we do have a finally clause which will run regardless, so the endofstream will always be written to the channel.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org