def _collect_agents_for_update(self,
tps: Set[TP]) -> Dict[AgentT, Set[TP]]:
by_agent: Dict[AgentT, Set[TP]] = defaultdict(set)
for topic, tps in tp_set_to_map(tps).items():
for agent in self._by_topic[topic]:
by_agent[agent].update(tps)
return by_agent
def _update_tp_index(self, assigned: Set[TP]) -> None:
assignmap = tp_set_to_map(assigned)
tp_index = self._tp_index
for topic in self._topics:
if topic.active_partitions is not None:
# Isolated Partitions: One agent per partition.
if topic.active_partitions:
assert topic.active_partitions.issubset(assigned)
for tp in topic.active_partitions:
tp_index[tp].add(topic)
else:
# Default: One agent receives messages for all partitions.
for subtopic in topic.topics:
for tp in assignmap[subtopic]:
tp_index[tp].add(topic)
