scales/asynchronous.py
class AsyncResult(g_AsyncResult):
  @staticmethod
  def WhenAll(ars):
    """Returns an AsyncResult representing the state of all AsyncResults passed.

    Args:
      ars - An enumerable of AsyncResults.
    Returns:
      An AsyncResult representing the completion of all ars passed in.  When all
      complete, the AsyncResult will be set to an array of the results of each
      AsyncResult, in the order they were enumerated in.
      If any AsyncResult fails, the return result will fail.
    """

    ret = AsyncResult()
    num_ars = len(ars)
    total = [num_ars]
    results = [None] * num_ars
    def complete(_n, _ar):
      if _ar.exception:
        ret.set_exception(_ar.exception)
      elif not ret.ready():
        total[0] -= 1
        results[_n] = _ar.value
        if total[0] == 0:
          ret.set(results)

    for n, ar in enumerate(ars):
      ar.rawlink(functools.partial(complete, n))
    return ret

  @staticmethod
  def WhenAny(ars):
    """Returns an AsyncResult representing the state of any AsyncResult passed in.
    The return value represents the state of the first AsyncResult to complete, or,
    if all fail, the last to fail.

    Args:
      ars - An enumerable of AsyncResults.
    Returns:
      An AsyncResult representing the state of the first AsyncResult to complete.
      The AsyncResult's value will be set to the value of the first result to
      complete, or, if all fail, the exception thrown by the last to fail.
    """
    ready_ars = [ar for ar in ars if ar.ready()]
    if ready_ars:
      return ready_ars[0]

    ret = AsyncResult()
    total = [len(ars)]
    def complete(_ar):
      total[0] -= 1
      if total[0] == 0 and _ar.exception:
        ret.set_exception(_ar.exception)
      elif not ret.ready() and _ar.successful():
        ret.set(_ar.value)

    for ar in ars:
      ar.rawlink(complete)
    return ret

  @staticmethod
  def FromValue(val):
    if val is None:
      return AsyncResult.Complete()
    else:
      ar = AsyncResult()
      ar.set(val)
      return ar

  @staticmethod
  def Complete():
    """Return an AsyncResult that has completed."""
    return _COMPLETE

  @staticmethod
  def CompleteIn(n):
    """Returns an AsyncResult that completes in <n> seconds

    Args:
      n - The number of seconds to wait before completing.
    """
    ar = AsyncResult()
    def helper():
      ar.set()
    g = Greenlet(helper)
    g.start_later(float(n))
    return ar

  def _SafeLinkHelper(self, fn):
    try:
      self.set(fn())
    except:
      self.set_exception(sys.exc_info()[1])

  def SafeLink(self, fn):
    """Propagate the result of calling fn() on a new greenlet to ar

    Args:
      ar - An AsyncResult.
      fn - The function to execute.
    """
    gevent.spawn(self._SafeLinkHelper, fn)

  def ContinueWith(self, fn, on_hub=True):
    cw_ar = AsyncResult()
    def continue_with_callback(_ar):
      def run():
        try:
          val = fn(_ar)
          cw_ar.set(val)
        except:
          cw_ar.set_exception(sys.exc_info()[1])
      if on_hub:
        run()
      else:
        gevent.spawn(run)
    self.rawlink(continue_with_callback)
    return cw_ar

  def Map(self, fn):
    def mapper(_):
      if self.exception:
        return self
      else:
        return fn(self.value)
    return self.ContinueWith(mapper).Unwrap()

  def _UnwrapHelper(self, target):
    if self.ready():
      # We're ready, propagate the result
      if self.exception:
        target.set_exception(self.exception)
      else:
        if isinstance(self.value, AsyncResult):
          self.value._UnwrapHelper(target)
        else:
          target.set(self.value)
    else:
      self.rawlink(
        functools.partial(AsyncResult._UnwrapHelper, target=target))

  def Unwrap(self):
    unwrapped_ar = AsyncResult()
    self._UnwrapHelper(unwrapped_ar)
    return unwrapped_ar

  @staticmethod
  def TryGet(val):
    if isinstance(val, AsyncResult):
      return val.get()
    else:
      return val

  @staticmethod
  def Run(fn):
    ar = AsyncResult()
    ar.SafeLink(fn)
    return ar

  @staticmethod
  def RunInline(fn):
    ar = AsyncResult()
    ar._SafeLinkHelper(fn)
    return ar
scales/constants.py
class Int(object):
  MaxValue = 2147483647
  MinValue = -2147483648
