fix shutdown behaviour of remote config fetcher

pidtree_bcc/main.py

@@ -22,6 +22,7 @@
 from pidtree_bcc.utils import self_restart
 from pidtree_bcc.utils import smart_open
 from pidtree_bcc.utils import StopFlagWrapper
+from pidtree_bcc.yaml_loader import FileIncludeLoader
 
 
 EXIT_CODE = 0
@@ -203,6 +204,7 @@ def main(args: argparse.Namespace):
             out.flush()
     except RestartSignal:
         stop_wrapper.stop()
+        FileIncludeLoader.cleanup()
         raise
     except Exception as e:
         # Terminate everything if something goes wrong

pidtree_bcc/yaml_loader.py

@@ -65,7 +65,7 @@ def include_file(self, loader: yaml.Loader, node: yaml.Node) -> Any:
         """
         name = loader.construct_scalar(node)
         filepath = (
-            self.include_remote(name)
+            self.include_remote(name, self.stop_flag)
             if re.match(r'^https?://', name)
             else (
                 os.path.join(os.path.dirname(loader.name), name)
@@ -82,33 +82,37 @@ def include_file(self, loader: yaml.Loader, node: yaml.Node) -> Any:
             _, value, traceback = sys.exc_info()
             raise yaml.YAMLError(value).with_traceback(traceback)
 
-    def include_remote(self, url: str) -> str:
+    @classmethod
+    def include_remote(cls, url: str, stop_flag: Optional[StopFlagWrapper] = None) -> str:
         """ Load remote configuration data
 
         :param str url: resource url
         :return: local filepath where data is stored
         """
-        if self.remote_fetcher is None or not self.remote_fetcher.is_alive():
-            self.remote_fetcher_fence = Condition()
-            self.remote_fetcher_outdir = tempfile.mkdtemp(prefix='pidtree-bcc-conf')
-            self.remote_fetcher = Thread(
+        if cls.remote_fetcher is None or not cls.remote_fetcher.is_alive():
+            logging.info('Setting up remote configuration fetcher thread')
+            cls.remote_fetcher_fence = Condition()
+            cls.remote_fetcher_outdir = tempfile.mkdtemp(prefix='pidtree-bcc-conf')
+            cls.remote_fetcher = Thread(
                 target=fetch_remote_configurations,
                 args=(
-                    self.REMOTE_FETCH_INTERVAL_SECONDS, self.remote_fetcher_fence,
-                    self.remote_fetch_workload, self.stop_flag,
+                    cls.REMOTE_FETCH_INTERVAL_SECONDS,
+                    cls.remote_fetcher_fence,
+                    cls.remote_fetch_workload,
+                    stop_flag,
                 ),
                 daemon=True,
             )
-            self.remote_fetcher.start()
+            FileIncludeLoader.remote_fetcher.start()
         logging.info(f'Loading remote configuration from {url}')
         ready = Condition()
         url_sha = hashlib.sha256(url.encode()).hexdigest()
-        output_path = os.path.join(self.remote_fetcher_outdir, f'{url_sha}.yaml')
-        self.remote_fetch_workload[url] = (output_path, ready)
-        with self.remote_fetcher_fence:
-            self.remote_fetcher_fence.notify()
+        output_path = os.path.join(cls.remote_fetcher_outdir, f'{url_sha}.yaml')
+        cls.remote_fetch_workload[url] = (output_path, ready)
+        with cls.remote_fetcher_fence:
+            cls.remote_fetcher_fence.notify()
         with ready:
-            if not ready.wait(timeout=self.REMOTE_FETCH_MAX_WAIT_SECONDS):
+            if not ready.wait(timeout=cls.REMOTE_FETCH_MAX_WAIT_SECONDS):
                 raise ValueError(f'Failed to load configuration at {url}')
         return output_path
 
@@ -122,6 +126,15 @@ def get_loader_instance(cls, stop_flag: Optional[StopFlagWrapper] = None) -> Tup
         included_files = []
         return partial(cls, included_files=included_files, stop_flag=stop_flag), included_files
 
+    @classmethod
+    def cleanup(cls):
+        """ Cleans eventual background configuration fetcher setup """
+        if cls.remote_fetcher_fence:
+            with cls.remote_fetcher_fence:
+                cls.remote_fetcher_fence.notify()
+            shutil.rmtree(cls.remote_fetcher_outdir, ignore_errors=True)
+        cls.remote_fetch_workload.clear()
+
 
 @never_crash
 def fetch_remote_configurations(

tests/yaml_loader_test.py

@@ -37,6 +37,7 @@ def test_file_include_remote(mock_request, mock_tempfile, tmp_path):
         with open('tests/fixtures/remote_config.yaml') as f:
             data = yaml.load(f, Loader=loader)
     stop_flag.stop()
+    FileIncludeLoader.cleanup()
     assert data == {
         'foo': [1, {'a': 2, 'b': 3}, 4],
         'bar': {'fizz': 'buzz'},
@@ -47,3 +48,5 @@ def test_file_include_remote(mock_request, mock_tempfile, tmp_path):
     mock_request.urlopen.assert_called_once_with(
         'https://raw.githubusercontent.com/Yelp/pidtree-bcc/master/tests/fixtures/child_config.yaml',
     )
+    assert not FileIncludeLoader.remote_fetch_workload
+    assert not FileIncludeLoader.remote_fetcher.is_alive()