forked from goodboy/tractor
				
			Add `tractor.wait_for_actor()` helper
Allows for waiting on another actor (by name) to register with the arbiter. This makes synchronized actor spawning and consecutive task coordination easier to accomplish from within sub-actors. Resolves #31wait_for_actor
							parent
							
								
									8e027ff571
								
							
						
					
					
						commit
						3f0c644768
					
				| 
						 | 
				
			
			@ -9,7 +9,7 @@ import trio
 | 
			
		|||
from .log import get_console_log, get_logger, get_loglevel
 | 
			
		||||
from ._ipc import _connect_chan, Channel
 | 
			
		||||
from ._actor import (
 | 
			
		||||
    Actor, _start_actor, Arbiter, get_arbiter, find_actor
 | 
			
		||||
    Actor, _start_actor, Arbiter, get_arbiter, find_actor, wait_for_actor
 | 
			
		||||
)
 | 
			
		||||
from ._trionics import open_nursery
 | 
			
		||||
from ._state import current_actor
 | 
			
		||||
| 
						 | 
				
			
			@ -17,8 +17,13 @@ from ._portal import RemoteActorError
 | 
			
		|||
 | 
			
		||||
 | 
			
		||||
__all__ = [
 | 
			
		||||
    'current_actor', 'find_actor', 'get_arbiter', 'open_nursery',
 | 
			
		||||
    'RemoteActorError', 'Channel',
 | 
			
		||||
    'current_actor',
 | 
			
		||||
    'find_actor',
 | 
			
		||||
    'get_arbiter',
 | 
			
		||||
    'wait_for_actor',
 | 
			
		||||
    'open_nursery',
 | 
			
		||||
    'RemoteActorError',
 | 
			
		||||
    'Channel',
 | 
			
		||||
]
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -14,8 +14,13 @@ from async_generator import asynccontextmanager, aclosing
 | 
			
		|||
 | 
			
		||||
from ._ipc import Channel, _connect_chan
 | 
			
		||||
from .log import get_console_log, get_logger
 | 
			
		||||
from ._portal import (Portal, open_portal, _do_handshake, LocalPortal,
 | 
			
		||||
                      maybe_open_nursery)
 | 
			
		||||
from ._portal import (
 | 
			
		||||
    Portal,
 | 
			
		||||
    open_portal,
 | 
			
		||||
    _do_handshake,
 | 
			
		||||
    LocalPortal,
 | 
			
		||||
    maybe_open_nursery
 | 
			
		||||
)
 | 
			
		||||
from . import _state
 | 
			
		||||
from ._state import current_actor
 | 
			
		||||
 | 
			
		||||
| 
						 | 
				
			
			@ -293,12 +298,12 @@ class Actor:
 | 
			
		|||
                cid = msg.get('cid')
 | 
			
		||||
                if cid:
 | 
			
		||||
                    if cid == 'internal':  # internal actor error
 | 
			
		||||
                        # import pdb; pdb.set_trace()
 | 
			
		||||
                        raise InternalActorError(
 | 
			
		||||
                            f"{chan.uid}\n" + msg['error'])
 | 
			
		||||
 | 
			
		||||
                    # deliver response to local caller/waiter
 | 
			
		||||
                    await self._push_result(chan.uid, cid, msg)
 | 
			
		||||
 | 
			
		||||
                    log.debug(
 | 
			
		||||
                        f"Waiting on next msg for {chan} from {chan.uid}")
 | 
			
		||||
                    continue
 | 
			
		||||
