Source code for artiq.protocols.sync_struct

"""This module helps synchronizing a mutable Python structure owned and
modified by one process (the *publisher*) with copies of it (the
*subscribers*) in different processes and possibly different machines.

Synchronization is achieved by sending a full copy of the structure to each
subscriber upon connection (*initialization*), followed by dictionaries
describing each modification made to the structure (*mods*, see
:class:`ModAction`).

Structures must be PYON serializable and contain only lists, dicts, and
immutable types. Lists and dicts can be nested arbitrarily.
"""

import asyncio
from enum import Enum, unique
from operator import getitem
from functools import partial

from artiq.monkey_patches import *
from artiq.protocols import pyon
from artiq.protocols.asyncio_server import AsyncioServer


_protocol_banner = b"ARTIQ sync_struct\n"


@unique
[docs]class ModAction(Enum): """Describes the type of incremental modification. `Mods` are represented by a dictionary ``m``. ``m["action"]`` describes the type of modification, as per this enum, serialized as a string if required. The path (member field) the change applies to is given in ``m["path"]`` as a list; elements give successive levels of indexing. (There is no ``path`` on initial initialization.) Details on the modification are stored in additional data fields specific to each type. For example, this represents appending the value ``42`` to an array ``data.counts[0]``: :: { "action": "append", "path": ["data", "counts", 0], "x": 42 } """ #: A full copy of the data is sent in `struct`; no `path` given. init = "init" #: Appends `x` to target list. append = "append" #: Inserts `x` into target list at index `i`. insert = "insert" #: Removes index `i` from target list. pop = "pop" #: Sets target's `key` to `value`. setitem = "setitem" #: Removes target's `key`. delitem = "delitem"
# Handlers to apply a given mod to a target dict, invoked with (target, mod). _mod_appliers = { ModAction.append: lambda t, m: t.append(m["x"]), ModAction.insert: lambda t, m: t.insert(m["i"], m["x"]), ModAction.pop: lambda t, m: t.pop(m["i"]), ModAction.setitem: lambda t, m: t.__setitem__(m["key"], m["value"]), ModAction.delitem: lambda t, m: t.__delitem__(m["key"]) }
[docs]def process_mod(target, mod): """Apply a *mod* to the target, mutating it.""" for key in mod["path"]: target = getitem(target, key) _mod_appliers[ModAction(mod["action"])](target, mod)
[docs]class Subscriber: """An asyncio-based client to connect to a ``Publisher``. :param notifier_name: Name of the notifier to subscribe to. :param target_builder: A function called during initialization that takes the object received from the publisher and returns the corresponding local structure to use. Can be identity. :param notify_cb: An optional function called every time a mod is received from the publisher. The mod is passed as parameter. The function is called after the mod has been processed. A list of functions may also be used, and they will be called in turn. :param disconnect_cb: An optional function called when disconnection happens from external causes (i.e. not when ``close`` is called). """ def __init__(self, notifier_name, target_builder, notify_cb=None, disconnect_cb=None): self.notifier_name = notifier_name self.target_builder = target_builder if notify_cb is None: notify_cb = [] if not isinstance(notify_cb, list): notify_cb = [notify_cb] self.notify_cbs = notify_cb self.disconnect_cb = disconnect_cb async def connect(self, host, port, before_receive_cb=None): self.reader, self.writer = \ await asyncio.open_connection(host, port, limit=100*1024*1024) try: if before_receive_cb is not None: before_receive_cb() self.writer.write(_protocol_banner) self.writer.write((self.notifier_name + "\n").encode()) self.receive_task = asyncio.ensure_future(self._receive_cr()) except: self.writer.close() del self.reader del self.writer raise async def close(self): self.disconnect_cb = None try: self.receive_task.cancel() try: await asyncio.wait_for(self.receive_task, None) except asyncio.CancelledError: pass finally: self.writer.close() del self.reader del self.writer async def _receive_cr(self): try: target = None while True: line = await self.reader.readline() if not line: return mod = pyon.decode(line.decode()) if mod["action"] == "init": target = self.target_builder(mod["struct"]) else: process_mod(target, mod) for notify_cb in self.notify_cbs: notify_cb(mod) finally: if self.disconnect_cb is not None: self.disconnect_cb()
[docs]class Notifier: """Encapsulates a structure whose changes need to be published. All mutations to the structure must be made through the :class:`.Notifier`. The original structure must only be accessed for reads. In addition to the list methods below, the :class:`.Notifier` supports the index syntax for modification and deletion of elements. Modification of nested structures can be also done using the index syntax, for example: >>> n = Notifier([]) >>> n.append([]) >>> n[0].append(42) >>> n.raw_view [[42]] This class does not perform any network I/O and is meant to be used with e.g. the :class:`.Publisher` for this purpose. Only one publisher at most can be associated with a :class:`.Notifier`. :param backing_struct: Structure to encapsulate. """ def __init__(self, backing_struct, root=None, path=[]): #: The raw data encapsulated (read-only!). self.raw_view = backing_struct if root is None: self.root = self self.publish = None else: self.root = root self._backing_struct = backing_struct self._path = path # Backing struct modification methods. # All modifications must go through them!
[docs] def append(self, x): """Append to a list.""" self._backing_struct.append(x) if self.root.publish is not None: self.root.publish({"action": ModAction.append.value, "path": self._path, "x": x})
[docs] def insert(self, i, x): """Insert an element into a list.""" self._backing_struct.insert(i, x) if self.root.publish is not None: self.root.publish({"action": ModAction.insert.value, "path": self._path, "i": i, "x": x})
[docs] def pop(self, i=-1): """Pop an element from a list. The returned element is not encapsulated in a :class:`.Notifier` and its mutations are no longer tracked.""" r = self._backing_struct.pop(i) if self.root.publish is not None: self.root.publish({"action": ModAction.pop.value, "path": self._path, "i": i}) return r
def __setitem__(self, key, value): self._backing_struct.__setitem__(key, value) if self.root.publish is not None: self.root.publish({"action": ModAction.setitem.value, "path": self._path, "key": key, "value": value}) def __delitem__(self, key): self._backing_struct.__delitem__(key) if self.root.publish is not None: self.root.publish({"action": ModAction.delitem.value, "path": self._path, "key": key}) def __getitem__(self, key): item = getitem(self._backing_struct, key) return Notifier(item, self.root, self._path + [key])
[docs]def update_from_dict(target, source): """Updates notifier contents from given source dictionary. Only the necessary changes are performed; unchanged fields are not written. (Currently, modifications are only performed at the top level. That is, whenever there is a change to a child array/struct the entire member is updated instead of choosing a more optimal set of mods.) """ curr = target.raw_view # Delete removed keys. for k in list(curr.keys()): if k not in source: del target[k] # Insert/update changed data. for k in source.keys(): if k not in curr or curr[k] != source[k]: target[k] = source[k]
[docs]class Publisher(AsyncioServer): """A network server that publish changes to structures encapsulated in a :class:`.Notifier`. :param notifiers: A dictionary containing the notifiers to associate with the :class:`.Publisher`. The keys of the dictionary are the names of the notifiers to be used with :class:`.Subscriber`. """ def __init__(self, notifiers): AsyncioServer.__init__(self) self.notifiers = notifiers self._recipients = {k: set() for k in notifiers.keys()} self._notifier_names = {id(v): k for k, v in notifiers.items()} for notifier in notifiers.values(): notifier.publish = partial(self.publish, notifier) async def _handle_connection_cr(self, reader, writer): try: line = await reader.readline() if line != _protocol_banner: return line = await reader.readline() if not line: return notifier_name = line.decode()[:-1] try: notifier = self.notifiers[notifier_name] except KeyError: return obj = {"action": ModAction.init.value, "struct": notifier.raw_view} line = pyon.encode(obj) + "\n" writer.write(line.encode()) queue = asyncio.Queue() self._recipients[notifier_name].add(queue) try: while True: line = await queue.get() writer.write(line) # raise exception on connection error await writer.drain() finally: self._recipients[notifier_name].remove(queue) except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError, TimeoutError): # subscribers disconnecting are a normal occurrence pass finally: writer.close() def publish(self, notifier, mod): line = pyon.encode(mod) + "\n" line = line.encode() notifier_name = self._notifier_names[id(notifier)] for recipient in self._recipients[notifier_name]: recipient.put_nowait(line)