import asyncio

from haystack import AsyncPipeline, component


class BranchGate:
    def __init__(self, branch_count: int):
        self.branch_count = branch_count
        self.ready = set()
        self.all_ready = asyncio.Event()

    async def wait_for_peers(self, branch_name: str):
        self.ready.add(branch_name)
        if len(self.ready) == self.branch_count:
            self.all_ready.set()
        await asyncio.wait_for(self.all_ready.wait(), timeout=1)