| 
						 | 
				
			
			@ -591,16 +596,44 @@ class Arbiter(Actor):
 | 
			
		|||
 | 
			
		||||
    def __init__(self, *args, **kwargs):
 | 
			
		||||
        self._registry = defaultdict(list)
 | 
			
		||||
        self._waiters = {}
 | 
			
		||||
        super().__init__(*args, **kwargs)
 | 
			
		||||
 | 
			
		||||
    def find_actor(self, name):
 | 
			
		||||
        for uid, actor in self._registry.items():
 | 
			
		||||
        for uid, sockaddr in self._registry.items():
 | 
			
		||||
            if name in uid:
 | 
			
		||||
                print('found it!')
 | 
			
		||||
                return actor
 | 
			
		||||
                return sockaddr
 | 
			
		||||
 | 
			
		||||
    async def wait_for_actor(self, name):
 | 
			
		||||
        """Wait for a particular actor to register.
 | 
			
		||||
 | 
			
		||||
        This is a blocking call if no actor by the provided name is currently
 | 
			
		||||
        registered.
 | 
			
		||||
        """
 | 
			
		||||
        sockaddrs = []
 | 
			
		||||
 | 
			
		||||
        for (aname, _), sockaddr in self._registry.items():
 | 
			
		||||
            if name == aname:
 | 
			
		||||
                sockaddrs.append(sockaddr)
 | 
			
		||||
 | 
			
		||||
        if not sockaddrs:
 | 
			
		||||
            waiter = trio.Event()
 | 
			
		||||
            self._waiters.setdefault(name, []).append(waiter)
 | 
			
		||||
            await waiter.wait()
 | 
			
		||||
            for uid in self._waiters[name]:
 | 
			
		||||
                sockaddrs.append(self._registry[uid])
 | 
			
		||||
 | 
			
		||||
        return sockaddrs
 | 
			
		||||
 | 
			
		||||
    def register_actor(self, uid, sockaddr):
 | 
			
		||||
        self._registry[uid].append(sockaddr)
 | 
			
		||||
        name, uuid = uid
 | 
			
		||||
        self._registry[uid] = sockaddr
 | 
			
		||||
 | 
			
		||||
        # pop and signal all waiter events
 | 
			
		||||
        events = self._waiters.pop(name, ())
 | 
			
		||||
        self._waiters.setdefault(name, []).append(uid)
 | 
			
		||||
        for event in events:
 | 
			
		||||
            event.set()
 | 
			
		||||
 | 
			
		||||
    def unregister_actor(self, uid):
 | 
			
		||||
        self._registry.pop(uid, None)
 | 
			
		||||
| 
						 | 
				
			
			@ -633,7 +666,7 @@ async def _start_actor(actor, main, host, port, arbiter_addr, nursery=None):
 | 
			
		|||
            result = await main()
 | 
			
		||||
            # XXX: If spawned with a dedicated "main function",
 | 
			
		||||
            # the actor is cancelled when this context is complete
 | 
			
		||||
            # given that there are no more active peer channels connected to it.
 | 
			
		||||
            # given that there are no more active peer channels connected
 | 
			
		||||
            actor.cancel_server()
 | 
			
		||||
 | 
			
		||||
    # block on actor to complete
 | 
			
		||||
| 
						 | 
				
			
			@ -675,17 +708,31 @@ async def find_actor(
 | 
			
		|||
    known to the arbiter.
 | 
			
		||||
    """
 | 
			
		||||
    actor = current_actor()
 | 
			
		||||
    if not actor:
 | 
			
		||||
        raise RuntimeError("No actor instance has been defined yet?")
 | 
			
		||||
 | 
			
		||||
    async with get_arbiter(*arbiter_sockaddr or actor._arb_addr) as arb_portal:
 | 
			
		||||
        sockaddrs = await arb_portal.run('self', 'find_actor', name=name)
 | 
			
		||||
        sockaddr = await arb_portal.run('self', 'find_actor', name=name)
 | 
			
		||||
        # TODO: return portals to all available actors - for now just
 | 
			
		||||
        # the last one that registered
 | 
			
		||||
        if sockaddrs:
 | 
			
		||||
            sockaddr = sockaddrs[-1]
 | 
			
		||||
        if sockaddr:
 | 
			
		||||
            async with _connect_chan(*sockaddr) as chan:
 | 
			
		||||
                async with open_portal(chan) as portal:
 | 
			
		||||
                    yield portal
 | 
			
		||||
        else:
 | 
			
		||||
            yield None
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@asynccontextmanager
 | 
			
		||||
async def wait_for_actor(
 | 
			
		||||
    name,
 | 
			
		||||
    arbiter_sockaddr=None,
 | 
			
		||||
):
 | 
			
		||||
    """Wait on an actor to register with the arbiter.
 | 
			
		||||
 | 
			
		||||
    A portal to the first actor which registered is be returned.
 | 
			
		||||
    """
 | 
			
		||||
    actor = current_actor()
 | 
			
		||||
    async with get_arbiter(*arbiter_sockaddr or actor._arb_addr) as arb_portal:
 | 
			
		||||
        sockaddrs = await arb_portal.run('self', 'wait_for_actor', name=name)
 | 
			
		||||
        sockaddr = sockaddrs[-1]
 | 
			
		||||
        async with _connect_chan(*sockaddr) as chan:
 | 
			
		||||
            async with open_portal(chan) as portal:
 | 
			
		||||
                yield portal
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -7,4 +7,6 @@ _current_actor = None
 | 
			
		|||
def current_actor() -> 'Actor':
 | 
			
		||||
    """Get the process-local actor instance.
 | 
			
		||||
    """
 | 
			
		||||
    if not _current_actor:
 | 
			
		||||
        raise RuntimeError("No actor instance has been defined yet?")
 | 
			
		||||
    return _current_actor
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
		Loading…
	
		Reference in New Issue