googleapis · dhermes · Dec 6, 2017 · Dec 6, 2017 · Dec 6, 2017 · Dec 6, 2017
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_helper_threads.py b/pubsub/google/cloud/pubsub_v1/subscriber/_helper_threads.py
@@ -57,9 +57,6 @@ class HelperThreadRegistry(object):
     def __init__(self):
         self._helper_threads = {}
 
-    def __contains__(self, needle):
-        return needle in self._helper_threads
-
     def start(self, name, queue_put, target):
         """Create and start a helper thread.
 

diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/policy/thread.py b/pubsub/google/cloud/pubsub_v1/subscriber/policy/thread.py
@@ -30,7 +30,7 @@
 
 
 _LOGGER = logging.getLogger(__name__)
-_CALLBACK_WORKER_NAME = 'CallbackRequestsWorker'
+_CALLBACK_WORKER_NAME = 'Thread-Consumer-CallbackRequestsWorker'
 
 
 def _callback_completed(future):
@@ -98,6 +98,9 @@ def __init__(self, client, subscription, flow_control=types.FlowControl(),
         self._request_queue = self._get_queue(queue)
         # Also maintain an executor.
         self._executor = self._get_executor(executor)
+        # The threads created in ``.open()``.
+        self._dispatch_thread = None
+        self._leases_thread = None
 
     @staticmethod
     def _get_queue(queue):
@@ -146,8 +149,12 @@ def _get_executor(executor):
     def close(self):
         """Close the existing connection."""
         # Stop consuming messages.
-        self._consumer.helper_threads.stop(_CALLBACK_WORKER_NAME)
+        self._request_queue.put(_helper_threads.STOP)
+        self._dispatch_thread.join()  # Wait until stopped.
+        self._dispatch_thread = None
         self._consumer.stop_consuming()
+        self._leases_thread.join()
+        self._leases_thread = None
         self._executor.shutdown()
 
         # The subscription is closing cleanly; resolve the future if it is not
@@ -156,6 +163,53 @@ def close(self):
             self._future.set_result(None)
         self._future = None
 
+    def _start_dispatch(self):
+        """Start a thread to dispatch requests queued up by callbacks.
+
+        .. note::
+
+            This assumes, but does not check, that ``_dispatch_thread``
+            is :data:`None`.
+
+        Spawns a thread to run :meth:`dispatch_callback` and sets the
+        "dispatch thread" member on the current policy.
+        """
+        _LOGGER.debug('Starting callback requests worker.')
+        dispatch_worker = _helper_threads.QueueCallbackWorker(
+            self._request_queue,
+            self.dispatch_callback,
+        )
+        # Create and start the helper thread.
+        thread = threading.Thread(
+            name=_CALLBACK_WORKER_NAME,
+            target=dispatch_worker,
+        )
+        thread.daemon = True
+        thread.start()
+        _LOGGER.debug('Started helper thread %s', thread.name)
+        self._dispatch_thread = thread
+
+    def _start_lease_worker(self):
+        """Spawn a helper thread that maintains all of leases for this policy.
+
+        .. note::
+
+            This assumes, but does not check, that ``_leases_thread`` is
+            :data:`None`.
+
+        Spawns a thread to run :meth:`maintain_leases` and sets the
+        "leases thread" member on the current policy.
+        """
+        _LOGGER.debug('Starting lease maintenance worker.')
+        thread = threading.Thread(
+            name='Thread-LeaseMaintenance',
+            target=self.maintain_leases,
+        )
+        thread.daemon = True
+        thread.start()
+
+        self._leases_thread = thread
+
     def open(self, callback):
         """Open a streaming pull connection and begin receiving messages.
 
@@ -177,30 +231,11 @@ def open(self, callback):
         self._future = Future(policy=self)
 
         # Start the thread to pass the requests.
-        _LOGGER.debug('Starting callback requests worker.')
         self._callback = callback
-        dispatch_worker = _helper_threads.QueueCallbackWorker(
-            self._request_queue,
-            self.dispatch_callback,
-        )
-        self._consumer.helper_threads.start(
-            _CALLBACK_WORKER_NAME,
-            self._request_queue.put,
-            dispatch_worker,
-        )
-
+        self._start_dispatch()
         # Actually start consuming messages.
         self._consumer.start_consuming()
