Fix pubsub Streaming Pull shutdown on RetryError (googleapis#7863)

plamut · web-flow · commit e00f6b36253b · 2019-05-07T19:18:00.000+02:00
If a RetryError occurs, it is time to stop waiting for the underlying
gRPC channel to recover from a transient failure, and a clean shutdown
needs to be triggered.

This commit assures that this indeed happens (it used to happen on
terminal channel errors only).
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
@@ -262,7 +262,11 @@ def _send_unary_request(self, request):
         _LOGGER.debug("Sent request(s) over unary RPC.")
 
     def send(self, request):
-        """Queue a request to be sent to the RPC."""
+        """Queue a request to be sent to the RPC.
+
+        If a RetryError occurs, the manager shutdown is triggered, and the
+        error is re-raised.
+        """
         if self._UNARY_REQUESTS:
             try:
                 self._send_unary_request(request)
@@ -272,6 +276,17 @@ def send(self, request):
                     "non-fatal as stream requests are best-effort.",
                     exc_info=True,
                 )
+            except exceptions.RetryError as exc:
+                _LOGGER.debug(
+                    "RetryError while sending unary RPC. Waiting on a transient "
+                    "error resolution for too long, will now trigger shutdown.",
+                    exc_info=False,
+                )
+                # The underlying channel has been suffering from a retryable error
+                # for too long, time to give up and shut the streaming pull down.
+                self._on_rpc_done(exc)
+                raise
+
         else:
             self._rpc.send(request)
 
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
@@ -245,7 +245,7 @@ def test_send_unary_empty():
     manager._client.modify_ack_deadline.assert_not_called()
 
 
-def test_send_unary_error(caplog):
+def test_send_unary_api_call_error(caplog):
     caplog.set_level(logging.DEBUG)
 
     manager = make_manager()
@@ -259,6 +259,24 @@ def test_send_unary_error(caplog):
     assert "The front fell off" in caplog.text
 
 
+def test_send_unary_retry_error(caplog):
+    caplog.set_level(logging.DEBUG)
+
+    manager, _, _, _, _, _ = make_running_manager()
+    manager._UNARY_REQUESTS = True
+
+    error = exceptions.RetryError(
+        "Too long a transient error", cause=Exception("Out of time!")
+    )
+    manager._client.acknowledge.side_effect = error
+
+    with pytest.raises(exceptions.RetryError):
+        manager.send(types.StreamingPullRequest(ack_ids=["ack_id1", "ack_id2"]))
+
+    assert "RetryError while sending unary RPC" in caplog.text
+    assert "signaled streaming pull manager shutdown" in caplog.text
+
+
 def test_send_streaming():
     manager = make_manager()
     manager._UNARY_REQUESTS = False
