From 71b5d8ef00cbad54557cd5b93bfc901713bca9f9 Mon Sep 17 00:00:00 2001 From: Matteo Piano Date: Fri, 26 Jan 2024 02:39:06 -0800 Subject: [PATCH] fix shutdown behaviour of remote config fetcher --- pidtree_bcc/main.py | 2 ++ pidtree_bcc/yaml_loader.py | 41 +++++++++++++++++++++++++------------- tests/yaml_loader_test.py | 3 +++ 3 files changed, 32 insertions(+), 14 deletions(-) diff --git a/pidtree_bcc/main.py b/pidtree_bcc/main.py index f6a4b2f..afd3896 100644 --- a/pidtree_bcc/main.py +++ b/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 diff --git a/pidtree_bcc/yaml_loader.py b/pidtree_bcc/yaml_loader.py index 7a436fd..e5b2152 100644 --- a/pidtree_bcc/yaml_loader.py +++ b/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( diff --git a/tests/yaml_loader_test.py b/tests/yaml_loader_test.py index 8c03998..8e972b3 100644 --- a/tests/yaml_loader_test.py +++ b/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()