Rework ohlc sampling to launch from .start()

Avoid bothering with a trio event and expect the caller to do manual shm
registering with the write loop. Provide OHLC sample period indexing
through a re-branded pub-sub func ``iter_ohlc_periods()``.

piker/data/_buffer.py

@@ -17,17 +17,25 @@
 """
 Data buffers for fast shared humpy.
 """
-from typing import Tuple, Callable, Dict
-# import time
+from typing import Dict, List
 
 import tractor
 import trio
+from trio_typing import TaskStatus
 
 from ._sharedmem import ShmArray
+from ..log import get_logger
 
 
-_shms: Dict[int, ShmArray] = {}
+log = get_logger(__name__)
+
+
+# TODO: we could stick these in a composed type to avoid
+# angering the "i hate module scoped variables crowd" (yawn).
+_shms: Dict[int, List[ShmArray]] = {}
 _start_increment: Dict[str, trio.Event] = {}
+_incrementers: Dict[int, trio.CancelScope] = {}
+_subscribers: Dict[str, tractor.Context] = {}
 
 
 def shm_incrementing(shm_token_name: str) -> trio.Event:
@@ -35,11 +43,9 @@ def shm_incrementing(shm_token_name: str) -> trio.Event:
     return _start_increment.setdefault(shm_token_name, trio.Event())
 
 
-@tractor.msg.pub
 async def increment_ohlc_buffer(
-    shm_token: dict,
-    get_topics: Callable[..., Tuple[str]],
-    # delay_s: Optional[float] = None,
+    delay_s: int,
+    task_status: TaskStatus[trio.CancelScope] = trio.TASK_STATUS_IGNORED,
 ):
     """Task which inserts new bars into the provide shared memory array
     every ``delay_s`` seconds.
@@ -54,14 +60,16 @@ async def increment_ohlc_buffer(
     the underlying buffers will actually be incremented.
     """
 
-    # wait for brokerd to signal we should start sampling
-    await shm_incrementing(shm_token['shm_name']).wait()
+    # # wait for brokerd to signal we should start sampling
+    # await shm_incrementing(shm_token['shm_name']).wait()
 
     # TODO: right now we'll spin printing bars if the last time stamp is
     # before a large period of no market activity.  Likely the best way
     # to solve this is to make this task aware of the instrument's
     # tradable hours?
 
+    global _incrementers
+
     # adjust delay to compensate for trio processing time
     ad = min(_shms.keys()) - 0.001
 
@@ -69,6 +77,12 @@ async def increment_ohlc_buffer(
     lowest = min(_shms.keys())
     ad = lowest - 0.001
 
+    with trio.CancelScope() as cs:
+
+        # register this time period step as active
+        _incrementers[delay_s] = cs
+        task_status.started(cs)
+
         while True:
             # TODO: do we want to support dynamically
             # adding a "lower" lowest increment period?
@@ -103,13 +117,35 @@ async def increment_ohlc_buffer(
                     shm.push(last)
 
                 # broadcast the buffer index step
-        yield {'index': shm._last.value}
+                # yield {'index': shm._last.value}
+                for ctx in _subscribers.get(delay_s, ()):
+                    try:
+                        await ctx.send_yield({'index': shm._last.value})
+                    except (
+                        trio.BrokenResourceError,
+                        trio.ClosedResourceError
+                    ):
+                        log.error(f'{ctx.chan.uid} dropped connection')
 
 
-def subscribe_ohlc_for_increment(
-    shm: ShmArray,
-    delay: int,
+@tractor.stream
+async def iter_ohlc_periods(
+    ctx: tractor.Context,
+    delay_s: int,
 ) -> None:
-    """Add an OHLC ``ShmArray`` to the increment set.
     """
-    _shms.setdefault(delay, []).append(shm)
+    Subscribe to OHLC sampling "step" events: when the time
+    aggregation period increments, this event stream emits an index
+    event.
+
+    """
+    # add our subscription
+    global _subscribers
+    subs = _subscribers.setdefault(delay_s, [])
+    subs.append(ctx)
+
+    try:
+        # stream and block until cancelled
+        await trio.sleep_forever()
+    finally:
+        subs.remove(ctx)