scales/constants.py
class SinkProperties(object):
  Endpoint = 'endpoint'
  ServiceInterface = 'service_iface'
  Label = 'label'
scales/loadbalancer/base.py
class NoMembersError(Exception): pass
scales/sink.py
class FailingMessageSink(ClientMessageSink):
  """A sink that always returns a failure message."""

  def __init__(self, ex):
    self._ex = ex
    super(FailingMessageSink, self).__init__()

  def AsyncProcessRequest(self, sink_stack, msg, stream, headers):
    msg = MethodReturnMessage(error=self._ex())
    sink_stack.AsyncProcessResponseMessage(msg)

  def AsyncProcessResponse(self, sink_stack, context, stream, msg):
    raise NotImplementedError("This should never be called")

  @property
  def state(self):
    return ChannelState.Open

  @property
  def endpoint(self):
    return None
scales/constants.py
class ChannelState(object):
  Idle = 1
  Open = 2
  Busy = 3
  Closed = 4
scales/sink.py
def SinkProvider(sink_cls, role=None, **defaults):
  """Factory for creating simple sink providers.

  Args:
    sink_cls - The type of sink to provide.
  Returns:
    A SinkProvider that provides sinks of type 'sink_cls'.
  """
  field_names = ' '.join(defaults.keys())
  params_cls = namedtuple('Params', field_names)

  def CreateSink(self, properties):
    return self.SINK_CLASS(self.next_provider, self.sink_properties, properties)

  def sink_class(self):
    return self.SINK_CLASS

  provider = type(
    sink_cls.__name__ + 'Provider',
    (SinkProviderBase, ),
    {
      'SINK_CLASS': sink_cls,
      'PARAMS_CLASS': params_cls,
      'Role': role,
      'CreateSink': CreateSink,
      'sink_class': property(sink_class),
      '_defaults': defaults
    }
  )
  return provider
scales/varz.py
class VarzType(object):
class Source(object):
class VarzMetric(object):
class Gauge(VarzMetric): VARZ_TYPE = VarzType.Gauge
class Rate(VarzMetric): VARZ_TYPE = VarzType.Rate
class AverageRate(VarzMetric): VARZ_TYPE = VarzType.AverageRate
class Counter(Rate): VARZ_TYPE = VarzType.Counter
class VarzTimerBase(VarzMetric):
class AverageTimer(VarzTimerBase): VARZ_TYPE = VarzType.AverageTimer
class AggregateTimer(VarzTimerBase): VARZ_TYPE = VarzType.AggregateTimer
class VarzMeta(type):
class _VarzBase(object):
class _SampleSet(object):
class VarzReceiver(object):
class VarzAggregator(object):
  class _Agg(object):
class VarzSocketWrapper(object):
  class Varz(VarzBase):
class MonoClock(object):
class Ema(object):
  def __init__(self, method=None, service=None, endpoint=None, client_id=None):
  def to_tuple(self):
  def to_dict(self):
  def __cmp__(self, other):
  def __hash__(self):
  def _Adapt(fn):
    def __Adapt(metric, source, amount=1):
  def __init__(self, metric, source):
  def __call__(self, *args):
  def ForSource(self, source):
  def Measure(self, source=None):
  def __new__(mcs, name, bases, dct):
def VerifySource(source):
  def __init__(self, source):
  def __getattr__(self, item):
  def __init__(self, max_size, data=None, p=.1):
  def Sample(self, value):
  def RegisterMetric(metric, varz_type):
  def IncrementVarz(source, metric, amount=1):
  def SetVarz(source, metric, value):
  def RecordPercentileSample(cls, source, metric, value):
def DefaultKeySelector(k):
    def __init__(self):
  def CalculatePercentile(values, pct):
  def _Downsample(lst, target_size):
  def Aggregate(varz, metrics, key_selector=None):
  def __init__(self, socket, varz_tag):
  def host(self):
  def port(self):
  def isOpen(self):
  def read(self, sz):
  def recv_into(self, buf, sz):
  def flush(self):
  def write(self, buff):
  def open(self):
  def close(self):
  def readAll(self, sz):
  def __init__(self):
  def Sample(self):
  def __init__(self, window):
  def Update(self, ts, sample):
  VARZ_TYPE = None
