Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 27 additions & 5 deletions src/omnibase_infra/event_bus/event_bus_kafka.py
Original file line number Diff line number Diff line change
Expand Up @@ -2275,7 +2275,11 @@ async def start_consuming(self) -> None:
if not self._started:
await self.start()

# Collect topics that need consumers while holding lock briefly
# Collect topics that need consumers while holding lock briefly.
# Pre-populate _pending_consumer_keys inside the lock so concurrent
# _start_consumer_for_topic_unlocked calls (via asyncio.gather below)
# cannot race and create duplicate consumers — the same reservation
# pattern used by subscribe().
topics_to_start: list[tuple[str, str]] = []
async with self._lock:
for topic in self._subscribers:
Expand All @@ -2284,12 +2288,30 @@ async def start_consuming(self) -> None:
if group_id in seen_group_ids:
continue
seen_group_ids.add(group_id)
if (topic, group_id) not in self._group_consumers:
consumer_key = (topic, group_id)
if (
consumer_key not in self._group_consumers
and consumer_key not in self._pending_consumer_keys
):
self._pending_consumer_keys.add(consumer_key)
topics_to_start.append((topic, group_id))

# Start consumers outside the lock to avoid blocking
for topic, group_id in topics_to_start:
await self._start_consumer_for_topic(topic, group_id)
# Start all consumers concurrently — wall-clock time becomes the
# slowest single consumer rather than the sum of all N consumers.
# return_exceptions=True lets every consumer attempt to start; we
# collect failures and re-raise the first one after all have settled
# so a single bad topic does not starve the rest.
if topics_to_start:
results = await asyncio.gather(
*[
self._start_consumer_for_topic_unlocked(topic, group_id)
for topic, group_id in topics_to_start
],
return_exceptions=True,
)
errors = [r for r in results if isinstance(r, BaseException)]
if errors:
raise errors[0]

# Block until shutdown
while not self._shutdown:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,3 +95,49 @@ async def fake_start(start_topic: str, group_id: str) -> None:
assert len(started) == 1
assert started[0][0] == topic
assert len(bus._subscribers[topic]) == 5


@pytest.mark.asyncio
async def test_start_consuming_starts_distinct_consumers_in_parallel() -> None:
"""start_consuming reserves all keys and starts distinct consumers concurrently."""
bus = _make_bus()
bus._started = True

active_starts = 0
max_active_starts = 0
started: list[tuple[str, str]] = []

async def fake_start(topic: str, group_id: str) -> None:
nonlocal active_starts, max_active_starts
active_starts += 1
max_active_starts = max(max_active_starts, active_starts)
await asyncio.sleep(0)
bus._group_consumers[(topic, group_id)] = AsyncMock()
bus._pending_consumer_keys.discard((topic, group_id))
started.append((topic, group_id))
active_starts -= 1

bus._start_consumer_for_topic_unlocked = AsyncMock( # type: ignore[method-assign]
side_effect=fake_start
)

group_a = "service-a"
group_b = "service-b"
async with bus._lock:
bus._subscribers["onex.evt.omnibase-infra.start-consuming-a.v1"] = [
(group_a, "sub-a-1", _handler),
(group_b, "sub-a-2", _handler),
]
bus._subscribers["onex.evt.omnibase-infra.start-consuming-b.v1"] = [
(group_a, "sub-b-1", _handler),
]

task = asyncio.create_task(bus.start_consuming())
await asyncio.sleep(0.1)
await bus.shutdown()
await asyncio.wait_for(task, timeout=2.0)

assert bus._start_consumer_for_topic_unlocked.await_count == 3
assert len(started) == 3
assert len(set(started)) == 3
assert max_active_starts > 1
54 changes: 54 additions & 0 deletions tests/unit/event_bus/test_kafka_event_bus.py
Original file line number Diff line number Diff line change
Expand Up @@ -2059,6 +2059,60 @@ async def consume_with_signal() -> None:
# Task should complete
await asyncio.wait_for(task, timeout=1.0)

@pytest.mark.asyncio
async def test_start_consuming_starts_consumers_concurrently(
self, mock_producer: AsyncMock
) -> None:
"""start_consuming() launches all consumers concurrently via asyncio.gather.

Verifies that _start_consumer_for_topic_unlocked is called once per
distinct (topic, group_id) pair, that pending keys are reserved inside
the lock before any concurrent call begins (no duplicates), and that
the wall-clock order of completion does not cause double-starts.
"""
call_log: list[tuple[str, str]] = []

async def fake_start_unlocked(topic: str, group_id: str) -> None:
call_log.append((topic, group_id))
# Yield control to simulate concurrent group-join latency
await asyncio.sleep(0)

with patch(
"omnibase_infra.event_bus.event_bus_kafka.AIOKafkaProducer",
return_value=mock_producer,
):
config = ModelKafkaEventBusConfig(bootstrap_servers=TEST_BOOTSTRAP_SERVERS)
event_bus = EventBusKafka(config=config)
await event_bus.start()

# Pre-register three distinct subscriptions without triggering
# consumer startup (bus is already started but we bypass it here by
# directly injecting into _subscribers so we can test start_consuming
# in isolation).
group_id_a = "svc-a"
group_id_b = "svc-b"

async with event_bus._lock:
event_bus._subscribers["topic-x"] = [
(group_id_a, "sub-1", AsyncMock()),
(group_id_b, "sub-2", AsyncMock()),
]
event_bus._subscribers["topic-y"] = [
(group_id_a, "sub-3", AsyncMock()),
]

event_bus._start_consumer_for_topic_unlocked = fake_start_unlocked # type: ignore[method-assign]

# Run start_consuming briefly then shut down
task = asyncio.create_task(event_bus.start_consuming())
await asyncio.sleep(0.1)
await event_bus.shutdown()
await asyncio.wait_for(task, timeout=2.0)

# Each distinct (topic, group_id) pair must be started exactly once
assert len(call_log) == 3
assert len(set(call_log)) == 3, "duplicate consumer starts detected"


class TestKafkaEventBusConfig:
"""Test suite for config-based EventBusKafka construction."""
Expand Down
Loading