BUG Ensure that polling thread shuts down when no longer needed (#45)

* BUG Ensure that polling thread shuts down when no longer needed

When a Python session shuts down, all threads are joined. That means that we need to wait for the polling thread to complete before the interpreter can shut down. It hadn't been noticable when polling was the primary way of getting results, but with the use of the notifications endpoint, it matters. I observed a very short job which used notifications wait 9.5 minutes to complete, because it had to wait for the `sleep` inside our polling thread to finish.

There's no way to shut down a thread in a ``ThreadPoolExecutor`` early, so replace it with our own thread. The ``cancel`` method lets us stop it early, and overloading ``join`` ensures that we never need to wait for a polling cycle to complete when exiting the interpreter. This also solves a problem I'd observed in testing, that the test session wouldn't exit until the poller shut down.

* MAINT Remove overly-paranoid exception catching

Stop wrapping `_check_result` in a try/except inside the thread. It's only necessary if we think there could be coding bugs inside `_check_result`. If there are, we should fix them rather than swallowing them.

* TST Conditionally select polling interval in tests

When we record tests with ``vcr``, we need to have normal polling intervals, so as not to break the Platform API. But when running tests later, we don't want to wait through the entire sequence of polling at the normal intervals. We had been solving this by patching some internals of `PollableResult`. That solution breaks when we swap in the new threading poller. Replace with conditional poll intervals. We need to remember to set `polling_interval` for all new tests, but in return we don't need to patch private methods of `PollableResult` (and it works with the new threading).

Author: Stephen Hoover

CHANGELOG.md

@@ -6,6 +6,9 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### API Changes
 - Deprecate ``api_key`` input to higher-level functions and classes in favor of an ``APIClient`` input. The ``api_key`` will be removed in v2.0.0. (#46)
 
+### Fixed
+- Improved threading implementation in ``PollableResult`` so that it no longer blocks interpreter shutdown.
+
 ### Added
 - Decorator function for deprecating parameters (#46)
 

civis/polling.py

@@ -1,5 +1,5 @@
-from concurrent import futures
 import time
+import threading
 
 from civis.base import CivisJobFailure, CivisAsyncResultBase, FAILED, DONE
 from civis.response import Response
@@ -8,6 +8,37 @@
 _DEFAULT_POLLING_INTERVAL = 15
 
 
+class _ResultPollingThread(threading.Thread):
+    """Poll a function until it returns a Response with a DONE state
+    """
+    # Inspired by `threading.Timer`
+
+    def __init__(self, poller, poller_args, polling_interval):
+        super().__init__()
+        self.polling_interval = polling_interval
+        self.poller = poller
+        self.poller_args = poller_args
+        self.finished = threading.Event()
+
+    def cancel(self):
+        """Stop the poller if it hasn't finished yet.
+        """
+        self.finished.set()
+
+    def join(self, timeout=None):
+        """Shut down the polling when the thread is terminated.
+        """
+        self.cancel()
+        super().join(timeout=timeout)
+
+    def run(self):
+        """Poll until done.
+        """
+        while not self.finished.wait(self.polling_interval):
+            if self.poller(*self.poller_args).state in DONE:
+                self.finished.set()
+
+
 class PollableResult(CivisAsyncResultBase):
     """A class for tracking pollable results.
 
@@ -55,7 +86,7 @@ class PollableResult(CivisAsyncResultBase):
     # Implementation notes: The `PollableResult` depends on some private
     # features of the `concurrent.futures.Future` class, so it's possible
     # that future versions of Python could break something here.
-    # (It works under at least 3.4 and 3.5.)
+    # (It works under at least 3.4, 3.5, and 3.6)
     # We use the following `Future` implementation details
     # - The `Future` checks its state against predefined strings. We use
     #   `STATE_TRANS` to translate from the Civis platform states to `Future`
@@ -76,6 +107,8 @@ def __init__(self, poller, poller_args,
                          api_key=api_key,
                          client=client,
                          poll_on_creation=poll_on_creation)
+        if self.polling_interval <= 0:
+            raise ValueError("The polling interval must be positive.")
 
         # Polling arguments. Never poll more often than the requested interval.
         if poll_on_creation:
@@ -84,37 +117,16 @@ def __init__(self, poller, poller_args,
             self._last_polled = time.time()
         self._last_result = None
 
-        self._self_polling_executor = None
-
-    def _wait_for_completion(self):
-        """Poll the job every `polling_interval` seconds. Blocks until the
-        job completes.
-        """
-        try:
-            while self._civis_state not in DONE:
-                time.sleep(self.polling_interval)
-        except Exception as e:
-            # Exceptions are caught in `_check_result`, so
-            # we should never get here. If there were to be a
-            # bug in `_check_result`, however, we would get stuck
-            # in an infinite loop without setting the `_result`.
-            with self._condition:
-                self._set_api_exception(exc=e)
-
-    def _poll_wait_elapsed(self, now):
-        # thie exists because it's easier to monkeypatch in testing
-        return (now - self._last_polled) >= self.polling_interval
+        self._polling_thread = _ResultPollingThread(self._check_result, (),
+                                                    polling_interval)
 
     def _check_result(self):
         """Return the job result from Civis. Once the job completes, store the
         result and never poll again."""
-
-        # If we haven't started the polling thread, do it now.
-        if self._self_polling_executor is None and self._result is None:
-            # Start a single thread continuously polling. It will stop once the
-            # job completes.
-            self._self_polling_executor = futures.ThreadPoolExecutor(1)
-            self._self_polling_executor.submit(self._wait_for_completion)
+        # Start a single thread continuously polling.
+        # It will stop once the job completes.
+        if not self._polling_thread.is_alive() and self._result is None:
+            self._polling_thread.start()
 
         with self._condition:
             if self._result is not None:
@@ -125,7 +137,8 @@ def _check_result(self):
             # Check to see if the job has finished, but don't poll more
             # frequently than the requested polling frequency.
             now = time.time()
-            if not self._last_polled or self._poll_wait_elapsed(now):
+            if (not self._last_polled or
+                    (now - self._last_polled) >= self.polling_interval):
                 # Poll for a new result
                 self._last_polled = now
                 try:
@@ -165,5 +178,7 @@ def _set_api_exception(self, exc, result=None):
             self.cleanup()
 
     def cleanup(self):
-        # This gets called after the result is set
-        pass
+        # This gets called after the result is set.
+        # Ensure that the polling thread shuts down when it's no longer needed.
+        if self._polling_thread.is_alive():
+            self._polling_thread.cancel()

civis/tests/test_io.py

@@ -18,19 +18,20 @@
 from civis.resources._resources import get_swagger_spec, generate_classes
 from civis.tests.testcase import (CivisVCRTestCase,
                                   cassette_dir,
-                                  conditionally_patch)
+                                  POLL_INTERVAL)
 
 swagger_import_str = 'civis.resources._resources.get_swagger_spec'
 THIS_DIR = os.path.dirname(os.path.realpath(__file__))
 with open(os.path.join(THIS_DIR, "civis_api_spec.json")) as f:
     civis_api_spec = json.load(f, object_pairs_hook=OrderedDict)
 
 
-@conditionally_patch('civis.polling.time.sleep', return_value=None)
-@conditionally_patch('civis.polling.PollableResult._poll_wait_elapsed',
-                     return_value=True)
 @patch(swagger_import_str, return_value=civis_api_spec)
 class ImportTests(CivisVCRTestCase):
+    # Note that all functions tested here should use a
+    # `polling_interval=POLL_INTERVAL` input. This lets us use
+    # sensible polling intervals when recording, but speed through
+    # the calls in the VCR cassette when testing later.
 
     @classmethod
     def setUpClass(cls):
@@ -43,9 +44,6 @@ def tearDownClass(cls):
         generate_classes.cache_clear()
 
     @classmethod
-    @conditionally_patch('civis.polling.time.sleep', return_value=None)
-    @conditionally_patch('civis.polling.PollableResult._poll_wait_elapsed',
-                         return_value=True)
     @patch(swagger_import_str, return_value=civis_api_spec)
     def setup_class(cls, *mocks):
         setup_vcr = vcr.VCR(filter_headers=['Authorization'])
@@ -71,14 +69,16 @@ def setup_class(cls, *mocks):
                 INSERT INTO scratch.api_client_test_fixture
                 VALUES (1,2,3);
             """
-            res = civis.io.query_civis(sql, 'redshift-general')
+            res = civis.io.query_civis(sql, 'redshift-general',
+                                       polling_interval=POLL_INTERVAL)
             res.result()  # block
 
             # create an export to check get_url. also tests export_csv
             with tempfile.NamedTemporaryFile() as tmp:
                 sql = "SELECT * FROM scratch.api_client_test_fixture"
                 database = 'redshift-general'
-                result = civis.io.civis_to_csv(tmp.name, sql, database)
+                result = civis.io.civis_to_csv(tmp.name, sql, database,
+                                               polling_interval=POLL_INTERVAL)
                 result = result.result()
                 assert result.state == 'succeeded'
 
@@ -106,7 +106,8 @@ def test_csv_to_civis(self, *mocks):
             table = "scratch.api_client_test_fixture"
             database = 'redshift-general'
             result = civis.io.csv_to_civis(tmp.name, database, table,
-                                           existing_table_rows='truncate')
+                                           existing_table_rows='truncate',
+                                           polling_interval=POLL_INTERVAL)
             result = result.result()  # block until done
 
         assert isinstance(result.id, int)
@@ -117,22 +118,25 @@ def test_csv_to_civis(self, *mocks):
     def test_read_civis_pandas(self, *mocks):
         expected = pd.DataFrame([[1, 2, 3]], columns=['a', 'b', 'c'])
         df = civis.io.read_civis('scratch.api_client_test_fixture',
-                                 'redshift-general', use_pandas=True)
+                                 'redshift-general', use_pandas=True,
+                                 polling_interval=POLL_INTERVAL)
         assert df.equals(expected)
 
     @patch(swagger_import_str, return_value=civis_api_spec)
     def test_read_civis_no_pandas(self, *mocks):
         expected = [['a', 'b', 'c'], ['1', '2', '3']]
         data = civis.io.read_civis('scratch.api_client_test_fixture',
-                                   'redshift-general', use_pandas=False)
+                                   'redshift-general', use_pandas=False,
+                                   polling_interval=POLL_INTERVAL)
         assert data == expected
 
     @patch(swagger_import_str, return_value=civis_api_spec)
     def test_read_civis_sql(self, *mocks):
         sql = "SELECT * FROM scratch.api_client_test_fixture"
         expected = [['a', 'b', 'c'], ['1', '2', '3']]
         data = civis.io.read_civis_sql(sql, 'redshift-general',
-                                       use_pandas=False)
+                                       use_pandas=False,
+                                       polling_interval=POLL_INTERVAL)
         assert data == expected
 
     @pytest.mark.skipif(not has_pandas, reason="pandas not installed")
