Cleaner shutdown for completion_refresher

changelog.md

@@ -1,6 +1,11 @@
 Upcoming (TBD)
 ==============
 
+Bugfixes
+---------
+* Avoid an exception when exiting while completions are being refreshed.
+
+
 Documentation
 ---------
 * Shorten/clarify `pyproject.toml` project description.

mycli/client.py

@@ -241,6 +241,10 @@ def _invalidate_prompt_session(self) -> None:
             self.prompt_session.app.invalidate()
 
     def close(self) -> None:
+        try:
+            self.completion_refresher.stop()
+        except Exception:
+            pass
         try:
             self.schema_prefetcher.stop()
         except Exception:

mycli/completion_refresher.py

@@ -18,6 +18,9 @@ class CompletionRefresher:
     def __init__(self, invalidate_app: Callable[[], None] | None = None) -> None:
         self._completer_thread: threading.Thread | None = None
         self._restart_refresh = threading.Event()
+        self._stop_refresh = threading.Event()
+        self._executor_lock = threading.Lock()
+        self._active_executor: SQLExecute | None = None
         self._refresh_visible_until = 0.0
         self._visibility_timer: threading.Timer | None = None
         self._invalidate_app = invalidate_app
@@ -49,6 +52,7 @@ def refresh(
             if self._visibility_timer is not None:
                 self._visibility_timer.cancel()
                 self._visibility_timer = None
+            self._stop_refresh.clear()
             self._refresh_visible_until = monotonic() + MIN_COMPLETION_REFRESH_MESSAGE_SECONDS
             self._completer_thread = threading.Thread(
                 target=self._bg_refresh, args=(executor, callbacks, completer_options), name="completion_refresh"
@@ -63,6 +67,29 @@ def is_refreshing(self) -> bool:
     def _thread_is_alive(self) -> bool:
         return bool(self._completer_thread and self._completer_thread.is_alive())
 
+    def stop(self) -> None:
+        """Stop and wait for an in-flight completion refresh."""
+        self._stop_refresh.set()
+        self._restart_refresh.clear()
+        self._refresh_visible_until = 0.0
+        if self._visibility_timer is not None:
+            self._visibility_timer.cancel()
+            self._visibility_timer = None
+
+        with self._executor_lock:
+            executor = self._active_executor
+        if executor is not None:
+            try:
+                executor.close()
+            except Exception:
+                pass
+
+        thread = self._completer_thread
+        if thread is not None and thread.is_alive() and thread is not threading.current_thread():
+            thread.join()
+        if thread is not None and not thread.is_alive():
+            self._completer_thread = None
+
     def _bg_refresh(
         self,
         sqlexecute: SQLExecute,
@@ -89,14 +116,21 @@ def _bg_refresh(
             self._finish_refreshing()
             return
 
+        with self._executor_lock:
+            self._active_executor = executor
         try:
+            if self._stop_refresh.is_set():
+                return
+
             # If callbacks is a single function then push it into a list.
             if callable(callbacks):
                 callbacks = [callbacks]
 
             while 1:
                 for refresher in self.refreshers.values():
                     refresher(completer, executor)
+                    if self._stop_refresh.is_set():
+                        return
                     if self._restart_refresh.is_set():
                         self._restart_refresh.clear()
                         break
@@ -109,11 +143,23 @@ def _bg_refresh(
                 # break statement.
                 continue
 
-            for callback in callbacks:
-                callback(completer)
+            if not self._stop_refresh.is_set():
+                for callback in callbacks:
+                    callback(completer)
+        except Exception:
+            if not self._stop_refresh.is_set():
+                raise
         finally:
-            executor.close()
-            self._finish_refreshing()
+            with self._executor_lock:
+                if self._active_executor is executor:
+                    self._active_executor = None
+            try:
+                executor.close()
+            except Exception:
+                if not self._stop_refresh.is_set():
+                    raise
+            finally:
+                self._finish_refreshing()
 
     def _finish_refreshing(self) -> None:
         self._invalidate()

test/pytests/test_client.py

@@ -391,23 +391,18 @@ def fake_read_config_file(handle: object, list_values: bool = True) -> dict[str,
     assert cli.mylogin_cnf == {'client': 'config'}
 
 
-def test_close_stops_schema_prefetcher_and_closes_sqlexecute() -> None:
+def test_close_stops_refreshers_before_closing_connection_and_tunnels() -> None:
     cli = MyCli.__new__(MyCli)
-    stopped: list[bool] = []
-    closed: list[bool] = []
-    ssh_tunnel_closed: list[bool] = []
-    boundary_tunnel_closed: list[bool] = []
-    cli.schema_prefetcher = SimpleNamespace(stop=lambda: stopped.append(True))
-    cli.sqlexecute = SimpleNamespace(close=lambda: closed.append(True))  # type: ignore[assignment]
-    cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: ssh_tunnel_closed.append(True))
-    cli.boundary_tunnel = SimpleNamespace(close=lambda: boundary_tunnel_closed.append(True))  # type: ignore[assignment]
+    calls: list[str] = []
+    cli.completion_refresher = SimpleNamespace(stop=lambda: calls.append('completion'))
+    cli.schema_prefetcher = SimpleNamespace(stop=lambda: calls.append('prefetch'))
+    cli.sqlexecute = SimpleNamespace(close=lambda: calls.append('connection'))  # type: ignore[assignment]
+    cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: calls.append('ssh'))
+    cli.boundary_tunnel = SimpleNamespace(close=lambda: calls.append('boundary'))  # type: ignore[assignment]
 
     MyCli.close(cli)
 
-    assert stopped == [True]
-    assert closed == [True]
-    assert ssh_tunnel_closed == [True]
-    assert boundary_tunnel_closed == [True]
+    assert calls == ['completion', 'prefetch', 'connection', 'ssh', 'boundary']
 
 
 def test_close_swallows_cleanup_errors() -> None:
@@ -416,6 +411,7 @@ def test_close_swallows_cleanup_errors() -> None:
     def fail() -> None:
         raise RuntimeError('cleanup failed')
 
+    cli.completion_refresher = SimpleNamespace(stop=fail)
     cli.schema_prefetcher = SimpleNamespace(stop=fail)
     cli.sqlexecute = SimpleNamespace(close=fail)  # type: ignore[assignment]
     cast(Any, cli).ssh_tunnel = SimpleNamespace(close=fail)
@@ -425,6 +421,8 @@ def fail() -> None:
 
 def test_close_swallows_boundary_tunnel_close_error() -> None:
     cli = MyCli.__new__(MyCli)
+    cli.completion_refresher = SimpleNamespace(stop=lambda: None)
+    cli.schema_prefetcher = SimpleNamespace(stop=lambda: None)
     cli.sqlexecute = None
     tunnel_closed: list[bool] = []
     cast(Any, cli).ssh_tunnel = SimpleNamespace(close=lambda: tunnel_closed.append(True))

test/pytests/test_completion_refresher.py

@@ -256,6 +256,108 @@ def test_refresh_cancels_pending_visibility_timer(monkeypatch, refresher) -> Non
     assert refresher._visibility_timer is None
 
 
+def test_stop_interrupts_and_joins_active_refresh(monkeypatch, refresher) -> None:
+    refresh_started = completion_refresher.threading.Event()
+    executor_closed = completion_refresher.threading.Event()
+    callback = Mock()
+    timer = Mock()
+
+    class FakeCompleter:
+        def __init__(self, **options) -> None:
+            pass
+
+    class FakeExecutor:
+        def __init__(self, *args) -> None:
+            pass
+
+        def close(self) -> None:
+            executor_closed.set()
+
+    def blocking_refresh(completer, executor) -> None:
+        refresh_started.set()
+        assert executor_closed.wait(timeout=1)
+        raise completion_refresher.pymysql.err.OperationalError(2013, 'connection closed')
+
+    monkeypatch.setattr(completion_refresher, 'SQLCompleter', FakeCompleter)
+    monkeypatch.setattr(completion_refresher, 'SQLExecute', FakeExecutor)
+    refresher.refreshers = {'blocking': blocking_refresh}
+
+    refresher.refresh(make_sqlexecute(), callback)
+    assert refresh_started.wait(timeout=1)
+    refresher._visibility_timer = timer
+
+    refresher.stop()
+
+    assert executor_closed.is_set()
+    assert refresher._completer_thread is None
+    assert refresher._active_executor is None
+    assert refresher._restart_refresh.is_set() is False
+    assert refresher.is_refreshing() is False
+    timer.cancel.assert_called_once_with()
+    callback.assert_not_called()
+
+
+def test_stop_before_executor_is_ready_prevents_refresh_and_callback(monkeypatch, refresher) -> None:
+    constructor_started = completion_refresher.threading.Event()
+    release_constructor = completion_refresher.threading.Event()
+    stop_finished = completion_refresher.threading.Event()
+    refresh = Mock()
+    callback = Mock()
+
+    class FakeCompleter:
+        def __init__(self, **options) -> None:
+            pass
+
+    class FakeExecutor:
+        def __init__(self, *args) -> None:
+            constructor_started.set()
+            assert release_constructor.wait(timeout=1)
+
+        def close(self) -> None:
+            pass
+
+    monkeypatch.setattr(completion_refresher, 'SQLCompleter', FakeCompleter)
+    monkeypatch.setattr(completion_refresher, 'SQLExecute', FakeExecutor)
+    refresher.refreshers = {'refresh': refresh}
+
+    refresher.refresh(make_sqlexecute(), callback)
+    assert constructor_started.wait(timeout=1)
+    stop_thread = completion_refresher.threading.Thread(target=lambda: (refresher.stop(), stop_finished.set()))
+    stop_thread.start()
+    assert refresher._stop_refresh.wait(timeout=1)
+    assert stop_finished.is_set() is False
+    release_constructor.set()
+    stop_thread.join(timeout=1)
+
+    assert stop_finished.is_set()
+    refresh.assert_not_called()
+    callback.assert_not_called()
+    assert refresher._completer_thread is None
+
+
+def test_stop_tolerates_executor_close_error_without_worker(refresher) -> None:
+    executor = Mock()
+    executor.close.side_effect = RuntimeError('close failed')
+    refresher._active_executor = executor
+
+    refresher.stop()
+
+    executor.close.assert_called_once_with()
+    assert refresher._completer_thread is None
+
+
+def test_stop_does_not_join_current_worker(monkeypatch, refresher) -> None:
+    thread = Mock()
+    thread.is_alive.return_value = True
+    refresher._completer_thread = thread
+    monkeypatch.setattr(completion_refresher.threading, 'current_thread', lambda: thread)
+
+    refresher.stop()
+
+    thread.join.assert_not_called()
+    assert refresher._completer_thread is thread
+
+
 def test_finish_refreshing_schedules_delayed_invalidation_before_deadline(monkeypatch, refresher) -> None:
     now = 10.0
     monkeypatch.setattr(completion_refresher, 'monotonic', lambda: now)
@@ -474,6 +576,72 @@ def __init__(self, *args) -> None:
     callback.assert_not_called()
 
 
+def test_bg_refresh_stops_after_current_refresher(monkeypatch, refresher) -> None:
+    callback = Mock()
+    executor = Mock()
+
+    def stop_refresh(completer, active_executor) -> None:
+        refresher._stop_refresh.set()
+
+    monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock())
+    monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor))
+    refresher.refreshers = {'stop': stop_refresh}
+
+    refresher._bg_refresh(make_sqlexecute(), callback, {})
+
+    callback.assert_not_called()
+    executor.close.assert_called_once_with()
+
+
+def test_bg_refresh_skips_callbacks_when_stopped_after_refresh(monkeypatch, refresher) -> None:
+    callback = Mock()
+    executor = Mock()
+    is_stopped = Mock(side_effect=[False, True])
+    monkeypatch.setattr(refresher._stop_refresh, 'is_set', is_stopped)
+    monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock())
+    monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor))
+    refresher.refreshers = {}
+
+    refresher._bg_refresh(make_sqlexecute(), callback, {})
+
+    callback.assert_not_called()
+    assert is_stopped.call_count == 2
+
+
+def test_bg_refresh_propagates_unexpected_refresher_error(monkeypatch, refresher) -> None:
+    executor = Mock()
+
+    def fail_refresh(completer, active_executor) -> None:
+        refresher._active_executor = None
+        raise RuntimeError('refresh failed')
+
+    monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock())
+    monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor))
+    refresher.refreshers = {'fail': fail_refresh}
+
+    with pytest.raises(RuntimeError, match='refresh failed'):
+        refresher._bg_refresh(make_sqlexecute(), Mock(), {})
+
+    executor.close.assert_called_once_with()
+
+
+@pytest.mark.parametrize('stopping', [False, True])
+def test_bg_refresh_only_suppresses_executor_close_error_when_stopping(monkeypatch, refresher, stopping) -> None:
+    executor = Mock()
+    executor.close.side_effect = RuntimeError('close failed')
+    monkeypatch.setattr(completion_refresher, 'SQLCompleter', Mock())
+    monkeypatch.setattr(completion_refresher, 'SQLExecute', Mock(return_value=executor))
+    refresher.refreshers = {}
+    if stopping:
+        refresher._stop_refresh.set()
+
+    if stopping:
+        refresher._bg_refresh(make_sqlexecute(), Mock(), {})
+    else:
+        with pytest.raises(RuntimeError, match='close failed'):
+            refresher._bg_refresh(make_sqlexecute(), Mock(), {})
+
+
 def test_refresher_decorator_registers_function() -> None:
     refreshers: dict[str, object] = {}