90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
import time
|
|
import random
|
|
|
|
try:
|
|
from Queue import PriorityQueue
|
|
except ImportError:
|
|
from queue import PriorityQueue
|
|
|
|
class ConnectionSelector(object):
|
|
" Base class for Selectors. "
|
|
def __init__(self, opts):
|
|
self.connection_opts = opts
|
|
|
|
|
|
class RoundRobinSelector(ConnectionSelector):
|
|
" Default selector using round-robin. "
|
|
def __init__(self, opts):
|
|
super(RoundRobinSelector, self).__init__(opts)
|
|
self.rr = -1
|
|
|
|
def select(self, connections):
|
|
self.rr += 1
|
|
self.rr %= len(connections)
|
|
return connections[self.rr]
|
|
|
|
|
|
class ConnectionPool(object):
|
|
def __init__(self, connections, dead_timeout=60, selector_class=RoundRobinSelector, randomize_hosts=True, **kwargs):
|
|
self.connections = [c for (c, opts) in connections]
|
|
self.dead = PriorityQueue(len(self.connections))
|
|
self.dead_count = {}
|
|
|
|
if randomize_hosts:
|
|
# randomize the connection list to avoid all clients hitting same node
|
|
# after startup/restart
|
|
random.shuffle(self.connections)
|
|
|
|
# default timeout after which to try resurrecting a connection
|
|
self.dead_timeout = dead_timeout
|
|
|
|
self.selector = selector_class(dict(connections))
|
|
|
|
def mark_dead(self, connection, dead_count, now=None):
|
|
# allow inject for testing purposes
|
|
now = now if now else time.time()
|
|
try:
|
|
self.connections.remove(connection)
|
|
except ValueError:
|
|
# connection not alive or another thread marked it already, ignore
|
|
return
|
|
else:
|
|
self.dead_count[connection] = dead_count
|
|
self.dead.put((now + self.dead_timeout * 2 ** (dead_count - 1), connection))
|
|
|
|
def mark_live(self, connection):
|
|
try:
|
|
del self.dead_count[connection]
|
|
except KeyError:
|
|
# race condition, safe to ignore
|
|
pass
|
|
|
|
def resurrect(self, force=False):
|
|
# no dead connections
|
|
if self.dead.empty():
|
|
return
|
|
|
|
# retrieve a connection to check
|
|
timeout, connection = self.dead.get()
|
|
|
|
if not force and timeout > time.time():
|
|
# return it back if not eligible and not forced
|
|
self.dead.put((timeout, connection))
|
|
return
|
|
|
|
# either we were forced or the connection is elligible to be retried
|
|
self.connections.append(connection)
|
|
|
|
def get_connection(self):
|
|
self.resurrect()
|
|
|
|
# no live nodes, resurrect one by force
|
|
if not self.connections:
|
|
self.resurrect(True)
|
|
|
|
connection = self.selector.select(self.connections)
|
|
|
|
return connection, self.dead_count.get(connection, 0)
|
|
|
|
|