From 6360c0367622e662e80728e34954c5df75865b94 Mon Sep 17 00:00:00 2001
From: Alexander Clausen <alex@gc-web.de>
Date: Thu, 22 Jun 2023 10:58:56 +0200
Subject: [PATCH] asi_tpx3: don't crash in cancelled scan

Happens often at the beginning of the acquisition, before
synchronization is complete
---
 .../detectors/asi_tpx3/acquisition.py         | 26 +++++++++++++------
 1 file changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/libertem_live/detectors/asi_tpx3/acquisition.py b/src/libertem_live/detectors/asi_tpx3/acquisition.py
index bf53db91..92b6ca2e 100644
--- a/src/libertem_live/detectors/asi_tpx3/acquisition.py
+++ b/src/libertem_live/detectors/asi_tpx3/acquisition.py
@@ -11,13 +11,18 @@
 
 from libertem.common import Shape, Slice
 from libertem.common.math import prod
-from libertem.common.executor import WorkerContext, TaskProtocol, WorkerQueue, TaskCommHandler
+from libertem.common.executor import (
+    WorkerContext, TaskProtocol, WorkerQueue, TaskCommHandler,
+    JobCancelledError,
+)
 from libertem.io.dataset.base import (
     DataTile, DataSetMeta, BasePartition, Partition, DataSet, TilingScheme,
 )
 from libertem.corrections.corrset import CorrectionSet
 
-from libertem_live.detectors.base.acquisition import AcquisitionMixin
+from libertem_live.detectors.base.acquisition import (
+    AcquisitionMixin,
+)
 from libertem_live.detectors.base.controller import AcquisitionController
 from libertem_live.hooks import ReadyForDataEnv, Hooks
 from .connection import AsiTpx3DetectorConnection, AsiTpx3PendingAcquisition
@@ -101,14 +106,17 @@ def handle_task(self, task: TaskProtocol, queue: WorkerQueue):
                 chunk_stack = self.conn.get_next_stack(
                     max_size=current_stack_size
                 )
+                if chunk_stack is None:
+                    if current_idx != end_idx:
+                        queue.put({
+                            "type": "END_PARTITION",
+                        })
+                        raise JobCancelledError("premature end of frame iterator")
+                    break
                 assert len(chunk_stack) <= current_stack_size,\
                     f"{len(chunk_stack)} <= {current_stack_size}"
                 t1 = time.perf_counter()
                 recv_time += t1 - t0
-                if chunk_stack is None:
-                    if current_idx != end_idx:
-                        raise RuntimeError("premature end of frame iterator")
-                    break
 
                 t0 = time.perf_counter()
                 serialized = chunk_stack.serialize()
@@ -343,8 +351,10 @@ def _get_tiles_straight(
             try:
                 stack = next(stacks)
             except StopIteration:
-                assert to_read == 0, f"we were still expecting to read {to_read} frames more!"
-
+                if to_read != 0:
+                    raise JobCancelledError(
+                        f"we were still expecting to read {to_read} frames more!"
+                    )
             for (
                 chunk_layout,
                 chunk_indptr,