@@ -141,15 +145,16 @@ def test_dataframe_to_civis(self, *mocks):
         df = pd.DataFrame([['1', '2', '3']], columns=['a', 'b', 'c'])
         result = civis.io.dataframe_to_civis(df, 'redshift-general',
                                              'scratch.api_client_test_fixture',
-                                             existing_table_rows='truncate')
+                                             existing_table_rows='truncate',
+                                             polling_interval=POLL_INTERVAL)
         result = result.result()
         assert result.state == 'succeeded'
 
     @patch(swagger_import_str, return_value=civis_api_spec)
     def test_civis_to_multifile_csv(self, *mocks):
         sql = "SELECT * FROM scratch.api_client_test_fixture"
-        result = civis.io.civis_to_multifile_csv(sql,
-                                                 database='redshift-general')
+        result = civis.io.civis_to_multifile_csv(
+            sql, database='redshift-general', polling_interval=POLL_INTERVAL)
         assert set(result.keys()) == {'entries', 'query', 'header'}
         assert result['query'] == sql
         assert result['header'] == ['a', 'b', 'c']
@@ -164,13 +169,15 @@ def test_civis_to_multifile_csv(self, *mocks):
     def test_transfer_table(self, *mocks):
         result = civis.io.transfer_table('redshift-general', 'redshift-test',
                                          'scratch.api_client_test_fixture',
-                                         'scratch.api_client_test_fixture')
+                                         'scratch.api_client_test_fixture',
+                                         polling_interval=POLL_INTERVAL)
         result = result.result()
         assert result.state == 'succeeded'
 
         # check for side effect
         sql = 'select * from scratch.api_client_test_fixture'
