Actually **stopping** the Pub / Sub consumer on an exception. (#4498)

Fixes #4463.

Author: dhermes

google/cloud/pubsub_v1/subscriber/_consumer.py

@@ -242,10 +242,10 @@ def _blocking_consume(self):
                 # case, break out of the while loop and exit this thread.
                 _LOGGER.debug('Clean RPC loop exit signalled consumer exit.')
                 break
-            except KeyboardInterrupt:
-                self.stop_consuming()
             except Exception as exc:
-                self.active = self._policy.on_exception(exc)
+                recover = self._policy.on_exception(exc)
+                if not recover:
+                    self.stop_consuming()
 
     def start_consuming(self):
         """Start consuming the stream."""

google/cloud/pubsub_v1/subscriber/_helper_threads.py

@@ -40,6 +40,19 @@
 STOP = uuid.uuid4()
 
 
+def _current_thread():
+    """Get the currently active thread.
+
+    This is provided as a test helper so that it can be mocked easily.
+    Mocking ``threading.current_thread()`` directly may have unintended
+    consequences on code that relies on it.
+
+    Returns:
+        threading.Thread: The current thread.
+    """
+    return threading.current_thread()
+
+
 class HelperThreadRegistry(object):
     def __init__(self):
         self._helper_threads = {}
@@ -84,6 +97,16 @@ def stop(self, name):
         if helper_thread is None:
             return
 
+        if helper_thread.thread is _current_thread():
+            # The current thread cannot ``join()`` itself but it can
+            # still send a signal to stop.
+            _LOGGER.debug('Cannot stop current thread %s', name)
+            helper_thread.queue.put(STOP)
+            # We return and stop short of ``pop()``-ing so that the
+            # thread that invoked the current helper can properly stop
+            # it.
+            return
+
         # Join the thread if it is still alive.
         if helper_thread.thread.is_alive():
             _LOGGER.debug('Stopping helper thread %s', name)

google/cloud/pubsub_v1/subscriber/policy/thread.py

@@ -180,12 +180,12 @@ def on_exception(self, exception):
         """Handle the exception.
 
         If the exception is one of the retryable exceptions, this will signal
-        to the consumer thread that it should remain active.
+        to the consumer thread that it should "recover" from the failure.
 
         This will cause the stream to exit when it returns :data:`False`.
 
         Returns:
-            bool: Indicates if the caller should remain active or shut down.
+            bool: Indicates if the caller should recover or shut down.
             Will be :data:`True` if the ``exception`` is "acceptable", i.e.
             in a list of retryable / idempotent exceptions.
         """

tests/unit/pubsub_v1/subscriber/test_consumer.py

@@ -76,19 +76,6 @@ def test_blocking_consume():
             assert on_res.mock_calls[1][1][1] == mock.sentinel.B
 
 
-def test_blocking_consume_keyboard_interrupt():
-    consumer = create_consumer()
-    Policy = type(consumer._policy)
-
-    # Establish that we get responses until we are sent the exiting event.
-    with mock.patch.object(Policy, 'call_rpc', autospec=True) as call_rpc:
-        call_rpc.return_value = (mock.sentinel.A, mock.sentinel.B)
-        with mock.patch.object(Policy, 'on_response', autospec=True) as on_res:
-            on_res.side_effect = KeyboardInterrupt
-            consumer._blocking_consume()
-            on_res.assert_called_once_with(consumer._policy, mock.sentinel.A)
-
-
 class OnException(object):
 
     def __init__(self, exiting_event, acceptable=None):

tests/unit/pubsub_v1/subscriber/test_helper_threads.py

@@ -36,6 +36,29 @@ def test_stop_noop():
     assert len(registry._helper_threads) == 0
 
 
+@mock.patch.object(
+    _helper_threads, '_current_thread', return_value=mock.sentinel.thread)
+def test_stop_current_thread(_current_thread):
+    registry = _helper_threads.HelperThreadRegistry()
+    queue_ = mock.Mock(spec=('put',))
+
+    name = 'here'
+    registry._helper_threads[name] = _helper_threads._HelperThread(
+        name=name,
+        queue=queue_,
+        thread=_current_thread.return_value,
+    )
+    assert list(registry._helper_threads.keys()) == [name]
+    registry.stop(name)
+    # Make sure it hasn't been removed from the registry ...
+    assert list(registry._helper_threads.keys()) == [name]
+    # ... but it did receive the STOP signal.
+    queue_.put.assert_called_once_with(_helper_threads.STOP)
+
+    # Verify that our mock was only called once.
+    _current_thread.assert_called_once_with()
+
+
 def test_stop_dead_thread():
     registry = _helper_threads.HelperThreadRegistry()
     registry._helper_threads['foo'] = _helper_threads._HelperThread(