jc211
/
tractor
forked from goodboy/tractor
1
0
Fork 0
Code Pull Requests Activity

Cast `defaultdict` to `dict` for registry get

optional_msgspec_support
Tyler Goodlet 2021-09-06 11:41:34 -04:00
parent 1ab495a64d
commit 486e983964
1 changed files with 18 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
tractor/_actor.py
View File

@ -800,8 +800,9 @@ class Actor:
# XXX: msgspec doesn't support serializing tuples # XXX: msgspec doesn't support serializing tuples
# so just cash manually here since it's what our # so just cash manually here since it's what our
# internals expect. # internals expect.
address: Tuple[str, int] = tuple(value) if value else value self._arb_addr: Tuple[str, int] = (
self._arb_addr = address tuple(value) if value else value
)
else: else:
setattr(self, attr, value) setattr(self, attr, value)
@ -1206,8 +1207,10 @@ class Arbiter(Actor):
is_arbiter = True is_arbiter = True
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
self._registry = defaultdict(list) self._registry = defaultdict(list)
self._waiters = {} self._waiters = {}
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
async def find_actor(self, name: str) -> Optional[Tuple[str, int]]: async def find_actor(self, name: str) -> Optional[Tuple[str, int]]:
@ -1220,22 +1223,25 @@ class Arbiter(Actor):
async def get_registry( async def get_registry(
self self
) -> Dict[str, Tuple[str, str]]: ) -> Dict[str, Tuple[str, str]]:
"""Return current name registry. '''Return current name registry.
"""
This method is async to allow for cross-actor invocation.
'''
# NOTE: requires ``strict_map_key=False`` to the msgpack # NOTE: requires ``strict_map_key=False`` to the msgpack
# unpacker since we have tuples as keys (not this makes the # unpacker since we have tuples as keys (not this makes the
# arbiter suscetible to hashdos): # arbiter suscetible to hashdos):
# https://github.com/msgpack/msgpack-python#major-breaking-changes-in-msgpack-10 # https://github.com/msgpack/msgpack-python#major-breaking-changes-in-msgpack-10
return self._registry return dict(self._registry)
async def wait_for_actor( async def wait_for_actor(
self, name: str self,
name: str,
) -> List[Tuple[str, int]]: ) -> List[Tuple[str, int]]:
"""Wait for a particular actor to register. '''Wait for a particular actor to register.
This is a blocking call if no actor by the provided name is currently This is a blocking call if no actor by the provided name is currently
registered. registered.
""" '''
sockaddrs = [] sockaddrs = []
for (aname, _), sockaddr in self._registry.items(): for (aname, _), sockaddr in self._registry.items():
@ -1267,5 +1273,8 @@ class Arbiter(Actor):
if isinstance(event, trio.Event): if isinstance(event, trio.Event):
event.set() event.set()
async def unregister_actor(self, uid: Tuple[str, str]) -> None: async def unregister_actor(
self,
uid: Tuple[str, str]
) -> None:
self._registry.pop(tuple(uid)) self._registry.pop(tuple(uid))