-        result = civis.io.query_civis(sql, 'redshift-test').result()
+        result = civis.io.query_civis(sql, 'redshift-test',
+                                      polling_interval=POLL_INTERVAL).result()
         assert result.state == 'succeeded'
 
     def test_get_sql_select(self, *mocks):

civis/tests/test_polling.py

@@ -15,12 +15,8 @@ def __init__(self, state):
         self.state = state
 
 
-def func():
-    pass
-
-
 def create_pollable_result(state, exception=None, result=None):
-    f = PollableResult(State, (state, ), polling_interval=0)
+    f = PollableResult(State, (state, ), polling_interval=0.001)
     f._exception = exception
     f._result = result
     return f
@@ -29,12 +25,9 @@ def create_pollable_result(state, exception=None, result=None):
 CANCELLED_RESULT = create_pollable_result(state='cancelled')
 FINISHED_RESULT = create_pollable_result(state='success')
 QUEUED_RESULT = create_pollable_result(state='queued')
-# avoid the polling thread hanging
-QUEUED_RESULT._wait_for_completion = func
 
 
 class TestPolling(unittest.TestCase):
-
     def test_as_completed(self):
         my_futures = [QUEUED_RESULT, CANCELLED_RESULT, FINISHED_RESULT]
         fs = futures.as_completed(my_futures)
