def __init__(self, endpoints, client_cls, size=DEFAULT_SIZE,
ttransport=TTransport.TBufferedTransport,
tprotocol=TBinaryProtocol.TBinaryProtocol,
use_ssl=False, ca_certs=None, cert=None, key=None):
"""
Initialize the thrift connection pool for a specific server and client
"""
self.endpoints = endpoints
self.client_cls = client_cls
self.ttransport = ttransport
self.tprotocol = tprotocol
self.use_ssl = use_ssl
self.ca_certs = ca_certs
self.cert = cert
self.key = key
self._closed = False
self.size = size
self._semaphore = threading.BoundedSemaphore(size)
self._connection_queue = LifoQueueWithEvict(size)
class ThriftConnectionPool(object):
"""
Thrift Conncetion pool
Attributes:
endpoints: list of endpoints that thrift server is listening on
client_cls: thrift client class to instantiate
size: optional size of the pool, equivalent to max number of
connections. Defaults to ThriftConnectionPool.DEFAULT_SIZE
ttransport: optional thrift transport to use, defaults to
TTransport.TBufferedTransport
tprotocol: optional thrift protocol to use, defaults to
TBinaryProtocol.TBinaryProtocol
ca_certs, cert, key: files needed by SSL
"""
DEFAULT_SIZE = 5
def __init__(self, endpoints, client_cls, size=DEFAULT_SIZE,
ttransport=TTransport.TBufferedTransport,
tprotocol=TBinaryProtocol.TBinaryProtocol,
use_ssl=False, ca_certs=None, cert=None, key=None):
"""
Initialize the thrift connection pool for a specific server and client
"""
self.endpoints = endpoints
self.client_cls = client_cls
self.ttransport = ttransport
self.tprotocol = tprotocol
self.use_ssl = use_ssl
self.ca_certs = ca_certs
self.cert = cert
self.key = key
self._closed = False
self.size = size
self._semaphore = threading.BoundedSemaphore(size)
self._connection_queue = LifoQueueWithEvict(size)
def evict_check(self, idle_threshold_millis):
self._connection_queue.evict(idle_threshold_millis, self._close_connection)
def close(self):
"""
Closes the connection pool. Outstanding connections will be closed on
the next call to return_connection
"""
self._closed = True
while not self._connection_queue.empty():
try:
(conn, t) = self._connection_queue.get(block=False)
self._connection_queue.task_done()
try:
self._close_connection(conn)
except:
# ignore thrift errors, nothing we can do here
pass
except Queue.Empty:
# ignore empty queue - that's what we want anyway
pass
def _create_connection(self):
"""
Helper that creates a thrift client connection to the instance host
and port.
Returns:
An instance of the thrift client class with an open transport
"""
endpoints = self.endpoints[:]
shuffle(endpoints)
for endpoint in endpoints:
try:
host, port = endpoint.split(':')
logger.debug('connecting to %s:%s. ssl? %s', host, port, self.use_ssl)
if self.use_ssl:
socket = TSSLSocket(host=host, port=int(port),
ca_certs=self.ca_certs, cert=self.cert, key=self.key)
else:
socket = TSocket.TSocket(host, int(port))
transport = self.ttransport(socket)
protocol = self.tprotocol(transport)
connection = self.client_cls(protocol)
transport.open()
except Exception as e:
logger.info("connection to %s failed, try other endpoints: %s", endpoint, e)
continue
return connection
raise Exception("Can't open a connection on any endpoints!")
@staticmethod
def _close_connection(conn):
"""
Helper that closes a thrift client input and output transports
Args:
conn: the thrift connection that should be closed
"""
try:
conn._iprot.trans.close()
except:
logger.warn('failed to close iprot transport on %s', conn)
try:
conn._oprot.trans.close()
except:
logger.warn('failed to close oprot transport on %s', conn)
def get_connection(self):
"""
Get a connection from the pool. If the pool is empty, it will create
a new one. If the pool is full, it blocks until one is available.
Returns:
A connected thrift client
"""
self._semaphore.acquire()
if self._closed:
raise RuntimeError('connection pool already closed')
try:
(conn, t) = self._connection_queue.get(block=False)
self._connection_queue.task_done()
return conn
except Queue.Empty:
try:
return self._create_connection()
except:
self._semaphore.release()
raise
def return_connection(self, conn):
"""
Return the given connection to the pool of available connections. If
the pool is closed, just close the connection.
Args:
conn: the thrift connection to be returned to the connection pool
"""
if self._closed:
self._close_connection(conn)
else:
self._connection_queue.put((conn, current_time_millis()))
self._semaphore.release()
def release_conn(self, conn):
"""
Release the connection without returning to the pool. Use when the
connection is bad, i.e. catching a TTransportException.
Args:
conn: the thrift connection being released
"""
try:
self._close_connection(conn)
except:
# Ignore, not returning to the pool anyway
pass
if not self._closed:
self._semaphore.release()
