From 7b653fe4f44587e42c28375fbed46b3c3be4b9f4 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Sun, 28 Aug 2022 20:28:42 -0400
Subject: [PATCH] Store shm array size in token schema, use for loading

---
 piker/data/_sharedmem.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/piker/data/_sharedmem.py b/piker/data/_sharedmem.py
index 82f61e79..8cf09aa5 100644
--- a/piker/data/_sharedmem.py
+++ b/piker/data/_sharedmem.py
@@ -118,6 +118,7 @@ class _Token(Struct, frozen=True):
     shm_first_index_name: str
     shm_last_index_name: str
     dtype_descr: tuple
+    size: int
 
     @property
     def dtype(self) -> np.dtype:
@@ -158,6 +159,7 @@ def get_shm_token(key: str) -> _Token:
 
 def _make_token(
     key: str,
+    size: int,
     dtype: Optional[np.dtype] = None,
 ) -> _Token:
     '''
@@ -170,7 +172,8 @@ def _make_token(
         shm_name=key,
         shm_first_index_name=key + "_first",
         shm_last_index_name=key + "_last",
-        dtype_descr=tuple(np.dtype(dtype).descr)
+        dtype_descr=tuple(np.dtype(dtype).descr),
+        size=size,
     )
 
 
@@ -222,6 +225,7 @@ class ShmArray:
             shm_first_index_name=self._first._shm.name,
             shm_last_index_name=self._last._shm.name,
             dtype_descr=tuple(self._array.dtype.descr),
+            size=self._len,
         )
 
     @property
@@ -467,7 +471,8 @@ def open_shm_array(
 
     token = _make_token(
         key=key,
-        dtype=dtype
+        size=size,
+        dtype=dtype,
     )
 
     # create single entry arrays for storing an first and last indices
@@ -527,7 +532,7 @@ def open_shm_array(
 
 def attach_shm_array(
     token: tuple[str, str, tuple[str, str]],
-    size: int = _default_size,
+    # size: int = _default_size,
     readonly: bool = True,
 
 ) -> ShmArray:
@@ -566,7 +571,7 @@ def attach_shm_array(
             raise _err
 
     shmarr = np.ndarray(
-        (size,),
+        (token.size,),
         dtype=token.dtype,
         buffer=shm.buf
     )
@@ -634,6 +639,7 @@ def maybe_open_shm_array(
     use ``attach_shm_array``.
 
     '''
+    size = kwargs.pop('size', _default_size)
     try:
         # see if we already know this key
         token = _known_tokens[key]
@@ -641,7 +647,11 @@ def maybe_open_shm_array(
     except KeyError:
         log.warning(f"Could not find {key} in shms cache")
         if dtype:
-            token = _make_token(key, dtype)
+            token = _make_token(
+                key,
+                size=size,
+                dtype=dtype,
+            )
             try:
                 return attach_shm_array(token=token, **kwargs), False
             except FileNotFoundError: