fix(pyamber): disable data sub-queues registered after disable_data

amber/src/main/python/core/models/internal_queue.py

@@ -20,9 +20,8 @@
 from dataclasses import dataclass
 from enum import Enum
 from threading import RLock
-from typing import TypeVar, Set
+from typing import Tuple, TypeVar, Set
 
-from core.models.internal_marker import InternalMarker
 from core.models.payload import DataPayload
 from core.util.customized_queue.linked_blocking_multi_queue import (
     LinkedBlockingMultiQueue,
@@ -77,11 +76,19 @@ def get(self) -> T:
     def put(self, item: T) -> None:
         if isinstance(item, InternalQueueElement):
             if item.tag not in self._queue_ids:
-                self._queue.add_sub_queue(item.tag, 1 if item.tag.is_control else 2)
-                self._queue_ids.add(item.tag)
-            if isinstance(item, (DataElement, InternalMarker, ECMElement)):
-                self._queue.put(item.tag, item)
-            elif isinstance(item, DCMElement):
+                # registration must not interleave with disable_data/enable_data
+                with self._lock:
+                    if item.tag not in self._queue_ids:
+                        self._queue.add_sub_queue(
+                            item.tag, 1 if item.tag.is_control else 2
+                        )
+                        # while data is disabled, a new data sub-queue must
+                        # start disabled too (before its first element is
+                        # enqueued), or it would leak data during pause/backpressure
+                        if not item.tag.is_control and self._queue_state:
+                            self._queue.disable(item.tag)
+                        self._queue_ids.add(item.tag)
+            if isinstance(item, (DataElement, ECMElement, DCMElement)):
                 self._queue.put(item.tag, item)
             else:
                 raise ValueError(f"item {item} is not recognized by internal queue")
@@ -94,19 +101,26 @@ def disable(self, channel_id: ChannelIdentity) -> None:
     def enable(self, channel_id: ChannelIdentity) -> None:
         self._queue.enable(channel_id)
 
+    def _control_queue_ids(self) -> Tuple[ChannelIdentity, ...]:
+        """Snapshot of the registered control channels.
+
+        put() can grow _queue_ids from another thread, and iterating the
+        live set while it grows raises RuntimeError, so queries must iterate
+        a snapshot taken through these helpers.
+        """
+        snapshot = tuple(self._queue_ids)
+        return tuple(queue_id for queue_id in snapshot if queue_id.is_control)
+
+    def _data_queue_ids(self) -> Tuple[ChannelIdentity, ...]:
+        """Snapshot of the registered data channels; see _control_queue_ids."""
+        snapshot = tuple(self._queue_ids)
+        return tuple(queue_id for queue_id in snapshot if not queue_id.is_control)
+
     def is_control_empty(self) -> bool:
-        return all(
-            self.is_empty(queue_id)
-            for queue_id in self._queue_ids
-            if queue_id.is_control
-        )
+        return all(self.is_empty(queue_id) for queue_id in self._control_queue_ids())
 
     def is_data_empty(self) -> bool:
-        return all(
-            self.is_empty(queue_id)
-            for queue_id in self._queue_ids
-            if not queue_id.is_control
-        )
+        return all(self.is_empty(queue_id) for queue_id in self._data_queue_ids())
 
     def __len__(self) -> int:
         return self.size()
@@ -115,18 +129,10 @@ def size(self) -> int:
         return self._queue.size()
 
     def size_control(self) -> int:
-        return sum(
-            self._queue.size(queue_id)
-            for queue_id in self._queue_ids
-            if queue_id.is_control
-        )
+        return sum(self._queue.size(queue_id) for queue_id in self._control_queue_ids())
 
     def size_data(self) -> int:
-        return sum(
-            self._queue.size(queue_id)
-            for queue_id in self._queue_ids
-            if not queue_id.is_control
-        )
+        return sum(self._queue.size(queue_id) for queue_id in self._data_queue_ids())
 
     def enable_data(self, disable_type: DisableType) -> bool:
         with self._lock:
@@ -148,14 +154,10 @@ def disable_data(self, disable_type: DisableType) -> None:
 
     def in_mem_size(self) -> int:
         return sum(
-            self._queue.in_mem_size(queue_id)
-            for queue_id in self._queue_ids
-            if not queue_id.is_control
+            self._queue.in_mem_size(queue_id) for queue_id in self._data_queue_ids()
         )
 
     def is_data_enabled(self) -> bool:
         return any(
-            self._queue.is_enabled(queue_id)
-            for queue_id in self._queue_ids
-            if not queue_id.is_control
+            self._queue.is_enabled(queue_id) for queue_id in self._data_queue_ids()
         )