"""
Implementation of proxy-related features.
"""
import os
import sys
import time
import Pyro4
from Pyro4.errors import ConnectionClosedError
from Pyro4.errors import NamingError
from Pyro4.message import FLAGS_ONEWAY
from . import config
from .address import SocketAddress
from .address import address_to_host_port
def locate_ns(nsaddr, timeout=3.0):
"""
Locate a name server to ensure it actually exists.
Parameters
----------
nsaddr : SocketAddress
The address where the name server should be up and running.
timeout : float
Timeout in seconds before aborting location.
Returns
-------
nsaddr
The address where the name server was located.
Raises
------
NamingError
If the name server could not be located.
"""
host, port = address_to_host_port(nsaddr)
time0 = time.time()
while True:
try:
Pyro4.locateNS(host, port)
return nsaddr
except NamingError:
if time.time() - time0 < timeout:
time.sleep(0.1)
continue
raise TimeoutError('Could not locate the name server!')
class Proxy(Pyro4.core.Proxy):
"""
A proxy to access remote agents.
Parameters
----------
name : str
Proxy name, as registered in the name server.
nsaddr : SocketAddress, str
Name server address.
timeout : float
Timeout, in seconds, to wait until the agent is discovered.
safe : bool, default is None
Use safe calls by default. When not set, osbrain default's
:py:data:`osbrain.config['SAFE']` is used.
"""
def __init__(self, name, nsaddr=None, timeout=3.0, safe=None):
if not nsaddr:
nsaddr = os.environ.get('OSBRAIN_NAMESERVER_ADDRESS')
nshost, nsport = address_to_host_port(nsaddr)
# Make sure name server exists
locate_ns(nsaddr)
time0 = time.time()
super().__init__('PYRONAME:%s@%s:%s' % (name, nshost, nsport))
if safe is not None:
self._default_safe = safe
else:
self._default_safe = config['SAFE']
self._safe = self._default_safe
self._next_oneway = False
while not self._ready_or_timeout(time0, timeout):
continue
def _ready_or_timeout(self, time0, timeout):
"""
Check if the proxy is ready or raise after a timeout.
Parameters
----------
time0 : float
Timestamp (in seconds) to take as the initial time.
timeout : float
Time (in seconds) allowed after `time0` before raising an
exception.
"""
try:
self.unsafe.ping()
except Exception:
time.sleep(0.1)
if time.time() - time0 < timeout:
return False
raise
return True
def wait_for_running(self, timeout=3.0):
"""
Wait until the agent is running.
Parameters
----------
timeout : float
Raise and exception if the agent is not running after this number
of seconds. Use a negative value to wait forever.
Raises
------
TimeoutError
If the agent is not running after the given timeout.
Returns
-------
Proxy
The object itself.
"""
time0 = time.time()
while not self.is_running():
if timeout >= 0 and time.time() - time0 > timeout:
msg = 'Timed out while waiting for the agent to be running'
raise TimeoutError(msg)
time.sleep(0.01)
return self
def __getstate__(self):
return super().__getstate__() + (
self._next_oneway,
self._default_safe,
self._safe,
)
def __setstate__(self, state):
super().__setstate__(state[:-3])
self._next_oneway = state[-3]
self._default_safe = state[-2]
self._safe = state[-1]
def __setattr__(self, name, value):
if name in ('_safe', '_default_safe', '_next_oneway'):
return super(Pyro4.core.Proxy, self).__setattr__(name, value)
if name.startswith('_'):
return super().__setattr__(name, value)
kwargs = {name: value}
return self.set_attr(**kwargs)
def __getattr__(self, name):
if name in self._pyroAttrs:
return self.get_attr(name)
return super().__getattr__(name)
def release(self):
"""
Release the connection to the Pyro daemon.
"""
self._pyroRelease()
def nsaddr(self):
"""
Get the socket address of the name server.
Returns
-------
SocketAddress
The socket address.
"""
return SocketAddress(self._pyroUri.host, self._pyroUri.port)
@property
def safe(self):
"""
Make the next remote method call be safe.
Returns
-------
The proxy itself.
"""
self._safe = True
return self
@property
def unsafe(self):
"""
Make the next remote method call be unsafe.
Returns
-------
The proxy itself.
"""
self._safe = False
return self
@property
def oneway(self):
"""
Make the next remote method call be one way.
Returns
-------
The proxy itself.
"""
self._next_oneway = True
return self
def _pyroInvoke( # noqa: N802
self, methodname, args, kwargs, flags=0, objectId=None # noqa: N803
):
"""
Wrapper around `_remote_call` to safely execute methods on remote
objects.
"""
try:
result = self._remote_call(
methodname, args, kwargs, flags, objectId
)
except Exception:
sys.stdout.write(''.join(Pyro4.util.getPyroTraceback()))
sys.stdout.flush()
raise
finally:
self._safe = self._default_safe
self._next_oneway = False
self._post_invoke(methodname, args, kwargs)
return result
def _is_safe_method(self, methodname):
"""
Check if a remote method can be called safely.
Parameters
----------
methodname : str
The name of the method to evaluate.
Returns
-------
bool
Whether the method can be safely called.
"""
return (
methodname in self._pyroMethods
and not methodname.startswith('_')
and methodname
not in (
'run',
'get_attr',
'kill',
'safe_call',
'concurrent',
'is_running',
)
)
def _remote_call(
self, methodname, args, kwargs, flags, objectId # noqa: N803
):
"""
Call a remote method from the proxy.
"""
if self._next_oneway:
flags |= FLAGS_ONEWAY
result = super()._pyroInvoke(
methodname, args, kwargs, flags=flags, objectId=objectId
)
return result
if self._safe and self._is_safe_method(methodname):
safe_args = [methodname] + list(args)
result = super()._pyroInvoke(
'safe_call', safe_args, kwargs, flags=flags, objectId=objectId
)
if isinstance(result, Exception):
raise result
else:
result = super()._pyroInvoke(
methodname, args, kwargs, flags=flags, objectId=objectId
)
return result
def _post_invoke(self, methodname, args, kwargs):
"""
After invoking a call, check if the proxy must be modified.
This could happen if the `set_method` or `set_attr` have been invoked.
In that case, the method(s) or attribute(s) are added to the proxy's
available method(s)/attributes(s).
"""
if methodname == 'set_method':
self._set_new_available_methods(args, kwargs)
elif methodname == 'set_attr':
self._set_new_available_attributes(kwargs)
def _set_new_available_methods(self, args, kwargs):
"""
Set new methods available from the proxy.
Parameters
----------
args : list
A list of new methods to be made available from the proxy.
kwargs : dict
A dictionary with the methods' names and their values.
"""
for method in args:
self._pyroMethods.add(method.__name__)
for name, _ in kwargs.items():
self._pyroMethods.add(name)
def _set_new_available_attributes(self, kwargs):
"""
Set new attributes available from the proxy.
Parameters
----------
kwargs : dict
A dictionary with the attributes' names and their values.
"""
for name in kwargs:
self._pyroAttrs.add(name)
class NSProxy(Pyro4.core.Proxy):
"""
A proxy to access a name server.
Parameters
----------
nsaddr : SocketAddress, str
Name server address.
timeout : float
Timeout, in seconds, to wait until the name server is discovered.
"""
def __init__(self, nsaddr=None, timeout=3):
if not nsaddr:
nsaddr = os.environ.get('OSBRAIN_NAMESERVER_ADDRESS')
nshost, nsport = address_to_host_port(nsaddr)
# Make sure name server exists
locate_ns(nsaddr, timeout)
ns_name = Pyro4.constants.NAMESERVER_NAME
super().__init__('PYRONAME:%s@%s:%d' % (ns_name, nshost, nsport))
def release(self):
"""
Release the connection to the Pyro daemon.
"""
self._pyroRelease()
def proxy(self, name, timeout=3.0):
"""
Get a proxy to access an agent registered in the name server.
Parameters
----------
name : str
Proxy name, as registered in the name server.
timeout : float
Timeout, in seconds, to wait until the agent is discovered.
Returns
-------
Proxy
A proxy to access an agent registered in the name server.
"""
return Proxy(name, nsaddr=self.addr(), timeout=timeout)
def addr(self, agent_alias=None, address_alias=None):
"""
Return the name server address or the address of an agent's socket.
Parameters
----------
agent_alias : str, default is None
The alias of the agent to retrieve its socket address.
address_alias : str, default is None
The alias of the socket address to retrieve from the agent.
Returns
-------
SocketAddress or AgentAddress
The name server or agent's socket address.
"""
if not agent_alias and not address_alias:
return SocketAddress(self._pyroUri.host, self._pyroUri.port)
agent = self.proxy(agent_alias)
addr = agent.addr(address_alias)
agent.release()
return addr
def shutdown_agents(self, timeout=10.0):
"""
Shutdown all agents registered in the name server.
Parameters
----------
timeout : float, default is 10.
Timeout, in seconds, to wait for the agents to shutdown.
"""
# Wait for all agents to be shutdown (unregistered)
time0 = time.time()
super()._pyroInvoke('async_shutdown_agents', (self.addr(),), {})
while time.time() - time0 <= timeout / 2.0:
if not len(self.agents()):
return
time.sleep(0.1)
super()._pyroInvoke('async_kill_agents', (self.addr(),), {})
while time.time() - time0 <= timeout:
if not len(self.agents()):
return
time.sleep(0.1)
raise TimeoutError(
'Chances are {} were not shutdown after {} s!'.format(
self.agents(), timeout
)
)
def shutdown(self, timeout=10.0):
"""
Shutdown the name server. All agents will be shutdown as well.
Parameters
----------
timeout : float, default is 10.
Timeout, in seconds, to wait for the agents to shutdown.
"""
self.shutdown_agents(timeout)
try:
super()._pyroInvoke('daemon_shutdown', (), {}, flags=0)
except ConnectionClosedError:
pass
