Dedup workers: one _PooledWorker base for the three pool fan-outs

DcPrecomputeWorker, Ch4MaskWorker and CrossCorrelateWorker shared the
same submit/as_completed/emit/shutdown skeleton; they are now
constructors plus _plan/_items/_one/_emit hooks. All three inherit the
cancellation-aware shutdown (wait=not stopped, cancel_futures) that
only DcPrecomputeWorker had before, and the per-run budget attributes
are initialised in __init__ instead of appearing mid-run.

Also removes the viewer's write-only _dc_precompute_worker plumbing
(job tracking already owns worker lifetime via _jobs).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

sras_viewer.py

@@ -1478,7 +1478,6 @@ class SrasViewerWindow(QMainWindow):
         # same combination is free.
         self._dc_cache: dict[tuple[int, int], np.ndarray] = {}
         self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {}
-        self._dc_precompute_worker: DcPrecomputeWorker | None = None
         self._dc_generation: int = 0
 
         # Angle alignment ("Fusion" menu)
@@ -2463,8 +2462,7 @@ class SrasViewerWindow(QMainWindow):
         n_angles = self._sras.n_angles
 
         worker = DcPrecomputeWorker(self._sras)
-        self._dc_precompute_worker = worker
+        self._run_worker(
-        started = self._run_worker(
             "dc_precompute", worker,
             connect=(
                 ("angle_done", lambda a, dc3, dc4, g=generation:
@@ -2473,10 +2471,7 @@ class SrasViewerWindow(QMainWindow):
                     f"DC precompute error: {msg}", 5000)),
             ),
             quit_on=("finished", "error"),
-            on_done=lambda: setattr(self, "_dc_precompute_worker", None),
         )
-        if not started:
-            self._dc_precompute_worker = None
 
     def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int,
                                      dc3_mv: np.ndarray, dc4_mv: np.ndarray,

sras_workers.py

@@ -54,6 +54,34 @@ class CancellableWorker(QObject):
         return self._stop
 
 
+class _PooledWorker(CancellableWorker):
+    """Fans a per-item computation across a thread pool, emitting each result
+    from this worker's own thread as it lands (never from a pool thread).
+
+    Subclasses provide _plan() -> n_workers (stashing whatever per-run
+    context they need), _items(), _one(item) -> result, and _emit(result).
+    On stop(): queued items are dropped, in-flight ones are not waited for —
+    that is what keeps closing the window responsive on a large scan.
+    """
+    finished = pyqtSignal()
+    error = pyqtSignal(str)
+
+    def run(self):
+        try:
+            pool = ThreadPoolExecutor(max_workers=max(1, self._plan()))
+            try:
+                futures = [pool.submit(self._one, it) for it in self._items()]
+                for fut in as_completed(futures):
+                    if self._stop:
+                        break
+                    self._emit(fut.result())
+            finally:
+                pool.shutdown(wait=not self._stop, cancel_futures=True)
+            self.finished.emit()
+        except Exception as exc:
+            self.error.emit(str(exc))
+
+
 class LoadWorker(QObject):
     finished = pyqtSignal(object)   # SrasFile | None
     error = pyqtSignal(str)
@@ -117,29 +145,29 @@ class ComputeWorker(CancellableWorker):
             self.error.emit(str(exc))
 
 
-class DcPrecomputeWorker(CancellableWorker):
+class DcPrecomputeWorker(_PooledWorker):
     """Computes CH3/CH4 DC images for every angle in the background.
 
     DC images are cheap (a per-waveform mean, no FFT) compared to the
     CH1/Velocity FFT, so precomputing them for the whole file right after load
     makes switching angles instant while on a DC channel, and also means the
     FFT masking step (which needs a DC4 image) rarely has to wait on anything.
-
-    Angles are computed on a thread pool — the work is a pure mean over the
-    waveform block, so it is I/O- and bandwidth-bound and embarrassingly
-    parallel. Results are emitted one at a time as they land (out of angle
-    order), and always from this worker's own thread: nothing emits a Qt
-    signal from a pool thread.
     """
     angle_done = pyqtSignal(int, np.ndarray, np.ndarray)   # angle_idx, dc3_mv, dc4_mv
-    finished = pyqtSignal()
-    error = pyqtSignal(str)
 
     def __init__(self, sras: SrasFile):
         super().__init__()
         self._sras = sras
+        self._angle_budget = 0
 
-    def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
+    def _plan(self) -> int:
+        n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
+        return n_workers
+
+    def _items(self):
+        return range(self._sras.n_angles)
+
+    def _one(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
         # max_workers=1 *and* a budget share: this call is one of several
         # concurrent angles, and both the thread count and the buffer size
         # have to be divided (see compute.plan_angle_level).
@@ -149,26 +177,8 @@ class DcPrecomputeWorker(CancellableWorker):
                 dc_image_mv(self._sras, a, CH3_IDX, **kw),
                 dc_image_mv(self._sras, a, CH4_IDX, **kw))
 