@@ -64,51 +57,33 @@ def test_error_setting(self):
         assert isinstance(pollable.exception(), ZeroDivisionError)
 
     def test_timeout(self):
-        # Note: Something about the test framework seems to prevent the
-        # Pollable result from being destroyed while the polling
-        # thread is running. The test will hang if the PollableResult
-        # never completes. I haven't seen the same problem in
-        # the interpreter.
         pollable = PollableResult(
-            mock.Mock(side_effect=[Response({"state": "running"}),
-                                   ValueError()]), (),
+            mock.Mock(return_value=Response({"state": "running"})),
+            poller_args=(),
             polling_interval=0.1)
         pytest.raises(futures.TimeoutError, pollable.result, timeout=0.05)
 
-    def test_no_hanging(self):
-        # Make sure that an error in the `_check_result` doesn't
-        # cause an infinite loop.
-        class PollableResultTester(PollableResult):
-            def __init__(self, *args, **kwargs):
-                self._poll_ct = 0
-                super().__init__(*args, **kwargs)
-
-            def _check_result(self):
-                if self._poll_ct is not None:
-                    self._poll_ct += 1
-                    if self._poll_ct > 10:
-                        self._poll_ct = None  # Disable the counter.
-                        # Make the _wait_for_completion loop fail.
-                        raise ZeroDivisionError()
-                return super()._check_result()
-
-        # The following should raise a CivisJobFailure before a TimeoutError.
-        pollable = PollableResultTester(
-            lambda: Response({"state": "running"}), (),
-            polling_interval=0.1)
-        pytest.raises(ZeroDivisionError, pollable.result, timeout=5)
-
     def test_poll_on_creation(self):
-        poller = mock.Mock(side_effect=Response({"state": "running"}))
+        poller = mock.Mock(return_value=Response({"state": "running"}))
         pollable = PollableResult(poller,
                                   (),
                                   polling_interval=0.01,
                                   poll_on_creation=False)
-        repr(pollable)
+        pollable.done()  # Check status once to start the polling thread
         assert poller.call_count == 0
-        time.sleep(0.02)
+        time.sleep(0.015)
         assert poller.call_count == 1
 
 
+def test_repeated_polling():
+    # Verify that we poll the expected number of times.
+    poller = mock.Mock(return_value=Response({"state": "running"}))
+    pollable = PollableResult(poller, (), polling_interval=0.1)
+    pollable.done()  # Check status once to start the polling thread
+    assert poller.call_count == 1, "Poll once on the first status check"
+    time.sleep(0.25)
+    assert poller.call_count == 3, "After waiting 2.5x the polling interval"
+
+
 if __name__ == '__main__':
     unittest.main()

civis/tests/testcase.py

@@ -1,23 +1,18 @@
-from functools import wraps
 import os
 
 from vcr import persist
 from vcr.serialize import CASSETTE_FORMAT_VERSION
 from vcr.serializers import compat
 from vcr_unittest import VCRTestCase
 
-
-def conditionally_patch(target, *args, **kwargs):
-    from unittest.mock import patch
-    if os.getenv('GENERATE_TESTS'):
-        def pass_func(func):
-            wraps(func)
-            return func
-        return pass_func
-    else:
-        def decorated_func(func):
-            return patch(target, *args, **kwargs)(func)
-        return decorated_func
+# The "GENERATE_TEST" environment variable indicates that
+# we're recording new cassettes.
+if os.environ.get('GENERATE_TESTS'):
+    # Use default polling intervals if generating new tests.
+    POLL_INTERVAL = None
+else:
+    # Speed through calls in pre-recorded VCR cassettes.
+    POLL_INTERVAL = 0.00001
 
 
 def cassette_dir():