244 lines
7.0 KiB
Python
244 lines
7.0 KiB
Python
from itertools import islice, cycle, groupby, repeat
|
|
from random import randint
|
|
from threading import RLock
|
|
|
|
from decoder import ConsistencyLevel
|
|
|
|
class HostDistance(object):
|
|
|
|
IGNORED = -1
|
|
LOCAL = 0
|
|
REMOTE = 1
|
|
|
|
class LoadBalancingPolicy(object):
|
|
|
|
def distance(self, host):
|
|
raise NotImplemented()
|
|
|
|
def make_query_plan(self, query):
|
|
raise NotImplemented()
|
|
|
|
|
|
class RoundRobinPolicy(LoadBalancingPolicy):
|
|
|
|
def populate(self, cluster, hosts):
|
|
self._live_hosts = set(hosts)
|
|
self._position = randint(0, len(hosts) - 1)
|
|
self._lock = RLock()
|
|
|
|
def distance(self, host):
|
|
return HostDistance.LOCAL
|
|
|
|
def make_query_plan(self, query):
|
|
with self._lock:
|
|
pos = self._position
|
|
self._position += 1
|
|
|
|
length = len(self._live_hosts)
|
|
pos %= length
|
|
return islice(cycle(self._live_hosts, pos, pos + length))
|
|
|
|
def on_up(self, host):
|
|
self._live_hosts.add(host)
|
|
|
|
def on_down(self, host):
|
|
self._live_hosts.discard(host)
|
|
|
|
def on_add(self, host):
|
|
self._live_hosts.add(host)
|
|
|
|
def on_remove(self, host):
|
|
self._live_hosts.remove(host)
|
|
|
|
|
|
class DCAwareRoundRobinPolicy(LoadBalancingPolicy):
|
|
|
|
def __init__(self, local_dc, used_hosts_per_remote_dc=0):
|
|
LoadBalancingPolicy.__init__(self)
|
|
self.local_dc = local_dc
|
|
self.used_hosts_per_remote_dc = used_hosts_per_remote_dc
|
|
self._dc_live_hosts = {}
|
|
|
|
def populate(self, cluster, hosts):
|
|
for dc, hosts in groupby(hosts, lambda h: h.dc):
|
|
self._dc_live_hosts[dc] = set(hosts)
|
|
|
|
self._position = randint(0, len(hosts) - 1)
|
|
|
|
def distance(self, host):
|
|
if host.dc == self.local_dc:
|
|
return HostDistance.LOCAL
|
|
|
|
if not self.used_hosts_per_remote_dc:
|
|
return HostDistance.IGNORE
|
|
else:
|
|
dc_hosts = self._dc_live_hosts.get(host.dc)
|
|
if not dc_hosts:
|
|
return HostDistance.IGNORE
|
|
|
|
if host in list(dc_hosts)[:self.used_hosts_per_remote_dc]:
|
|
return HostDistance.REMOTE
|
|
else:
|
|
return HostDistance.IGNORE
|
|
|
|
def make_query_plan(self, query):
|
|
with self._lock:
|
|
pos = self._position
|
|
self._position += 1
|
|
|
|
local_live = list(self._dc_live_hosts.get(self.local_dc))
|
|
pos %= len(local_live)
|
|
for host in islice(cycle(local_live, pos, pos + len(local_live))):
|
|
yield host
|
|
|
|
for dc, current_dc_hosts in self._dc_live_hosts.iteritems():
|
|
if dc == self.local_dc:
|
|
continue
|
|
|
|
for host in current_dc_hosts:
|
|
yield host
|
|
|
|
def on_up(self, host):
|
|
self._dc_live_hosts.setdefault(host.dc, set()).add(host)
|
|
|
|
def on_down(self, host):
|
|
self._dc_live_hosts.setdefault(host.dc, set()).discard(host)
|
|
|
|
def on_add(self, host):
|
|
self._dc_live_hosts.setdefault(host.dc, set()).add(host)
|
|
|
|
def on_remove(self, host):
|
|
self._dc_live_hosts.setdefault(host.dc, set()).discard(host)
|
|
|
|
|
|
class SimpleConvictionPolicy(object):
|
|
|
|
def __init__(self, host):
|
|
self.host = host
|
|
|
|
def add_failure(connection_exc):
|
|
return True
|
|
|
|
def reset(self):
|
|
pass
|
|
|
|
|
|
class ConstantReconnectionPolicy(object):
|
|
|
|
def __init__(self, delay):
|
|
if delay < 0:
|
|
raise ValueError("Delay may not be negative")
|
|
|
|
self.delay = delay
|
|
|
|
def new_schedule(self):
|
|
return repeat(self.delay)
|
|
|
|
|
|
class ExponentialReconnectionPolicy(object):
|
|
|
|
def __init__(self, base_delay, max_delay):
|
|
if base_delay < 0 or max_delay < 0:
|
|
raise ValueError("Delays may not be negative")
|
|
|
|
if max_delay < base_delay:
|
|
raise ValueError("Max delay must be greater than base delay")
|
|
|
|
self.base_delay = base_delay
|
|
self.max_delay = max_delay
|
|
|
|
def new_schedule(self):
|
|
return (min(self.base_delay * (i ** 2), self.max_delay) for i in xrange(64))
|
|
|
|
|
|
class WriteType(object):
|
|
|
|
SIMPLE = 0
|
|
BATCH = 1
|
|
UNLOGGED_BATCH = 2
|
|
COUNTER = 3
|
|
BATCH_LOG = 4
|
|
|
|
|
|
class RetryPolicy(object):
|
|
|
|
RETRY = 0
|
|
RETHROW = 1
|
|
IGNORE = 2
|
|
|
|
def on_read_timeout(self, query, consistency, required_responses,
|
|
received_responses, data_retrieved, attempt_num):
|
|
if attempt_num != 0:
|
|
return (self.RETHROW, None)
|
|
elif received_responses >= required_responses and not data_retrieved:
|
|
return (self.RETRY, consistency)
|
|
else:
|
|
return (self.RETHROW, None)
|
|
|
|
def on_write_timeout(self, query, consistency, write_type,
|
|
required_responses, received_responses, attempt_num):
|
|
if attempt_num != 0:
|
|
return (self.RETHROW, None)
|
|
elif write_type == WriteType.BATCH_LOG:
|
|
return (self.RETRY, consistency)
|
|
else:
|
|
return (self.RETHROW, None)
|
|
|
|
def on_unavailable(self, query, consistency, required_replicas, alive_replicas, attempt_num):
|
|
return (self.RETHROW, None)
|
|
|
|
|
|
class FallthroughRetryPolicy(RetryPolicy):
|
|
|
|
def on_read_timeout(self, *args, **kwargs):
|
|
return (self.RETHROW, None)
|
|
|
|
def on_write_timeout(self, *args, **kwargs):
|
|
return (self.RETHROW, None)
|
|
|
|
def on_unavailable(self, *args, **kwargs):
|
|
return (self.RETHROW, None)
|
|
|
|
|
|
class FallthroughRetryPolicy(RetryPolicy):
|
|
|
|
def _pick_consistency(self, num_responses):
|
|
if num_responses >= 3:
|
|
return (self.RETRY, ConsistencyLevel.name_to_value["THREE"])
|
|
elif num_responses >= 2:
|
|
return (self.RETRY, ConsistencyLevel.name_to_value["TWO"])
|
|
elif num_responses >= 1:
|
|
return (self.RETRY, ConsistencyLevel.name_to_value["ONE"])
|
|
else:
|
|
return (self.RETHROW, None)
|
|
|
|
def on_read_timeout(self, query, consistency, required_responses,
|
|
received_responses, data_retrieved, attempt_num):
|
|
if attempt_num != 0:
|
|
return (self.RETHROW, None)
|
|
elif received_responses < required_responses:
|
|
return self._pick_consistency(received_responses)
|
|
elif not data_retrieved:
|
|
return (self.RETRY, consistency)
|
|
else:
|
|
return (self.RETHROW, None)
|
|
|
|
def on_write_timeout(self, query, consistency, write_type,
|
|
required_responses, received_responses, attempt_num):
|
|
if attempt_num != 0:
|
|
return (self.RETHROW, None)
|
|
elif write_type in (WriteType.SIMPLE, WriteType.BATCH, WriteType.COUNTER):
|
|
return (self.IGNORE, None)
|
|
elif write_type == WriteType.UNLOGGED_BATCH:
|
|
return self._pick_consistency(received_responses)
|
|
elif write_type == WriteType.BATCH_LOG:
|
|
return (self.RETRY, consistency)
|
|
else:
|
|
return (self.RETHROW, None)
|
|
|
|
def on_unavailable(self, query, consistency, required_replicas, alive_replicas, attempt_num):
|
|
if attempt_num != 0:
|
|
return (self.RETHROW, None)
|
|
else:
|
|
return self._pick_consistency(alive_replicas)
|