class Gauge(VarzMetric): VARZ_TYPE = VarzType.Gauge
class Rate(VarzMetric): VARZ_TYPE = VarzType.Rate
class AverageRate(VarzMetric): VARZ_TYPE = VarzType.AverageRate
class Counter(Rate): VARZ_TYPE = VarzType.Counter
class AverageTimer(VarzTimerBase): VARZ_TYPE = VarzType.AverageTimer
class AggregateTimer(VarzTimerBase): VARZ_TYPE = VarzType.AggregateTimer
  _VARZ = {}
  _VARZ_BASE_NAME = None
  VARZ_METRICS = {}
  VARZ_DATA = defaultdict(lambda: defaultdict(int))
  VARZ_PERCENTILES = [.5, .90, .99, .999, .9999]
  _MAX_PERCENTILE_SIZE = 1000
  MAX_AGG_AGE = 5 * 60
    _VARZ_BASE_NAME = 'scales.socket'
    _VARZ = {
      'bytes_recv': Rate,
      'bytes_sent': Rate,
      'num_connections': Counter,
      'tests_failed': Counter,
      'connects': Rate,
      'open_latency': AverageTimer
    }
scales/constants.py
class SinkRole(object):
  Transport = 'transport'
  Pool = 'pool'
  LoadBalancer = 'loadbalancer'
  Formatter = 'formatter'
scales/loadbalancer/base.py
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)))
scales/constants.py
class MessageProperties(object):
  Endpoint = '__Endpoint'


Question: import functools
import random
  from gevent.lock import RLock # pylint: disable=E0611
  from gevent.coros import RLock
from .base import (
  LoadBalancerSink,
  NoMembersError
)
from ..asynchronous import AsyncResult
from ..constants import (
  ChannelState,
  Int,
  MessageProperties,
  SinkProperties,
  SinkRole
)
from ..sink import (
  FailingMessageSink,
  SinkProvider,
)
from ..varz import (
  Counter,
  Gauge,
  Source,
  VarzBase
)
"""Heap Load Balancer.

Based on the heap balancer from finagle, see (https://github.com/twitter/finagle/blob/master/finagle-core/src/main/scala/com/twitter/finagle/loadbalancer/HeapBalancer.scala)

The heap load balancer maintains all nodes in a min-heap.  The heap is adjusted
as the load on a node increases or decreases.  Nodes start at Zero (min int) and
increase load as they are used.  Load decreases when they are released back into
the pool.

Downed nodes are tracked by setting a node's load to > 0.  A linked list of downed
nodes is kept to resurrect downed nodes if they become active again.
"""


try:
except ImportError:


class Heap(object):
  """A utility class to perform heap functions"""
  @staticmethod
  def Swap(heap, i, j):
    """Swap two elements in the heap.

    Args:
      heap - The heap array.
      i, j - The indexes int the array to swap.
    """
    if i == 0 or j == 0:
      raise Exception("heap swapping of element 0 should never happen.")

    heap[i], heap[j] = heap[j], heap[i]
    heap[i].index = i
    heap[j].index = j

  @staticmethod
  def FixUp(heap, i):
    """Traverse up the heap, ensuring the invariant is maintained.

    Args:
      heap - The heap array.
      i - The index to start at.
    """
    while True:
      if i != 1 and heap[i] < heap[i//2]:
        Heap.Swap(heap, i, i//2)
        i //= 2 # FixUp(heap, i/2)
      else:
        break

  @staticmethod
  def FixDown(heap, i, j):
    """Traverse down the heap, ensuring the invariant is maintained.

    Args:
      heap - The heap array.
      i, j - The node index to traverse from -> to.
    """
    while True:
      if j < i * 2: break

      m = 2 * i if (j == i * 2 or heap[2*i] < heap[2*i+1]) else 2*i+1
      if heap[m] < heap[i]:
        Heap.Swap(heap, i, m)
        i = m  # FixDown(heap, m, j)
      else:
        break


def synchronized(fn):
  """Runs the wrapped method under a lock.
  The self parameter to the wrapped function is expected to have a __heap_lock
  attribute.
  """
  @functools.wraps(fn)
  def wrapper(self, *args, **kwargs):
    with self._heap_lock:
      return fn(self, *args, **kwargs)
  return wrapper


class HeapBalancerSink(LoadBalancerSink):
  """A sink that implements a heap load balancer."""
  Penalty = Int.MaxValue
  Idle = Int.MinValue + 1

  class HeapVarz(VarzBase):
    """
    size - The number of nodes in the pool
    no_members - The number of times the balancer served a failing requests
                 because there were no members in the pool.
    """
    _VARZ_BASE_NAME = 'scales.loadbalancer.Heap'
    _VARZ = {
      'size': Gauge,

Answer:       'no_members': Counter