diff --git a/hermes_cli/plugins_loader.py b/hermes_cli/plugins_loader.py index e916a8633b..4ca05947f3 100644 --- a/hermes_cli/plugins_loader.py +++ b/hermes_cli/plugins_loader.py @@ -116,7 +116,10 @@ def run_with_load_deadline(plugin_key: str, ctx: "PluginContext", fn: Callable[[ scope = _IN_PLUGIN_LOAD.get() if scope is not None: scope.append(ctx) - return fn() + try: + return fn() + finally: + scope.remove(ctx) timeout = _resolve_plugin_load_timeout() if timeout <= 0: return fn() diff --git a/tests/hermes_cli/test_plugins_loader.py b/tests/hermes_cli/test_plugins_loader.py index dede98d3b2..82ca101db9 100644 --- a/tests/hermes_cli/test_plugins_loader.py +++ b/tests/hermes_cli/test_plugins_loader.py @@ -1,24 +1,36 @@ -from hermes_cli import plugins_loader import threading +import pytest + +from hermes_cli import plugins_loader + def test_nested_plugin_load_runs_inline_on_deadline_worker(monkeypatch): - monkeypatch.setattr(plugins_loader, "_resolve_plugin_load_timeout", lambda: 1.0) + monkeypatch.setattr(plugins_loader, "_resolve_plugin_load_timeout", lambda: 0.3) + abandoned = [] class Context: + def __init__(self, name): + self.name = name + def _abandon_load(self): - raise AssertionError("unexpected timeout") + abandoned.append(self.name) observed_threads = [] - - def nested_load(): - observed_threads.append(threading.current_thread()) + release = threading.Event() def outer_load(): observed_threads.append(threading.current_thread()) - plugins_loader.run_with_load_deadline("nested", Context(), nested_load) + plugins_loader.run_with_load_deadline( + "done", Context("done"), lambda: observed_threads.append(threading.current_thread()), + ) + plugins_loader.run_with_load_deadline("hung", Context("hung"), release.wait) - plugins_loader.run_with_load_deadline("outer", Context(), outer_load) + with pytest.raises(plugins_loader.PluginLoadTimeout): + plugins_loader.run_with_load_deadline("outer", Context("outer"), outer_load) + release.set() assert len(observed_threads) == 2 assert observed_threads[0] is observed_threads[1] + # The outer timeout abandons only contexts still loading, not a nested load that already finished. + assert abandoned == ["outer", "hung"]