diff --git a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py index fcf175d6bcd..561b1680cf5 100644 --- a/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py +++ b/amber/src/main/python/core/util/customized_queue/linked_blocking_multi_queue.py @@ -227,16 +227,13 @@ def get_next_sub_queue(self) -> Optional[LinkedBlockingMultiQueue.SubQueue[T]]: return None def peek(self) -> Optional[T]: - start_idx = self.next_idx + idx = self.next_idx while True: - child = self.queues[self.next_idx] + child = self.queues[idx] if child.enabled and child.size() > 0: return child.head.next.item - else: - self.next_idx += 1 - if self.next_idx == len(self.queues): - self.next_idx = 0 - if self.next_idx == start_idx: + idx = (idx + 1) % len(self.queues) + if idx == self.next_idx: break return None diff --git a/amber/src/test/python/core/util/customized_queue/test_linked_blocking_multi_queue.py b/amber/src/test/python/core/util/customized_queue/test_linked_blocking_multi_queue.py index 8cdaada6bd8..6b68c6796af 100644 --- a/amber/src/test/python/core/util/customized_queue/test_linked_blocking_multi_queue.py +++ b/amber/src/test/python/core/util/customized_queue/test_linked_blocking_multi_queue.py @@ -273,7 +273,7 @@ def test_returns_none_when_the_only_non_empty_sub_queue_is_disabled(self, queue) queue.disable("data") assert queue.peek() is None - def test_priority_group_peek_perturbs_round_robin_state(self): + def test_priority_group_peek_preserves_round_robin_state(self): lbmq = LinkedBlockingMultiQueue() lbmq.add_sub_queue("first", 1) lbmq.add_sub_queue("second", 1) @@ -282,9 +282,10 @@ def test_priority_group_peek_perturbs_round_robin_state(self): assert group.next_idx == 0 assert group.peek() == "s" - # Skipping the empty `first` queue advanced next_idx, so peek is not - # side-effect free on the group's round-robin cursor. - assert group.next_idx == 1 + assert group.next_idx == 0 + + lbmq.put("first", "f") + assert lbmq.get() == "f" def test_priority_group_peek_returns_none_when_all_queues_empty(self, queue): group = queue.get_sub_queue("control").priority_group