def Open(self):
if self.__open_ar:
return self.__open_ar
self.__open_ar = AsyncResult()
self.__open_greenlet = gevent.spawn(self._OpenImpl)
return self.__open_ar
def _SendPingMessage(self):
"""Constructs and sends a Tping message."""
self._log.debug('Sending ping message.')
self._ping_ar = AsyncResult()
self._last_ping_start = time.time()
self._send_queue.put((self._ping_msg, self._EMPTY_DCT))
gevent.spawn(self._PingTimeoutHelper)
return self._ping_ar
def Open(self):
"""Open the underlying sink and increase the ref count."""
self._ref_count += 1
if self._ref_count > 1:
return AsyncResult.Complete()
def TryGet():
self._Get()
return True
# We don't want to link _Get directly as it'll hold a reference
# to the sink returned forever.
return AsyncResult.Run(TryGet)
def __init__(self, service, sink_provider, default_timeout, properties):
"""
Args:
service - The service interface class this dispatcher is serving.
sink_provider - An instance of a SinkProvider.
properties - The properties associated with this service and dispatcher.
"""
super(MessageDispatcher, self).__init__()
self.next_sink = sink_provider.CreateSink(properties)
self._dispatch_timeout = default_timeout
self._service = service
self._name = properties[SinkProperties.Label]
self._open_ar = AsyncResult()
def async_request(self, handler, *args, **kwargs):
"""Send an asynchronous request (does not wait for it to finish)
:returns: an :class:`rpyc.core.async.AsyncResult` object, which will
eventually hold the result (or exception)
"""
timeout = kwargs.pop("timeout", None)
if kwargs:
raise TypeError("got unexpected keyword argument(s) %s" % (list(kwargs.keys()),))
res = AsyncResult(weakref.proxy(self))
self._async_request(handler, args, res)
if timeout is not None:
res.set_expiry(timeout)
return res
def _TryExpandAperture(self, leave_pending=False):
"""Attempt to expand the aperture. By calling this it's assumed the aperture
needs to be expanded.
The aperture can be expanded if there are idle sinks available.
"""
endpoints = list(self._idle_endpoints)
added_node = None
new_endpoint = None
if endpoints:
new_endpoint = random.choice(endpoints)
self._idle_endpoints.discard(new_endpoint)
self._log.debug('Expanding aperture to include %s.' %
str(new_endpoint))
new_sink = self._servers[new_endpoint]
self._pending_endpoints.add(new_endpoint)
added_node = super(ApertureBalancerSink,
self)._AddSink(new_endpoint, new_sink)
self._UpdateSizeVarz()
if added_node:
if not leave_pending:
added_node.ContinueWith(
lambda ar: self._pending_endpoints.discard(new_endpoint))
return added_node, new_endpoint
else:
return AsyncResult.Complete(), None
def _SendPingMessage(self):
"""Constructs and sends a Tping message."""
self._log.debug('Sending ping message.')
self._ping_ar = AsyncResult()
self._last_ping_start = time.time()
self._send_queue.put((self._ping_msg, self._EMPTY_DCT))
gevent.spawn(self._PingTimeoutHelper)
return self._ping_ar
def _OpenInitialChannels(self):
"""Open the sink and all underlying nodes."""
self._open = True
if self._size > 0:
# Ignore the first sink, it's the FailingChannelSink.
AsyncResult.WhenAny([self._OpenNode(n) for n in self._heap[1:]])\
.ContinueWith(lambda _ar: self._OnOpenComplete())
else:
self._OnOpenComplete()
def _OnNodeDown(self, node):
"""Invoked by the base class when a node is marked down.
In this case, if the downed node is currently in the aperture, we want to
remove if, and then attempt to adjust the aperture.
"""
if node.channel.state != ChannelState.Idle:
ar, _ = self._TryExpandAperture()
return ar
else:
return AsyncResult.Complete()
def __init__(self, next_provider, sink_properties, global_properties):
super(HttpTransportSinkBase, self).__init__()
self._endpoint = global_properties[SinkProperties.Endpoint]
name = global_properties[SinkProperties.Label]
self._varz = self.Varz(
Source(service=name,
endpoint='%s:%d' %
(self._endpoint.host, self._endpoint.port)))
self._open_result = AsyncResult.Complete()
self._session = requests.Session()
self._raise_on_http_error = sink_properties.raise_on_http_error
def StaticDispatchMessage(sink, source, start_time, deadline, disp_msg):
# Init the properties dictionary w/ an empty endpoint
disp_msg.properties[MessageProperties.Endpoint] = None
if deadline:
disp_msg.properties[Deadline.KEY] = deadline
ar = AsyncResult()
sink_stack = ClientMessageSinkStack()
repsonse_sink = _AsyncResponseSink()
sink_stack.Push(repsonse_sink,
(source, start_time, ar, disp_msg.properties))
gevent.spawn(sink.AsyncProcessRequest, sink_stack, disp_msg, None, {})
return ar
def _AddSink(self, endpoint, sink_factory):
"""Add a sink to the heap.
The sink is immediately opened and initialized to Zero load.
Args:
sink - The sink that was just added.
"""
self._size += 1
self.__varz.size(self._size)
new_node = self.Node(sink_factory(), self.Idle, self._size, endpoint)
self._heap.append(new_node)
Heap.FixUp(self._heap, self._size)
# Adding an Open() in here allows us to optimistically assume it'll be opened
# before the next message attempts to get it. However, the Open() will likely
# yield, so other code paths need to be aware there is a potentially un-open
# sink on the heap.
if self._open:
return self._OpenNode(new_node)
else:
return AsyncResult.Complete()
class MessageDispatcher(ClientMessageSink):
"""Handles dispatching incoming and outgoing messages to a client sink stack."""
class Varz(VarzBase):
"""
dispatch_messages - The number of messages dispatched.
success_messages - The number of successful responses processed.
exception_messages - The number of exception responses processed.
request_latency - The average time taken to receive a response to a request.
"""
_VARZ_BASE_NAME = 'scales.MessageDispatcher'
_VARZ = {
'dispatch_messages': Rate,
'success_messages': Rate,
'exception_messages': Rate,
'request_latency': AverageTimer
}
def __init__(self, service, sink_provider, default_timeout, properties):
"""
Args:
service - The service interface class this dispatcher is serving.
sink_provider - An instance of a SinkProvider.
properties - The properties associated with this service and dispatcher.
"""
super(MessageDispatcher, self).__init__()
self.next_sink = sink_provider.CreateSink(properties)
self._dispatch_timeout = default_timeout
self._service = service
self._name = properties[SinkProperties.Label]
self._open_ar = AsyncResult()
def Open(self):
self._open_ar = self.next_sink.Open()
return self._open_ar
def Close(self):
self.next_sink.Close()
self._open_ar = None
def DispatchMethodCall(self, method, args, kwargs, timeout=None):
"""Creates and posts a Tdispatch message to a client sink stack.
Args:
service - The service interface originating this call.
method - The method being called.
args - The parameters passed to the method.
kwargs - The keyword parameters passed to the method.
timeout - An optional timeout. If not set, the global dispatch timeout
will be applied.
Returns:
An AsyncResult representing the status of the method call. This will be
signaled when either the call completes (successfully or from failure),
or after [timeout] seconds elapse.
"""
if not self._open_ar:
raise Exception('Dispatcher not open.')
timeout = timeout or self._dispatch_timeout
start_time = time.time()
if self._open_ar.ready():
return self._DispatchMethod(method, args, kwargs, timeout,
start_time)
else:
# _DispatchMethod returns an AsyncResult, so we end up with an
# AsyncResult<AsyncResult<TRet>>, Unwrap() removes one layer, yielding
# an AsyncResult<TRet>
return self._open_ar.ContinueWith(lambda ar: self._DispatchMethod(
method, args, kwargs, timeout, start_time)).Unwrap()
@staticmethod
def StaticDispatchMessage(sink, source, start_time, deadline, disp_msg):
# Init the properties dictionary w/ an empty endpoint
disp_msg.properties[MessageProperties.Endpoint] = None
if deadline:
disp_msg.properties[Deadline.KEY] = deadline
ar = AsyncResult()
sink_stack = ClientMessageSinkStack()
repsonse_sink = _AsyncResponseSink()
sink_stack.Push(repsonse_sink,
(source, start_time, ar, disp_msg.properties))
gevent.spawn(sink.AsyncProcessRequest, sink_stack, disp_msg, None, {})
return ar
def _DispatchMethod(self, method, args, kwargs, timeout, start_time):
open_time = time.time()
open_latency = open_time - start_time
if timeout:
# Calculate the deadline for this method call.
# Reduce it by the time it took for the open() to complete.
deadline = start_time + timeout - open_latency
else:
deadline = None
disp_msg = MethodCallMessage(self._service, method, args, kwargs)
source = Source(method=method, service=self._name)
self.Varz.dispatch_messages(source) # pylint: disable=no-member
return self.StaticDispatchMessage(self.next_sink, source, start_time,
deadline, disp_msg)
def AsyncProcessRequest(self, sink_stack, msg, stream, headers):
raise NotImplementedError("This should never be called.")
def AsyncProcessResponse(self, sink_stack, context, stream, msg):
raise NotImplementedError("This should never be called.")
def Open(self):
if self.next_sink:
return self.next_sink.Open()
else:
return AsyncResult.Complete()
class SocketTransportSink(MuxSocketTransportSink):
def __init__(self, socket, service):
self._ping_timeout = 5
self._ping_msg = self._BuildHeader(1, MessageType.Tping, 0)
self._last_ping_start = 0
super(SocketTransportSink, self).__init__(socket, service)
def _Init(self):
self._ping_ar = None
super(SocketTransportSink, self)._Init()
@staticmethod
def _EncodeTag(tag):
return [tag >> 16 & 0xff, tag >> 8 & 0xff, tag & 0xff] # Tag
def _BuildHeader(self, tag, msg_type, data_len):
total_len = 1 + 3 + data_len
return pack('!ibBBB',
total_len,
msg_type,
*self._EncodeTag(tag))
def _PingLoop(self):
"""Periodically pings the remote server."""
while self.isActive:
gevent.sleep(random.randint(30, 40))
if self.isActive:
self._SendPingMessage()
else:
break
def _SendPingMessage(self):
"""Constructs and sends a Tping message."""
self._log.debug('Sending ping message.')
self._ping_ar = AsyncResult()
self._last_ping_start = time.time()
self._send_queue.put((self._ping_msg, self._EMPTY_DCT))
gevent.spawn(self._PingTimeoutHelper)
return self._ping_ar
def _OnPingResponse(self, msg_type, stream):
"""Handles the response to a ping. On failure, shuts down the dispatcher.
"""
ar, self._ping_ar = self._ping_ar, None
if msg_type == MessageType.Rping:
ar.set()
ping_duration = time.time() - self._last_ping_start
self._log.debug('Got ping response in %d ms' % int(ping_duration * 1000))
else:
self._log.error('Unexpected response for tag 1 (msg_type was %d)' % msg_type)
ar.set_exception(Exception("Invalid ping response"))
def _PingTimeoutHelper(self):
ar = self._ping_ar
ar.wait(self._ping_timeout)
if not ar.successful():
ar.set_exception(Exception('Ping timed out'))
self._Shutdown('Ping Timeout')
def _CheckInitialConnection(self):
ar = self._SendPingMessage()
ar.get()
self._log.debug('Ping successful')
self._greenlets.append(self._SpawnNamedGreenlet('Ping Loop', self._PingLoop))
@staticmethod
def _CreateDiscardMessage(tag):
"""Create a Tdiscarded message for 'tag'
Args:
tag - The message tag to discard.
Returns
A (message, buffer, headers) tuple suitable for passing to AsyncProcessRequest.
"""
discard_message = MethodDiscardMessage(tag, 'Client timeout')
discard_message.which = tag
buf = StringIO()
headers = {}
MessageSerializer(None).Marshal(discard_message, buf, headers)
return discard_message, buf, headers
def _OnTimeout(self, tag):
if tag:
msg, buf, headers = self._CreateDiscardMessage(tag)
self.AsyncProcessRequest(None, msg, buf, headers)
def _ProcessReply(self, stream):
try:
msg_type, tag = ThriftMuxMessageSerializerSink.ReadHeader(stream)
if tag == 1 and msg_type == MessageType.Rping: #Ping
self._OnPingResponse(msg_type, stream)
elif tag != 0:
self._ProcessTaggedReply(tag, stream)
else:
self._log.error('Unexpected message, msg_type = %d, tag = %d' % (msg_type, tag))
except Exception:
self._log.exception('Exception processing reply message.')
def _Shutdown(self, reason, fault=True):
super(SocketTransportSink, self)._Shutdown(reason, fault)
if self._ping_ar:
self._ping_ar.set_exception(reason)
class LoadBalancerSink(ClientMessageSink):
"""Base class for all load balancer sinks."""
Role = SinkRole.LoadBalancer
def __init__(self, next_provider, sink_properties, global_properties):
self._properties = global_properties
service_name = global_properties[SinkProperties.Label]
server_set_provider = sink_properties.server_set_provider
log_name = self.__class__.__name__.replace('ChannelSink', '')
self._log = ROOT_LOG.getChild('%s.[%s]' % (log_name, service_name))
self.__init_done = Event()
self.__open_ar = None
self.__open_greenlet = None
self._server_set_provider = server_set_provider
self._endpoint_name = server_set_provider.endpoint_name
self._next_sink_provider = next_provider
self._state = ChannelState.Idle
self._servers = {}
super(LoadBalancerSink, self).__init__()
@property
def state(self):
return self._state
def Open(self):
if self.__open_ar:
return self.__open_ar
self.__open_ar = AsyncResult()
self.__open_greenlet = gevent.spawn(self._OpenImpl)
return self.__open_ar
def _OpenImpl(self):
while self._state != ChannelState.Closed:
try:
self._server_set_provider.Initialize(
self.__OnServerSetJoin,
self.__OnServerSetLeave)
server_set = self._server_set_provider.GetServers()
except gevent.GreenletExit:
return
except:
self._log.exception("Unable to initialize serverset, retrying in 5 seconds.")
gevent.sleep(5)
continue
random.shuffle(server_set)
self._servers = {}
[self.__AddServer(m) for m in server_set]
self.__init_done.set()
self._OpenInitialChannels()
self._open_greenlet = None
self._state = ChannelState.Open
return True
@abstractmethod
def _OpenInitialChannels(self):
"""To be overriden by subclasses. Called after the ServerSet is initialized
and the initial set of servers has been loaded.
"""
pass
def _OnOpenComplete(self):
"""To be called by subclasses when they've completed (successfully or not)
opening their sinks.
"""
self.__open_ar.set(True)
def WaitForOpenComplete(self, timeout=None):
self.__open_ar.wait(timeout)
def Close(self):
self._server_set_provider.Close()
self._state = ChannelState.Closed
self.__init_done.clear()
if self.__open_greenlet:
self.__open_greenlet.kill(block=False)
self.__open_greenlet = None
self.__open_ar = None
@abstractmethod
def _AsyncProcessRequestImpl(self, sink_stack, msg, stream, headers):
pass
def AsyncProcessRequest(self, sink_stack, msg, stream, headers):
if not self.__open_ar.ready():
def _on_open_done(_):
timeout_event = msg.properties.get(Deadline.EVENT_KEY, None)
# Either there is no timeout, or the timeout hasn't expired yet.
if not timeout_event or not timeout_event.Get():
self._AsyncProcessRequestImpl(sink_stack, msg, stream, headers)
self.__open_ar.rawlink(_on_open_done)
else:
self._AsyncProcessRequestImpl(sink_stack, msg, stream, headers)
def __GetEndpoint(self, instance):
if self._endpoint_name:
aeps = instance.additional_endpoints
ep = aeps.get(self._endpoint_name, None)
if not ep:
raise ValueError(
"Endpoint name %s not found in endpoints", self._endpoint_name)
return ep
else:
return instance.service_endpoint
def _OnServersChanged(self, endpoint, channel_factory, added):
"""Overridable by child classes. Invoked when servers in the server set are
added or removed.
Args:
instance - The server set member being added or removed.
added - True if the instance is being added, False if it's being removed.
"""
pass
def __AddServer(self, instance):
"""Adds a servers to the set of servers available to the load balancer.
Note: The new sink is not opened at this time.
Args:
instance - A Member object to be added to the pool.
"""
ep = self.__GetEndpoint(instance)
if not ep in self._servers:
new_props = self._properties.copy()
new_props.update({ SinkProperties.Endpoint: ep })
channel_factory = functools.partial(self._next_sink_provider.CreateSink, new_props)
self._servers[ep] = channel_factory
self._log.info("Instance %s joined (%d members)" % (
ep, len(self._servers)))
self._OnServersChanged(ep, channel_factory, True)
def __RemoveServer(self, instance):
"""Removes a server from the load balancer.
Args:
instance - A Member object to be removed from the pool.
"""
ep = self.__GetEndpoint(instance)
channel_factory = self._servers.pop(ep, None)
self._OnServersChanged(ep, channel_factory, False)
def __OnServerSetJoin(self, instance):
"""Invoked when an instance joins the server set.
Args:
instance - Instance added to the cluster.
"""
# callbacks from the ServerSet are delivered serially, so we can guarantee
# that once this unblocks, we'll still get the notifications delivered in
# the order that they arrived. Ex: OnJoin(a) -> OnLeave(a)
self.__init_done.wait()
# OnJoin notifications are delivered at startup, however we already
# pre-populate our copy of the ServerSet, so it's fine to ignore duplicates.
if self.__GetEndpoint(instance) in self._servers:
return
self.__AddServer(instance)
def __OnServerSetLeave(self, instance):
"""Invoked when an instance leaves the server set.
Args:
instance - Instance leaving the cluster.
"""
self.__init_done.wait()
self.__RemoveServer(instance)
self._log.info("Instance %s left (%d members)" % (
self.__GetEndpoint(instance), len(self._servers)))
class SocketTransportSink(MuxSocketTransportSink):
def __init__(self, socket, service):
self._ping_timeout = 5
self._ping_msg = self._BuildHeader(1, MessageType.Tping, 0)
self._last_ping_start = 0
super(SocketTransportSink, self).__init__(socket, service)
def _Init(self):
self._ping_ar = None
super(SocketTransportSink, self)._Init()
@staticmethod
def _EncodeTag(tag):
return [tag >> 16 & 0xff, tag >> 8 & 0xff, tag & 0xff] # Tag
def _BuildHeader(self, tag, msg_type, data_len):
total_len = 1 + 3 + data_len
return pack('!ibBBB', total_len, msg_type, *self._EncodeTag(tag))
def _PingLoop(self):
"""Periodically pings the remote server."""
while self.isActive:
gevent.sleep(random.randint(30, 40))
if self.isActive:
self._SendPingMessage()
else:
break
def _SendPingMessage(self):
"""Constructs and sends a Tping message."""
self._log.debug('Sending ping message.')
self._ping_ar = AsyncResult()
self._last_ping_start = time.time()
self._send_queue.put((self._ping_msg, self._EMPTY_DCT))
gevent.spawn(self._PingTimeoutHelper)
return self._ping_ar
def _OnPingResponse(self, msg_type, stream):
"""Handles the response to a ping. On failure, shuts down the dispatcher.
"""
ar, self._ping_ar = self._ping_ar, None
if msg_type == MessageType.Rping:
ar.set()
ping_duration = time.time() - self._last_ping_start
self._log.debug('Got ping response in %d ms' %
int(ping_duration * 1000))
else:
self._log.error('Unexpected response for tag 1 (msg_type was %d)' %
msg_type)
ar.set_exception(Exception("Invalid ping response"))
def _PingTimeoutHelper(self):
ar = self._ping_ar
ar.wait(self._ping_timeout)
if not ar.successful():
ar.set_exception(Exception('Ping timed out'))
self._Shutdown('Ping Timeout')
def _CheckInitialConnection(self):
ar = self._SendPingMessage()
ar.get()
self._log.debug('Ping successful')
self._greenlets.append(
self._SpawnNamedGreenlet('Ping Loop', self._PingLoop))
@staticmethod
def _CreateDiscardMessage(tag):
"""Create a Tdiscarded message for 'tag'
Args:
tag - The message tag to discard.
Returns
A (message, buffer, headers) tuple suitable for passing to AsyncProcessRequest.
"""
discard_message = MethodDiscardMessage(tag, 'Client timeout')
discard_message.which = tag
buf = StringIO()
headers = {}
MessageSerializer(None).Marshal(discard_message, buf, headers)
return discard_message, buf, headers
def _OnTimeout(self, tag):
if tag:
msg, buf, headers = self._CreateDiscardMessage(tag)
self.AsyncProcessRequest(None, msg, buf, headers)
def _ProcessReply(self, stream):
try:
msg_type, tag = ThriftMuxMessageSerializerSink.ReadHeader(stream)
if tag == 1 and msg_type == MessageType.Rping: #Ping
self._OnPingResponse(msg_type, stream)
elif tag != 0:
self._ProcessTaggedReply(tag, stream)
else:
self._log.error('Unexpected message, msg_type = %d, tag = %d' %
(msg_type, tag))
except Exception:
self._log.exception('Exception processing reply message.')
def _Shutdown(self, reason, fault=True):
super(SocketTransportSink, self)._Shutdown(reason, fault)
if self._ping_ar:
self._ping_ar.set_exception(reason)
def Open(self):
ar = AsyncResult()
ar.SafeLink(self._OpenImpl)
return ar
def _OnNodeDown(self, node):
return AsyncResult.Complete()