fix(heatmap): prevent threadleak on shutdown

bec_widgets/widgets/plots/heatmap/heatmap.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import atexit
 import json
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Literal
@@ -158,6 +159,9 @@ class _StepInterpolationWorker(QObject):
             request(_InterpolationRequest): The interpolation request payload.
             data_version(int): The data version for the request.
         """
+        thread = QThread.currentThread()
+        if thread.isInterruptionRequested():
+            return
         self._active_request = request
         self._processing = True
         try:
@@ -170,11 +174,29 @@ class _StepInterpolationWorker(QObject):
             )
         except Exception as exc:  # pragma: no cover - defensive
             logger.warning(f"Step-scan interpolation failed with: {exc}")
-            self.failed.emit(str(exc), data_version, request.scan_id)
+            if not thread.isInterruptionRequested():
+                self.failed.emit(str(exc), data_version, request.scan_id)
             self._processing = False
             return
         self._processing = False
-        self.finished.emit(image, transform, data_version, request.scan_id)
+        if not thread.isInterruptionRequested():
+            self.finished.emit(image, transform, data_version, request.scan_id)
+
+
+# Neither QObject may be owned by the widget: parent destruction can happen while
+# SciPy is still running. Retain both Python wrappers until Qt deletes the stopped thread.
+_interpolation_threads: dict[QThread, _StepInterpolationWorker] = {}
+
+
+@atexit.register
+def _wait_for_interpolation_threads() -> None:
+    """Drain surviving workers at interpreter exit, after the GUI event loop has stopped."""
+    threads = list(_interpolation_threads)
+    for thread in threads:
+        thread.requestInterruption()
+        thread.quit()
+    for thread in threads:
+        thread.wait()
 
 
 class Heatmap(ImageBase):
@@ -269,6 +291,7 @@ class Heatmap(ImageBase):
         self._latest_interpolation_version = -1
         self._interpolation_thread: QThread | None = None
         self._interpolation_worker: _StepInterpolationWorker | None = None
+        self._interpolation_shutdown = False
         self._pending_interpolation_request: _InterpolationRequest | None = None
         self.heatmap_dialog = None
         self.scan_history_dialog = None
@@ -813,6 +836,8 @@ class Heatmap(ImageBase):
             z_data(list[float]): Z values at each point
             msg(messages.ScanStatusMessage): Scan status message containing scan metadata
         """
