def pub( wrapped: typing.Callable = None, *, tasks: Set[str] = set(), ): """Publisher async generator decorator. A publisher can be called multiple times from different actors but will only spawn a finite set of internal tasks to stream values to each caller. The ``tasks: Set[str]`` argument to the decorator specifies the names of the mutex set of publisher tasks. When the publisher function is called, an argument ``task_name`` must be passed to specify which task (of the set named in ``tasks``) should be used. This allows for using the same publisher with different input (arguments) without allowing more concurrent tasks then necessary. Values yielded from the decorated async generator must be ``Dict[str, Dict[str, Any]]`` where the fist level key is the topic string and determines which subscription the packet will be delivered to and the value is a packet ``Dict[str, Any]`` by default of the form: .. ::python {topic: str: value: Any} The caller can instead opt to pass a ``packetizer`` callback who's return value will be delivered as the published response. The decorated async generator function must accept an argument :func:`get_topics` which dynamically returns the tuple of current subscriber topics: .. code:: python from tractor.msg import pub @pub(tasks={'source_1', 'source_2'}) async def pub_service(get_topics): data = await web_request(endpoints=get_topics()) for item in data: yield data['key'], data The publisher must be called passing in the following arguments: - ``topics: Set[str]`` the topic sequence or "subscriptions" - ``task_name: str`` the task to use (if ``tasks`` was passed) - ``ctx: Context`` the tractor context (only needed if calling the pub func without a nursery, otherwise this is provided implicitly) - packetizer: ``Callable[[str, Any], Any]`` a callback who receives the topic and value from the publisher function each ``yield`` such that whatever is returned is sent as the published value to subscribers of that topic. By default this is a dict ``{topic: str: value: Any}``. As an example, to make a subscriber call the above function: .. code:: python from functools import partial import tractor async with tractor.open_nursery() as n: portal = n.run_in_actor( 'publisher', # actor name partial( # func to execute in it pub_service, topics=('clicks', 'users'), task_name='source1', ) ) async for value in await portal.result(): print(f"Subscriber received {value}") Here, you don't need to provide the ``ctx`` argument since the remote actor provides it automatically to the spawned task. If you were to call ``pub_service()`` directly from a wrapping function you would need to provide this explicitly. Remember you only need this if you need *a finite set of tasks* running in a single actor to stream data to an arbitrary number of subscribers. If you are ok to have a new task running for every call to ``pub_service()`` then probably don't need this. """ global _pubtask2lock # handle the decorator not called with () case if wrapped is None: return partial(pub, tasks=tasks) task2lock: Dict[str, trio.StrictFIFOLock] = {} for name in tasks: task2lock[name] = trio.StrictFIFOLock() @wrapt.decorator async def wrapper(agen, instance, args, kwargs): # XXX: this is used to extract arguments properly as per the # `wrapt` docs async def _execute( ctx: Context, topics: Set[str], *args, # *, task_name: str = None, # default: only one task allocated packetizer: Callable = None, **kwargs, ): if task_name is None: task_name = trio.lowlevel.current_task().name if tasks and task_name not in tasks: raise TypeError( f"{agen} must be called with a `task_name` named " f"argument with a value from {tasks}") elif not tasks and not task2lock: # add a default root-task lock if none defined task2lock[task_name] = trio.StrictFIFOLock() _pubtask2lock.update(task2lock) topics = set(topics) lock = _pubtask2lock[task_name] all_subs = _pub_state.setdefault('_subs', {}) topics2ctxs = all_subs.setdefault(task_name, {}) try: modify_subs(topics2ctxs, topics, ctx) # block and let existing feed task deliver # stream data until it is cancelled in which case # the next waiting task will take over and spawn it again async with lock: # no data feeder task yet; so start one respawn = True while respawn: respawn = False log.info(f"Spawning data feed task for {funcname}") try: # unblocks when no more symbols subscriptions exist # and the streamer task terminates await fan_out_to_ctxs( pub_async_gen_func=partial( agen, *args, **kwargs), topics2ctxs=topics2ctxs, packetizer=packetizer, ) log.info( f"Terminating stream task {task_name or ''}" f" for {agen.__name__}") except trio.BrokenResourceError: log.exception("Respawning failed data feed task") respawn = True finally: # remove all subs for this context modify_subs(topics2ctxs, set(), ctx) # if there are truly no more subscriptions with this broker # drop from broker subs dict if not any(topics2ctxs.values()): log.info( f"No more subscriptions for publisher {task_name}") # invoke it await _execute(*args, **kwargs) funcname = wrapped.__name__ if not inspect.isasyncgenfunction(wrapped): raise TypeError( f"Publisher {funcname} must be an async generator function") if 'get_topics' not in inspect.signature(wrapped).parameters: raise TypeError(f"Publisher async gen {funcname} must define a " "`get_topics` argument") # XXX: manually monkey the wrapped function since # ``wrapt.decorator`` doesn't seem to want to play nice with its # whole "adapter" thing... wrapped._tractor_stream_function = True # type: ignore return wrapper(wrapped)