Add re-open throttling to ResumableBidiRpc

plamut · plamut · commit 2839d0be959d · 2019-05-30T15:38:42.000+02:00
diff --git a/api_core/google/api_core/bidi.py b/api_core/google/api_core/bidi.py
@@ -397,6 +397,7 @@ def __init__(self, start_rpc, should_recover, initial_request=None, metadata=Non
         self._operational_lock = threading.RLock()
         self._finalized = False
         self._finalize_lock = threading.Lock()
+        self._reopen_throttle = _Throttle(time_window=10, entry_cap=5)
 
     def _finalize(self, result):
         with self._finalize_lock:
@@ -440,7 +441,8 @@ def _reopen(self):
             # retryable error.
 
             try:
-                self.open()
+                with self._reopen_throttle:
+                    self.open()
             # If re-opening or re-calling the method fails for any reason,
             # consider it a terminal error and finalize the stream.
             except Exception as exc:
@@ -639,7 +641,7 @@ def start(self):
             thread = threading.Thread(
                 name=_BIDIRECTIONAL_CONSUMER_NAME,
                 target=self._thread_main,
-                args=(ready,)
+                args=(ready,),
             )
             thread.daemon = True
             thread.start()
diff --git a/api_core/tests/unit/test_bidi.py b/api_core/tests/unit/test_bidi.py
@@ -513,6 +513,20 @@ def test_reopen_failure_on_rpc_restart(self):
         assert bidi_rpc.is_active is False
         callback.assert_called_once_with(error2)
 
+    def test_throttles_reopen_requests(self):
+        call = CallStub([])
+        start_rpc = mock.create_autospec(
+            grpc.StreamStreamMultiCallable, instance=True, return_value=call
+        )
+        should_recover = mock.Mock(spec=["__call__"], return_value=True)
+        bidi_rpc = bidi.ResumableBidiRpc(start_rpc, should_recover)
+
+        patcher = mock.patch.object(bidi_rpc._reopen_throttle.__class__, "__enter__")
+        with patcher as mock_enter:
+            bidi_rpc._reopen()
+
+        mock_enter.assert_called_once()
+
     def test_send_not_open(self):
         bidi_rpc = bidi.ResumableBidiRpc(None, lambda _: False)
 