+        if self._interpolation_shutdown:
+            return
         request = _InterpolationRequest(
             x_data=list(x_data),
             y_data=list(y_data),
@@ -830,10 +855,21 @@ class Heatmap(ImageBase):
         self._start_step_scan_interpolation(request)
 
     def _ensure_interpolation_thread(self):
+        if self._interpolation_shutdown:
+            return
         if self._interpolation_thread is None:
             self._interpolation_thread = QThread()
             self._interpolation_worker = _StepInterpolationWorker()
-            self._interpolation_worker.moveToThread(self._interpolation_thread)
+            thread = self._interpolation_thread
+            worker = self._interpolation_worker
+            worker.moveToThread(thread)
+            _interpolation_threads[thread] = worker
+            thread.finished.connect(worker.deleteLater)
+            thread.finished.connect(thread.deleteLater)
+            thread.destroyed.connect(lambda: _interpolation_threads.pop(thread, None))
+            # These receivers outlive the widget and never access its deleted Qt wrapper.
+            self.destroyed.connect(thread.requestInterruption)
+            self.destroyed.connect(thread.quit)
             self.interpolation_requested.connect(
                 self._interpolation_worker.process, Qt.ConnectionType.QueuedConnection
             )
@@ -847,6 +883,8 @@ class Heatmap(ImageBase):
             self._interpolation_thread.start()
 
     def _start_step_scan_interpolation(self, request: _InterpolationRequest):
+        if self._interpolation_shutdown:
+            return
         # data_version = len(z_data) at the time of the request; keep the latest to gate results.
         self._ensure_interpolation_thread()
         if self._interpolation_thread is not None and not self._interpolation_thread.isRunning():
@@ -854,9 +892,12 @@ class Heatmap(ImageBase):
         self._latest_interpolation_version = request.data_version
         self.interpolation_requested.emit(request, request.data_version)
 
+    @SafeSlot(object, object, int, str)
     def _on_interpolation_finished(
         self, img: np.ndarray, transform: QTransform, data_version: int, scan_id: str
     ):
+        if self._interpolation_shutdown:
+            return
         # Only accept results that match the latest dispatched version for the active scan.
         if data_version == self._latest_interpolation_version and scan_id == self.scan_id:
             self._apply_image_update(img, transform)
@@ -864,33 +905,32 @@ class Heatmap(ImageBase):
             logger.info("Discarding outdated interpolation result.")
         self._maybe_start_pending_interpolation()
 
+    @SafeSlot(str, int, str)
     def _on_interpolation_failed(self, error: str, data_version: int, scan_id: str):
+        if self._interpolation_shutdown:
+            return
         logger.warning(f"Interpolation failed for scan {scan_id} (version {data_version}): {error}")
         self._maybe_start_pending_interpolation()
 
     def _finish_interpolation_thread(self):
-        self._pending_interpolation_request = None
+        """Request shutdown without blocking the GUI or deleting a running thread."""
+        self._interpolation_shutdown = True
+        self._invalidate_interpolation_generation()
         if self._interpolation_worker is not None:
             try:
                 self.interpolation_requested.disconnect(self._interpolation_worker.process)
             except (TypeError, RuntimeError) as ext:
                 logger.warning(f"Processing thread already disconnected: {ext}")
                 pass
-            self._interpolation_worker.deleteLater()
             self._interpolation_worker = None
         if self._interpolation_thread is not None:
-            if self._interpolation_thread.isRunning():
-                self._interpolation_thread.quit()
-                if not self._interpolation_thread.wait(3000):  # 3s timeout
-                    logger.error(
-                        f"Interpolation thread of widget {self.gui_id} did not stop within timeout 3s; leaving it dangling."
-                    )
-            self._interpolation_thread.deleteLater()
+            self._interpolation_thread.requestInterruption()
+            self._interpolation_thread.quit()
             self._interpolation_thread = None
-        logger.info(f"Interpolation thread finished of widget {self.gui_id}")
+        logger.info(f"Interpolation shutdown requested for widget {self.gui_id}")
 
     def _maybe_start_pending_interpolation(self):
-        if self._pending_interpolation_request is None:
+        if self._interpolation_shutdown or self._pending_interpolation_request is None:
             return
         if self._pending_interpolation_request.scan_id != self.scan_id:
             self._pending_interpolation_request = None

tests/unit_tests/test_heatmap_widget.py

@@ -1,11 +1,16 @@
+import gc
+from threading import Event
+from time import monotonic
 from unittest import mock
 
 import numpy as np
 import pytest
+import shiboken6
 from bec_lib import messages
 from bec_lib.scan_history import ScanHistory
-from qtpy.QtCore import QPointF
+from qtpy.QtCore import QEvent, QPointF
 from qtpy.QtGui import QTransform
+from qtpy.QtWidgets import QApplication, QWidget
 
 from bec_widgets.widgets.plots.heatmap.heatmap import (
     Heatmap,
@@ -20,16 +25,133 @@ from bec_widgets.widgets.plots.plot_info_label import TextOnlyLegendSample
 from tests.unit_tests.client_mocks import mocked_client
 
 from .client_mocks import create_dummy_scan_item
+from .conftest import create_widget
 
 
 @pytest.fixture
 def heatmap_widget(qtbot, mocked_client):
-    widget = Heatmap(client=mocked_client)
-    qtbot.addWidget(widget)
-    qtbot.waitExposed(widget)
+    widget = create_widget(qtbot, Heatmap, client=mocked_client)
     yield widget
 
 
+@pytest.fixture
+def blocked_interpolation(monkeypatch):
+    """Keep native interpolation active until the test explicitly releases it."""
+    entered = Event()
+    release = Event()
+
+    def compute(**kwargs):
+        entered.set()
+        if not release.wait(10):
+            raise TimeoutError("Test did not release interpolation")
+        return np.ones((2, 2)), QTransform()
+
+    monkeypatch.setattr(Heatmap, "compute_step_scan_image", compute)
+    yield entered, release
+    release.set()
+
+
+def _start_blocked_interpolation(widget, qtbot, entered):
+    widget.scan_id = "shutdown-scan"
+    request = _InterpolationRequest(
+        x_data=[0, 1, 0, 1],
+        y_data=[0, 0, 1, 1],
+        z_data=[1, 2, 3, 4],
+        data_version=4,
+        scan_id=widget.scan_id,
+        interpolation="linear",
+        oversampling_factor=1.0,
+    )
+    widget._start_step_scan_interpolation(request)
+    qtbot.waitUntil(entered.is_set)
+    return request
+
+
+@pytest.mark.parametrize("shutdown", ["close", "deleteLater"])
+def test_shutdown_preserves_running_interpolation(
+    mocked_client, qtbot, monkeypatch, blocked_interpolation, shutdown
+):
+    # These destruction tests delete their own widgets, so qtbot must not close them again.
+    heatmap_widget = Heatmap(client=mocked_client)
+    entered, release = blocked_interpolation
+    request = _start_blocked_interpolation(heatmap_widget, qtbot, entered)
+    heatmap_widget._pending_interpolation_request = request
+    thread = heatmap_widget._interpolation_thread
+    worker = heatmap_widget._interpolation_worker
+    unsafe_deletions = []
+    delete_thread = thread.deleteLater
+
+    def guarded_delete():
+        # Record the regression without letting Qt abort the entire test process.
+        if thread.isRunning():
+            unsafe_deletions.append(True)
+        else:
+            delete_thread()
+
+    monkeypatch.setattr(thread, "deleteLater", guarded_delete)
+    try:
+        start = monotonic()
+        getattr(heatmap_widget, shutdown)()
+        elapsed = monotonic() - start
+        assert not unsafe_deletions, "Shutdown scheduled deletion of a running QThread"
+        assert elapsed < 1, "Shutdown blocked the GUI waiting for interpolation"
+        QApplication.sendPostedEvents(None, QEvent.Type.DeferredDelete)
+        assert shiboken6.isValid(thread) and thread.isRunning()
+        assert shiboken6.isValid(worker)
+        assert heatmap_widget._pending_interpolation_request is None
+
+        release.set()
+        qtbot.waitUntil(lambda: not shiboken6.isValid(worker))
+        qtbot.waitUntil(lambda: not shiboken6.isValid(thread))
+    finally:
+        # Also leave the unpatched implementation safe after a failed assertion.
+        heatmap_widget._invalidate_interpolation_generation()
+        release.set()
+        if shiboken6.isValid(thread):
+            thread.quit()
+            assert thread.wait(2000)
+            delete_thread()
+        if shiboken6.isValid(heatmap_widget):
+            heatmap_widget.deleteLater()
+
+
+@pytest.mark.parametrize("busy", [False, True], ids=["idle", "interpolating"])
+def test_parent_destruction_stops_interpolation(qtbot, mocked_client, blocked_interpolation, busy):
+    parent = QWidget()
+    widget = Heatmap(parent=parent, client=mocked_client)
+    entered, release = blocked_interpolation
+    if busy:
+        _start_blocked_interpolation(widget, qtbot, entered)
+    else:
+        widget._ensure_interpolation_thread()
+    thread = widget._interpolation_thread
+    worker = widget._interpolation_worker
+    try:
+        parent.deleteLater()
+        QApplication.sendPostedEvents(None, QEvent.Type.DeferredDelete)
+        assert not shiboken6.isValid(widget)
+        # Keep the Python child wrapper alive: Qt destruction alone must stop the thread.
+        release.set()
+        qtbot.waitUntil(
+            lambda: not shiboken6.isValid(thread) or not thread.isRunning(), timeout=1000
+        )
+        qtbot.waitUntil(lambda: not shiboken6.isValid(worker))
+        qtbot.waitUntil(lambda: not shiboken6.isValid(thread))
+    finally:
+        release.set()
+        if shiboken6.isValid(thread):
+            if shiboken6.isValid(worker):
+                worker.deleteLater()
+            thread.quit()
+            assert thread.wait(2000)
+            thread.deleteLater()
+        if shiboken6.isValid(parent):
+            parent.deleteLater()
+        # Pyqtgraph's unparented menus are owned by Python objects in the deleted plot.
+        del widget, parent
+        gc.collect()
+
+
 def test_heatmap_plot(heatmap_widget):
     heatmap_widget.plot(device_x="samx", device_y="samy", device_z="bpm4i")
 
@@ -651,14 +773,39 @@ def test_finish_interpolation_thread_cleans_references(heatmap_widget):
 
     heatmap_widget._finish_interpolation_thread()
 
-    worker_mock.deleteLater.assert_called_once()
+    worker_mock.deleteLater.assert_not_called()
+    thread_mock.requestInterruption.assert_called_once()
     thread_mock.quit.assert_called_once()
-    thread_mock.wait.assert_called_once()
-    thread_mock.deleteLater.assert_called_once()
+    thread_mock.wait.assert_not_called()
+    thread_mock.deleteLater.assert_not_called()
     assert heatmap_widget._interpolation_worker is None
     assert heatmap_widget._interpolation_thread is None
 
 
+@pytest.mark.parametrize("failed", [False, True], ids=["finished", "failed"])
+def test_queued_interpolation_callback_after_shutdown_is_ignored(heatmap_widget, qtbot, failed):
+    heatmap_widget.scan_id = "shutdown-scan"
+    heatmap_widget._ensure_interpolation_thread()
+    worker = heatmap_widget._interpolation_worker
+    thread = heatmap_widget._interpolation_thread
+    # The signal is queued for the widget's GUI thread and arrives only after cleanup.
+    if failed:
+        worker.failed.emit("interpolation failed", 4, heatmap_widget.scan_id)
+    else:
+        worker.finished.emit(np.ones((2, 2)), QTransform(), 4, heatmap_widget.scan_id)
+    heatmap_widget._finish_interpolation_thread()
+    heatmap_widget._finish_interpolation_thread()  # Shutdown is safe to repeat.
+    with (
+        mock.patch.object(heatmap_widget, "_apply_image_update") as apply_image,
+        mock.patch.object(heatmap_widget, "_maybe_start_pending_interpolation") as start_pending,
+    ):
+        QApplication.processEvents()
+        apply_image.assert_not_called()
+        start_pending.assert_not_called()
+    qtbot.waitUntil(lambda: not shiboken6.isValid(worker))
+    qtbot.waitUntil(lambda: not shiboken6.isValid(thread))
+
+
 def test_device_safe_properties_get(heatmap_widget):
     """Test that device SafeProperty getters work correctly."""
     # Initially devices should be empty