Repository: bluesky/bluesky
Pull request #1648 (2023-12-26 17:11:06)

FIX: handle case of rewinding to before the beginning of a stream
<!--- Provide a general summary of your changes in the Title above -->

## Description
<!--- Describe your changes in detail -->

The issue is that:

 - we track the event count (for seq_num which is non-monotonic and
   can repeat if rewound) in `Bundler._sequence_counters`
 - we snapshot the counters at checkpoints
 - on rewind we clear the _event_counters dictionary and then .update it
   from the copy

In the case of implicitly creating a stream between a checkpoint and a
pause (so there is `declare_stream` message) the above process effectively
drops our tracking in the event counters.

In `_prepare_stream` we ensure that the needed counter exists (and starts at 1)
so replicate that logic here.

## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. -->

We should not crash this way.

## How Has This Been Tested?
<!--- Please describe in detail how you tested your changes. -->
<!--- Include details of your testing environment, and the tests you ran to -->
<!--- see how your change affects other areas of the code, etc. -->

Manually + adding a test



Patch:
diff --git a/bluesky/bundlers.py b/bluesky/bundlers.py
index f7ae5db9..aabbc085 100644
--- a/bluesky/bundlers.py
+++ b/bluesky/bundlers.py
@@ -2,7 +2,7 @@ import asyncio
 import inspect
 import time as ttime
 from collections import deque
-from typing import (Any, Deque, Dict, FrozenSet, Iterable, List, Optional, Set,
+from typing import (Any, Deque, Dict, FrozenSet, Iterable, List, Optional,
                     Tuple, Union)
 
 from event_model import (ComposeDescriptorBundle, DataKey, Datum,
@@ -42,7 +42,7 @@ class RunBundler:
         self._config_ts_cache: ObjDict[Any] = dict()  # " obj.read_configuration() timestamps
         # cache of {name: (doc, compose_event, compose_event_page)}
         self._descriptors: Dict[Any, ComposeDescriptorBundle] = dict()
-        self._descriptor_objs: Dict[str, Set[HasName]] = dict()
+        self._descriptor_objs: Dict[str, Dict[HasName, Dict[str, DataKey]]] = dict()
         # cache of {obj: {objs_frozen_set: (doc, compose_event, compose_event_page)}
         self._local_descriptors: Dict[Any, Dict[FrozenSet[str], ComposeDescriptorBundle]] = dict()
         # a seq_num counter per stream
@@ -139,7 +139,7 @@ class RunBundler:
     async def _prepare_stream(
         self,
         desc_key: str,
-        objs_dks: Dict[Any, Dict[str, DataKey]],
+        objs_dks: Dict[HasName, Dict[str, DataKey]],
     ):
         # We do not have an Event Descriptor for this set
         # so one must be created.
@@ -430,6 +430,13 @@ class RunBundler:
     def rewind(self):
         self._sequence_counters.clear()
         self._sequence_counters.update(self._sequence_counters_copy)
+        # make sure we do not forget about streams we roll back to the
+        # very beginning of
+        for desc_key in self._descriptor_objs:
+            if desc_key not in self._sequence_counters:
+                self._sequence_counters[desc_key] = 1
+                self._sequence_counters_copy[desc_key] = 1
+
         # This is needed to 'cancel' an open bundling (e.g. create) if
         # the pause happens after a 'checkpoint', after a 'create', but
         # before the paired 'save'.


Test patch:
diff --git a/bluesky/tests/test_run_engine.py b/bluesky/tests/test_run_engine.py
index 1d9f2274..6abdfc08 100644
--- a/bluesky/tests/test_run_engine.py
+++ b/bluesky/tests/test_run_engine.py
@@ -15,7 +15,7 @@ from event_model import DocumentNames
 
 from bluesky import Msg
 from bluesky.plan_stubs import (abs_set, checkpoint, declare_stream,
-                                trigger_and_read, wait)
+                                trigger_and_read, wait, pause)
 from bluesky.plans import count, grid_scan
 from bluesky.preprocessors import (SupplementalData, baseline_wrapper,
                                    finalize_wrapper, reset_positions_decorator,
@@ -1903,3 +1903,60 @@ def test_wait_with_timeout(set_finished, RE):
     else:
         with pytest.raises(TimeoutError):
             RE(plan())
+
+
+def test_1event_rewind(RE, hw):
+    cache = []
+
+    def hook(msg):
+        cache.append(msg)
+
+    doc_cache = []
+
+    def cb(name, doc):
+        doc_cache.append((name, doc))
+
+    @run_decorator()
+    def broken():
+        yield from checkpoint()
+        yield from trigger_and_read([hw.det], name="primary")
+        yield from pause()
+        yield from trigger_and_read([hw.det], name="primary")
+
+    RE.msg_hook = hook
+    with pytest.raises(RunEngineInterrupted):
+        RE(broken(), cb)
+    RE.resume()
+
+    assert [_[1]["seq_num"] for _ in doc_cache if _[0] == "event"] == [1, 1, 2]
+
+    assert [_[0] for _ in doc_cache] == [
+        "start",
+        "descriptor",
+        "event",
+        "event",
+        "event",
+        "stop",
+    ]
+
+    assert [msg.command for msg in cache] == [
+        "open_run",
+        "checkpoint",
+        "trigger",
+        "wait",
+        "create",
+        "read",
+        "save",
+        "pause",
+        "trigger",
+        "wait",
+        "create",
+        "read",
+        "save",
+        "trigger",
+        "wait",
+        "create",
+        "read",
+        "save",
+        "close_run",
+    ]