-    def run(self):
+    def _emit(self, result):
-        try:
+        self.angle_done.emit(*result)
-            n = self._sras.n_angles
-            n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
-            pool = ThreadPoolExecutor(max_workers=n_workers)
-            try:
-                futures = {pool.submit(self._one_angle, a): a for a in range(n)}
-                for fut in as_completed(futures):
-                    if self._stop:
-                        break
-                    a, dc3, dc4 = fut.result()
-                    self.angle_done.emit(a, dc3, dc4)
-            finally:
-                # cancel_futures drops the queued angles; should_stop lets the
-                # in-flight ones bail within a chunk. Not waiting here is what
-                # keeps closing the window responsive on a large scan.
-                pool.shutdown(wait=not self._stop, cancel_futures=True)
-            self.finished.emit()
-        except Exception as exc:
-            self.error.emit(str(exc))
 
 
 class BatchCacheWorker(QObject):
@@ -306,7 +316,7 @@ class AngleAlignmentWorker(QObject):
             self.finished.emit(None, str(exc))
 
 
-class Ch4MaskWorker(QObject):
+class Ch4MaskWorker(_PooledWorker):
     """Fetches each requested angle's CH4 (Bias B) DC image in mV, for
     ManualAlignmentDialog's initial threshold-mask overlay.
 
@@ -320,35 +330,29 @@ class Ch4MaskWorker(QObject):
     ManualAlignmentDialog._start_mask_prep).
     """
     angle_done = pyqtSignal(int, np.ndarray)   # angle_idx, dc4_mv
-    finished = pyqtSignal()
-    error = pyqtSignal(str)
 
     def __init__(self, sras: SrasFile, angle_indices: list[int]):
         super().__init__()
         self._sras = sras
         self._angles = angle_indices
+        self._budget = 0
 
-    def run(self):
+    def _plan(self) -> int:
-        try:
+        n_workers, self._budget = compute.plan_angle_level(self._sras)
-            n_workers, budget = compute.plan_angle_level(self._sras)
+        return n_workers
-            pool = ThreadPoolExecutor(max_workers=n_workers)
+
-            try:
+    def _items(self):
-                futures = {
+        return self._angles
-                    pool.submit(dc_image_mv, self._sras, a, CH4_IDX,
+
-                               max_workers=1, budget=budget): a
+    def _one(self, a: int) -> tuple[int, np.ndarray]:
-                    for a in self._angles
+        return a, dc_image_mv(self._sras, a, CH4_IDX,
-                }
+                              max_workers=1, budget=self._budget)
-                for fut in as_completed(futures):
+
-                    a = futures[fut]
+    def _emit(self, result):
-                    self.angle_done.emit(a, fut.result())
+        self.angle_done.emit(*result)
-            finally:
-                pool.shutdown(wait=True)
-            self.finished.emit()
-        except Exception as exc:
-            self.error.emit(str(exc))
 
 
-class CrossCorrelateWorker(QObject):
+class CrossCorrelateWorker(_PooledWorker):
     """Rigid registration (rotation + translation, never scale) of each of
     *angle_indices* against *ref_angle_idx*, for ManualAlignmentDialog's Auto
     Cross-Correlate button.
@@ -363,8 +367,6 @@ class CrossCorrelateWorker(QObject):
     """
     # angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source
     angle_done = pyqtSignal(int, float, float, float, float, str)
-    finished = pyqtSignal()
-    error = pyqtSignal(str)
 
     def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
                 dc4_mv: dict[int, np.ndarray], *,
@@ -379,25 +381,19 @@ class CrossCorrelateWorker(QObject):
         self._threshold = dc_threshold_mv
         self._search_deg = search_deg
 
+    def _plan(self) -> int:
+        return compute._registration_workers(self._sras, compute._DEFAULT_FINE_DIM)
+
+    def _items(self):
+        return self._angles
+
     def _one(self, a: int) -> tuple[int, compute.RigidFit]:
         return a, compute.register_angle_to_reference(
             self._sras, a, self._ref, self._dc4_mv,
             dc_threshold_mv=self._threshold, sources=self._sources,
             search_deg=self._search_deg)
 
-    def run(self):
+    def _emit(self, result):
-        try:
+        a, fit = result
-            n_workers = compute._registration_workers(
-                self._sras, compute._DEFAULT_FINE_DIM)
-            pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
-            try:
-                futures = [pool.submit(self._one, a) for a in self._angles]
-                for fut in as_completed(futures):
-                    a, fit = fut.result()
         self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
                              fit.shift_mm[1], fit.score, fit.source)
-            finally:
-                pool.shutdown(wait=True)
-            self.finished.emit()
-        except Exception as exc:
-            self.error.emit(str(exc))