From 7a5c43d01ae72b4a15f108153cf4cafc0efd6f2d Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Tue, 6 Jun 2023 17:29:03 -0400
Subject: [PATCH] Support injecting a `info: dict` to `Sampler.broadcast_all()`
 calls

---
 piker/data/_sampling.py | 39 ++++++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/piker/data/_sampling.py b/piker/data/_sampling.py
index a8dba30b..98a7603f 100644
--- a/piker/data/_sampling.py
+++ b/piker/data/_sampling.py
@@ -230,6 +230,7 @@ class Sampler:
         self,
         period_s: float,
         time_stamp: float | None = None,
+        info: dict | None = None,
 
     ) -> None:
         '''
@@ -258,10 +259,14 @@ class Sampler:
             try:
                 for stream in (subs - sent):
                     try:
-                        await stream.send({
+                        msg = {
                             'index': time_stamp or last_ts,
                             'period': period_s,
-                        })
+                        }
+                        if info:
+                            msg.update(info)
+
+                        await stream.send(msg)
                         sent.add(stream)
 
                     except (
@@ -287,9 +292,15 @@ class Sampler:
                 )
 
     @classmethod
-    async def broadcast_all(self) -> None:
+    async def broadcast_all(
+        self,
+        info: dict | None = None,
+    ) -> None:
         for period_s in self.subscribers:
-            await self.broadcast(period_s)
+            await self.broadcast(
+                period_s,
+                info=info,
+            )
 
 
 @tractor.context
@@ -359,8 +370,10 @@ async def register_with_sampler(
 
                         # except broadcast requests from the subscriber
                         async for msg in stream:
-                            if msg == 'broadcast_all':
-                                await Sampler.broadcast_all()
+                            if 'broadcast_all' in msg:
+                                await Sampler.broadcast_all(
+                                    info=msg['broadcast_all'],
+                                )
                 finally:
                     if (
                         sub_for_broadcasts
@@ -468,6 +481,8 @@ async def open_sample_stream(
     cache_key: str | None = None,
     allow_new_sampler: bool = True,
 
+    ensure_is_active: bool = False,
+
 ) -> AsyncIterator[dict[str, float]]:
     '''
     Subscribe to OHLC sampling "step" events: when the time aggregation
@@ -510,12 +525,18 @@ async def open_sample_stream(
             },
         ) as (ctx, first)
     ):
-        assert len(first) > 1
+        if ensure_is_active:
+            assert len(first) > 1
+
         async with (
             ctx.open_stream() as istream,
 
-            # TODO: we don't need this task-bcasting right?
-            # istream.subscribe() as istream,
+            # TODO: we DO need this task-bcasting so that
+            # for eg. the history chart update loop eventually
+            # receceives all backfilling event msgs such that
+            # the underlying graphics format arrays are
+            # re-allocated until all history is loaded!
+            istream.subscribe() as istream,
         ):
             yield istream