-
-        # Spawn a helper thread that maintains all of the leases for
-        # this policy.
-        _LOGGER.debug('Starting lease maintenance worker.')
-        self._leaser = threading.Thread(
-            name='Thread-LeaseMaintenance',
-            target=self.maintain_leases,
-        )
-        self._leaser.daemon = True
-        self._leaser.start()
+        self._start_lease_worker()
 
         # Return the future.
         return self._future

diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_base.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_base.py
@@ -180,8 +180,10 @@ def test_load():
     assert policy._load == 0.2
 
     # Returning a number above 100% is fine.
-    policy.lease(ack_id='three', byte_size=1000)
-    assert policy._load == 1.16
+    with mock.patch.object(policy, 'close') as close:
+        policy.lease(ack_id='three', byte_size=1000)
+        assert policy._load == 1.16
+        close.assert_called_once_with()
 
 
 def test_modify_ack_deadline():

diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_thread.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_policy_thread.py
@@ -21,6 +21,7 @@
 from google.auth import credentials
 import mock
 import pytest
+import six
 from six.moves import queue
 
 from google.cloud.pubsub_v1 import subscriber
@@ -49,36 +50,69 @@ def test_init_with_executor():
 
 
 def test_close():
+    dispatch_thread = mock.Mock(spec=threading.Thread)
+    leases_thread = mock.Mock(spec=threading.Thread)
+
     policy = create_policy()
+    policy._dispatch_thread = dispatch_thread
+    policy._leases_thread = leases_thread
     consumer = policy._consumer
     with mock.patch.object(consumer, 'stop_consuming') as stop_consuming:
         policy.close()
         stop_consuming.assert_called_once_with()
-    assert 'callback request worker' not in policy._consumer.helper_threads
+
+    assert policy._dispatch_thread is None
+    dispatch_thread.join.assert_called_once_with()
+    assert policy._leases_thread is None
+    leases_thread.join.assert_called_once_with()
 
 
 def test_close_with_future():
+    dispatch_thread = mock.Mock(spec=threading.Thread)
+    leases_thread = mock.Mock(spec=threading.Thread)
+
     policy = create_policy()
+    policy._dispatch_thread = dispatch_thread
+    policy._leases_thread = leases_thread
     policy._future = Future(policy=policy)
     consumer = policy._consumer
     with mock.patch.object(consumer, 'stop_consuming') as stop_consuming:
         future = policy.future
         policy.close()
         stop_consuming.assert_called_once_with()
+
+    assert policy._dispatch_thread is None
+    dispatch_thread.join.assert_called_once_with()
+    assert policy._leases_thread is None
+    leases_thread.join.assert_called_once_with()
     assert policy.future != future
     assert future.result() is None
 
 
-@mock.patch.object(_helper_threads.HelperThreadRegistry, 'start')
-@mock.patch.object(threading.Thread, 'start')
-def test_open(thread_start, htr_start):
+def test_open():
     policy = create_policy()
-    with mock.patch.object(policy._consumer, 'start_consuming') as consuming:
+    consumer = policy._consumer
+    threads = (
+        mock.Mock(spec=('name', 'start')),
+        mock.Mock(spec=('name', 'start')),
+        mock.Mock(spec=('name', 'start')),
+    )
+    with mock.patch.object(threading, 'Thread', side_effect=threads):
         policy.open(mock.sentinel.CALLBACK)
-        assert policy._callback is mock.sentinel.CALLBACK
-        consuming.assert_called_once_with()
-        htr_start.assert_called()
-        thread_start.assert_called()
+
+    assert policy._callback is mock.sentinel.CALLBACK
+
+    assert policy._dispatch_thread is threads[0]
+    threads[0].start.assert_called_once_with()
+
+    threads_dict = consumer.helper_threads._helper_threads
+    assert len(threads_dict) == 1
+    helper_thread = next(six.itervalues(threads_dict))
+    assert helper_thread.thread is threads[1]
+    threads[1].start.assert_called_once_with()
+
+    assert policy._leases_thread is threads[2]
+    threads[2].start.assert_called_once_with()
 
 
 def test_dispatch_callback_valid_actions():