def connect_to_thrift(conf):
"""
Connect to a thrift endpoint as determined by the 'conf' parameter.
Note that this does *not* open the transport.
Returns a tuple of (service, protocol, transport)
"""
sock = TSocket(conf.host, conf.port)
if conf.timeout_seconds:
# Thrift trivia: You can do this after the fact with
# _grab_transport_from_wrapper(self.wrapped.transport).setTimeout(seconds*1000)
sock.setTimeout(conf.timeout_seconds * 1000.0)
if conf.use_sasl:
def sasl_factory():
saslc = sasl.Client()
saslc.setAttr("host", conf.host)
saslc.setAttr("service", conf.kerberos_principal)
saslc.init()
return saslc
transport = TSaslClientTransport(sasl_factory, "GSSAPI", sock)
else:
transport = TBufferedTransport(sock)
protocol = TBinaryProtocol(transport)
service = conf.klass(protocol)
return service, protocol, transport
def __init__(self, host=None, port=10000, authMechanism=None, user=None, password=None, database=None, configuration=None, timeout=None):
authMechanisms = set(['NOSASL', 'PLAIN', 'KERBEROS', 'LDAP'])
if authMechanism not in authMechanisms:
raise NotImplementedError('authMechanism is either not supported or not implemented')
#Must set a password for thrift, even if it doesn't need one
#Open issue with python-sasl
if authMechanism == 'PLAIN' and (password is None or len(password) == 0):
password = '******'
socket = TSocket(host, port)
socket.setTimeout(timeout)
if authMechanism == 'NOSASL':
transport = TBufferedTransport(socket)
else:
sasl_mech = 'PLAIN'
saslc = sasl.Client()
saslc.setAttr("username", user)
saslc.setAttr("password", password)
if authMechanism == 'KERBEROS':
krb_host,krb_service = self._get_krb_settings(host, configuration)
sasl_mech = 'GSSAPI'
saslc.setAttr("host", krb_host)
saslc.setAttr("service", krb_service)
saslc.init()
transport = TSaslClientTransport(saslc, sasl_mech, socket)
self.client = TCLIService.Client(TBinaryProtocol(transport))
transport.open()
res = self.client.OpenSession(TOpenSessionReq(username=user, password=password, configuration=configuration))
self.session = res.sessionHandle
if database is not None:
with self.cursor() as cur:
query = "USE {0}".format(database)
cur.execute(query)
def get_fast_transport(endpoint, timeout=5000):
"""
采用cython实现的transport
:param endpoint:
:param timeout:
:return:
"""
global _fast_transport
if not _fast_transport:
if endpoint.find(":") != -1:
hostport = endpoint.split(":")
host = hostport[0]
port = int(hostport[1])
unix_socket = None
else:
host = None
port = None
unix_socket = endpoint
socket = TSocket(host=host, port=port, unix_socket=unix_socket)
socket.setTimeout(timeout)
socket.open()
_fast_transport = TCyFramedTransport(
socket, maxIdleTime=1200) # 20分钟没有写数据，则重新打开transport
return _fast_transport
def connect_to_thrift(conf):
"""
Connect to a thrift endpoint as determined by the 'conf' parameter.
Note that this does *not* open the transport.
Returns a tuple of (service, protocol, transport)
"""
sock = TSocket(conf.host, conf.port)
if conf.timeout_seconds:
# Thrift trivia: You can do this after the fact with
# _grab_transport_from_wrapper(self.wrapped.transport).setTimeout(seconds*1000)
sock.setTimeout(conf.timeout_seconds*1000.0)
if conf.use_sasl:
def sasl_factory():
saslc = sasl.Client()
saslc.setAttr("host", conf.host)
saslc.setAttr("service", conf.kerberos_principal)
saslc.init()
return saslc
transport = TSaslClientTransport(sasl_factory, "GSSAPI", sock)
else:
transport = TBufferedTransport(sock)
protocol = TBinaryProtocol(transport)
service = conf.klass(protocol)
return service, protocol, transport
def _connect(self):
if hasattr(self.context.config, 'HBASE_STORAGE_SERVER_HOSTS'):
host = self.context.config.HBASE_STORAGE_SERVER_HOSTS[
(self.context.server.port + self.hbase_server_offset) %
len(self.context.config.HBASE_STORAGE_SERVER_HOSTS)]
else:
host = self.context.config.HBASE_STORAGE_SERVER_HOST
transport = TBufferedTransport(
TSocket(host=host,
port=self.context.config.HBASE_STORAGE_SERVER_PORT))
socket = TSocket(host=host,
port=self.context.config.HBASE_STORAGE_SERVER_PORT)
# Timeout is sum of HTTP timeouts, plus a bit.
try:
timeout = 5
socket.setTimeout(timeout * 1000)
except:
pass
try:
transport = TBufferedTransport(socket)
transport.open()
protocol = TBinaryProtocol.TBinaryProtocol(transport)
self.storage = Hbase.Client(protocol)
logger.info("Connected to HBase server " + host + ":" +
str(self.context.config.HBASE_STORAGE_SERVER_PORT))
except:
logger.error("Error connecting to HBase server " + host + ":" +
str(self.context.config.HBASE_STORAGE_SERVER_PORT))
self.hbase_server_offset = self.hbase_server_offset + 1
class QsfpServiceClient(QsfpService.Client):
DEFAULT_PORT = 5910
DEFAULT_TIMEOUT = 10.0
# we ignore the value of port
def __init__(self, host, port=None, timeout=None):
# In a box with all 32 QSFP ports populated, it takes about 7.5s right
# now to read all 32 QSFP ports. So, put the defaut timeout to 10s.
self.host = host
timeout = timeout or self.DEFAULT_TIMEOUT
self._socket = TSocket(host, self.DEFAULT_PORT)
# TSocket.setTimeout() takes a value in milliseconds
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
QsfpService.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
def create_client(
client_klass,
host=None,
port=None,
client_type=None,
path=None,
timeout=None,
):
"""
Given a thrift client class, and a host/port
return a client using HeaderTransport
"""
from thrift.transport.TSocket import TSocket
from thrift.protocol.THeaderProtocol import THeaderProtocol
sock = TSocket(host=host, port=port, unix_socket=path)
sock.setTimeout(timeout)
protocol = THeaderProtocol(
sock,
client_types=[client_type]
if client_type
else None, # We accept the same as our inital send_
client_type=client_type, # Used for the inital send_
)
sock.open()
return client_klass(protocol)
def get_fast_transport(endpoint, timeout=5000):
"""
采用cython实现的transport
:param endpoint:
:param timeout:
:return:
"""
global _fast_transport
if not _fast_transport:
if endpoint.find(":") != -1:
hostport = endpoint.split(":")
host = hostport[0]
port = int(hostport[1])
unix_socket = None
else:
host = None
port = None
unix_socket = endpoint
socket = TSocket(host=host, port=port, unix_socket=unix_socket)
socket.setTimeout(timeout)
socket.open()
_fast_transport = TCyFramedTransport(socket, maxIdleTime=1200) # 20分钟没有写数据，则重新打开transport
return _fast_transport
class QsfpServiceClient(QsfpService.Client):
DEFAULT_PORT = 5910
DEFAULT_TIMEOUT = 10.0
# we ignore the value of port
def __init__(self, host, port=None, timeout=None):
# In a box with all 32 QSFP ports populated, it takes about 7.5s right
# now to read all 32 QSFP ports. So, put the defaut timeout to 10s.
self.host = host
timeout = timeout or self.DEFAULT_TIMEOUT
self._socket = TSocket(host, self.DEFAULT_PORT)
# TSocket.setTimeout() takes a value in milliseconds
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
QsfpService.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSocket(self.host, self.port)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(self.transport)
self.client = Hbase.Client(protocol)
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSocket(self.host, self.port)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
protocol = self._protocol_class(self.transport)
self.client = Client(protocol)
def connect(server='localhost', port=9090, timeout=None):
socket = TSocket(server, int(port))
if timeout is not None:
socket.setTimeout(timeout)
transport = TBufferedTransport(socket)
transport.open()
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
client = Hbase.Client(protocol)
return client
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSocket(self.host, self.port)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(self.transport)
self.client = Hbase.Client(protocol)
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSocket(self.host, self.port)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
protocol = self._protocol_class(self.transport)
self.client = Client(protocol)
def connect(server='localhost', port=9090, timeout=None):
socket = TSocket(server, int(port))
if timeout is not None:
socket.setTimeout(timeout)
transport = TBufferedTransport(socket)
transport.open()
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
client = Hbase.Client(protocol)
return client
def connection_to_lb(self):
if self.unix_socket:
info_logger.info("Prepare open a socket to lb: %s, pid: %s", self.unix_socket, self.pid)
else:
info_logger.info("Prepare open a socket to lb: %s:%s, pid: %s", self.host, self.port, self.pid)
# 1. 创建一个到lb的连接，然后开始读取Frame, 并且返回数据
socket = TSocket(host=self.host, port=self.port, unix_socket=self.unix_socket)
try:
if not socket.isOpen():
socket.open()
socket.setTimeout(5000) # 出现异常，会自己重启
except TTransportException:
info_logger.info("Sleep %ds for another retry, pid: %s", self.reconnect_interval, self.pid)
time.sleep(self.reconnect_interval)
print_exception(info_logger)
if self.reconnect_interval < 4:
self.reconnect_interval *= 2
return
# 2. 连接创建成功
self.reconnect_interval = 1
self.socket = socket
# 每次建立连接都重新构建
self.queue = gevent.queue.Queue()
self.connection_ok = True
info_logger.info("Begin request loop....")
# 3. 在同一个transport上进行读写数据
transport = TCyFramedTransportEx(socket)
#
# 关注 transport的接口:
# flush_frame_buff
# read_frame
#
g1 = gevent.spawn(self.loop_reader, transport, self.queue)
g2 = gevent.spawn(self.loop_writer, transport, self.queue)
g3 = gevent.spawn(self.loop_hb_detect, transport)
gevent.joinall([g1, g2, g3])
# 4. 关闭连接
try:
# 什么情况下会关闭呢? 连接断开了,
print time.strftime(ISOTIMEFORMAT, time.localtime()), "Trans Closed, queue size: ", self.queue.qsize(), ", pid: ", self.pid
self.queue = None
self.socket = None
transport.close() # 关闭transport(而且transport也不会继续复用)
except:
print_exception(info_logger)
pass
def _refresh_thrift_client(self):
"""Refresh the Thrift socket, transport, and client."""
socket = TSocket(self.host, self.port)
if self.timeout is not None:
socket.setTimeout(self.timeout)
self.transport = self._transport_class(socket)
if self.use_kerberos:
self.transport = TSaslClientTransport(self.transport, self.host, self.sasl_service_name)
protocol = self._protocol_class(self.transport)
self.client = Hbase.Client(protocol)
def _create_conn(self):
self._host_index += 1
self._host_index %= len(self._host_list)
host = self._host_list[self._host_index]
parts = host.split(':')
host = parts[0]
port = int(parts[1])
conn = TSocket(host, port)
conn.setTimeout(self._time_out)
conn.open()
return conn
class Client(object):
__metaclass__ = EnsureConnectionClient
def connect(self, select_new=False):
# 已连接,设置请求超时时间
if self._connected:
self._socket.setTimeout(self._req_timeout)
return self._client
if select_new:
host = self._host_selector.get_host()
ip, port = host.split(':')
self._ip = ip
self._port = int(port)
self._socket = TSocket(self._ip, self._port)
# 设置socket连接超时时间
self._socket.setTimeout(self._socket_connection_timeout)
self._transport = self._transport_factory.getTransport(self._socket)
self._transport.open()
self._protocol = self._protocol_factory.getProtocol(self._transport)
thrift_client_class = self.__class__.__bases__[0]
self._client = thrift_client_class(self._protocol)
self._connected = True
self._connected_at = time.time()
# 设置请求超时时间
self._socket.setTimeout(self._req_timeout)
return self._client
def disconnect(self):
if self._connected:
self._transport.close()
self._connected = False
self._connected_at = 0
self._socket = None
self._transport = None
self._protocol = None
self._client = None
self._request_served_num = 0
def refresh_connection(self, request_num_for_disconnect=REQUEST_NUM_FOR_DISCONNECT):
self._request_served_num += 1
if self._connected:
# 连接超期或者已达请求数则关闭连接
if time.time() - self._connected_at > SECS_FOR_DISCONNECT \
or self._request_served_num == request_num_for_disconnect:
self.disconnect()
def __init__(self,
unix_socket=None,
host=None,
port=10000,
authMechanism=None,
user=None,
password=None,
database=None,
configuration=None,
timeout=None):
authMechanisms = set(['NOSASL', 'PLAIN', 'KERBEROS', 'LDAP'])
if authMechanism not in authMechanisms:
raise NotImplementedError(
'authMechanism is either not supported or not implemented')
#Must set a password for thrift, even if it doesn't need one
#Open issue with python-sasl
if authMechanism == 'PLAIN' and (password is None
or len(password) == 0):
password = '******'
if unix_socket is not None:
socket = TSocket(unix_socket=unix_socket)
else:
socket = TSocket(host, port)
socket.setTimeout(timeout)
if authMechanism == 'NOSASL':
transport = TBufferedTransport(socket)
else:
sasl_mech = 'PLAIN'
saslc = sasl.Client()
saslc.setAttr("username", user)
saslc.setAttr("password", password)
if authMechanism == 'KERBEROS':
krb_host, krb_service = self._get_krb_settings(
host, configuration)
sasl_mech = 'GSSAPI'
saslc.setAttr("host", krb_host)
saslc.setAttr("service", krb_service)
saslc.init()
transport = TSaslClientTransport(saslc, sasl_mech, socket)
self.client = TCLIService.Client(TBinaryProtocol(transport))
transport.open()
res = self.client.OpenSession(
TOpenSessionReq(username=user,
password=password,
configuration=configuration))
self.session = res.sessionHandle
if database is not None:
with self.cursor() as cur:
query = "USE {0}".format(database)
cur.execute(query)
def construct_client(klass, host, port, service_name, timeout_seconds=45):
"""
Constructs a thrift client, lazily.
"""
sock = TSocket(host, port)
if timeout_seconds:
# Thrift trivia: You can do this after the fact with
# self.wrapped.transport._TBufferedTransport__trans.setTimeout(seconds*1000)
sock.setTimeout(timeout_seconds*1000.0)
transport = TBufferedTransport(sock)
protocol = TBinaryProtocol(transport)
service = klass(protocol)
return SuperClient(service, transport, timeout_seconds=timeout_seconds)
class FbossAgentClient(FbossCtrl.Client):
DEFAULT_PORT = 5909
def __init__(self, host, port=None, timeout=5.0):
self.host = host
if port is None:
port = self.DEFAULT_PORT
self._socket = TSocket(host, port)
# TSocket.setTimeout() takes a value in milliseconds
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
FbossCtrl.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
#
# The getPortStats() thrift API was unfortunately renamed to getPortInfo().
# Here's a hacky workaround that tries to do the right thing regardless of
# whether the switch we are talking to supports getPortStats() or
# getPortInfo().
#
def getPortStats(self, *args, **kwargs):
return self.getPortInfo(*args, **kwargs)
def getAllPortStats(self, *args, **kwargs):
return self.getAllPortInfo(*args, **kwargs)
def getPortInfo(self, *args, **kwargs):
try:
return FbossCtrl.Client.getPortInfo(self, *args, **kwargs)
except TApplicationException as ex:
if 'Method name getPortInfo not found' in str(ex):
return FbossCtrl.Client.getPortStats(self, *args, **kwargs)
raise
def getAllPortInfo(self, *args, **kwargs):
try:
return FbossCtrl.Client.getAllPortInfo(self, *args, **kwargs)
except TApplicationException as ex:
if 'Method name getAllPortInfo not found' in str(ex):
return FbossCtrl.Client.getAllPortStats(self, *args, **kwargs)
raise
def init_protocol(args):
sock = TSocket(args.host, args.port, socket_family=socket.AF_INET)
sock.setTimeout(500)
trans = {
'buffered': TBufferedTransport,
'framed': TFramedTransport,
'http': THttpClient,
}[args.transport](sock)
trans.open()
return {
'binary': TBinaryProtocol,
'compact': TCompactProtocol,
'json': TJSONProtocol,
}[args.protocol](trans)
def init_protocol(args):
sock = TSocket(args.host, args.port, socket_family=socket.AF_INET)
sock.setTimeout(500)
trans = {
'buffered': TBufferedTransport,
'framed': TFramedTransport,
'http': THttpClient,
}[args.transport](sock)
trans.open()
return {
'binary': TBinaryProtocol,
'compact': TCompactProtocol,
'json': TJSONProtocol,
}[args.protocol](trans)
def test_isOpen_checks_for_readability(self):
# https://docs.python.org/3/library/socket.html#notes-on-socket-timeouts
# https://docs.python.org/3/library/socket.html#socket.socket.settimeout
timeouts = [
None, # blocking mode
0, # non-blocking mode
1.0, # timeout mode
]
for timeout in timeouts:
acc = ServerAcceptor(TServerSocket(port=0))
acc.start()
sock = TSocket(host="localhost", port=acc.port)
self.assertFalse(sock.isOpen())
sock.open()
sock.setTimeout(timeout)
# the socket shows as open immediately after connecting
self.assertTrue(sock.isOpen())
# and remains open during usage
sock.write(b"hello")
self.assertTrue(sock.isOpen())
while True:
try:
sock.read(5)
except TTransportException as exc:
if exc.inner.errno == errno.EAGAIN:
# try again when we're in non-blocking mode
continue
raise
break
self.assertTrue(sock.isOpen())
# once the server side closes, it no longer shows open
acc.client.close(
) # this also blocks until the other thread is done
acc.close()
self.assertFalse(sock.isOpen())
sock.close()
def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None,
autoconnect=True, table_prefix=None,
table_prefix_separator='_', compat='0.92',
transport='buffered'):
# Allow host and port to be None, which may be easier for
# applications wrapping a Connection instance.
self.host = host or DEFAULT_HOST
self.port = port or DEFAULT_PORT
self.timeout = timeout
if compat not in COMPAT_MODES:
raise ValueError("'compat' must be one of %s"
% ", ".join(COMPAT_MODES))
if transport not in THRIFT_TRANSPORTS:
raise ValueError("'transport' must be one of %s"
% ", ".join(THRIFT_TRANSPORTS.keys()))
if table_prefix is not None and not isinstance(table_prefix, basestring):
raise TypeError("'table_prefix' must be a string")
if not isinstance(table_prefix_separator, basestring):
raise TypeError("'table_prefix_separator' must be a string")
self.compat = compat
self.table_prefix = table_prefix
self.table_prefix_separator = table_prefix_separator
socket = TSocket(self.host, self.port)
if timeout is not None:
socket.setTimeout(timeout)
self.transport = THRIFT_TRANSPORTS[transport](socket)
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(self.transport)
self.client = Hbase.Client(protocol)
if autoconnect:
self.open()
self._initialized = True
class PcapPushSubClient(PcapPushSubscriber.Client):
DEFAULT_PORT = 5911
def __init__(self, host, port=None, timeout=5.0):
self.host = host
if port is None:
port = self.DEFAULT_PORT
self._socket = TSocket(host, port)
# TSocket.setTimeout() takes a value in milliseconds
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
PcapPushSubscriber.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
class NetlinkManagerClient(NetlinkManagerService.Client):
DEFAULT_PORT = 5912
def __init__(self, host, port=None, timeout=5.0):
self.host = host
if port is None:
port = self.DEFAULT_PORT
self._socket = TSocket(host, port)
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
NetlinkManagerService.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
class TestClient(TestService.Client):
DEFAULT_PORT = fboss.system_tests.test.constants.DEFAULT_PORT
def __init__(self, host, port=None, timeout=10.0):
self.host = host
if port is None:
port = self.DEFAULT_PORT
self._socket = TSocket(host, port)
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
TestService.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
class TestClient(TestService.Client):
DEFAULT_PORT = fboss.system_tests.test.constants.DEFAULT_PORT
def __init__(self, host, port=None, timeout=10.0):
self.host = host
if port is None:
port = self.DEFAULT_PORT
self._socket = TSocket(host, port)
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
TestService.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
class QsfpServiceClient(QsfpService.Client):
DEFAULT_PORT = 5910
# we ignore the value of port
def __init__(self, host, port=None, timeout=2.0):
self.host = host
self._socket = TSocket(host, self.DEFAULT_PORT)
# TSocket.setTimeout() takes a value in milliseconds
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
QsfpService.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
class EdenClient(EdenService.Client):
'''
EdenClient is a subclass of EdenService.Client that provides
a few additional conveniences:
- Smarter constructor
- Implement the context manager __enter__ and __exit__ methods, so it can
be used in with statements.
'''
def __init__(self, eden_dir=None, mounted_path=None):
self._eden_dir = eden_dir
if mounted_path:
sock_path = os.path.join(mounted_path, '.eden', 'socket')
else:
sock_path = os.path.join(self._eden_dir, SOCKET_PATH)
self._socket = TSocket(unix_socket=sock_path)
self._socket.setTimeout(60000) # in milliseconds
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
super(EdenClient, self).__init__(self._protocol)
def __enter__(self):
self.open()
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
self.close()
def open(self):
try:
self._transport.open()
except TTransportException as ex:
self._transport.close()
if ex.type == TTransportException.NOT_OPEN:
raise EdenNotRunningError(self._eden_dir)
def close(self):
if self._transport is not None:
self._transport.close()
self._transport = None
class PlainTextFbossAgentClientDontUseInFb(FbossCtrl.Client):
DEFAULT_PORT = 5909
def __init__(self, host, port=None, timeout=5.0):
self.host = host
if port is None:
port = self.DEFAULT_PORT
self._socket = TSocket(host, port)
# TSocket.setTimeout() takes a value in milliseconds
self._socket.setTimeout(timeout * 1000)
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
self._transport.open()
FbossCtrl.Client.__init__(self, self._protocol)
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
self._transport.close()
class EdenClient(EdenService.Client):
"""
EdenClient is a subclass of EdenService.Client that provides
a few additional conveniences:
- Smarter constructor
- Implement the context manager __enter__ and __exit__ methods, so it can
be used in with statements.
"""
def __init__(self, eden_dir=None, socket_path=None):
# type: (Optional[str], Optional[str]) -> None
if socket_path is not None:
self._socket_path = socket_path
elif eden_dir is not None:
self._socket_path = os.path.join(eden_dir, SOCKET_PATH)
else:
raise TypeError("one of eden_dir or socket_path is required")
self._socket = TSocket(unix_socket=self._socket_path)
# We used to set a timeout here, but picking the right duration is hard,
# and safely retrying an arbitrary thrift call may not be safe. So we
# just leave the client with no timeout.
# self.set_timeout(60)
self.set_timeout(None)
transport = THeaderTransport(self._socket)
self._transport = transport # type: Optional[THeaderTransport]
self._protocol = THeaderProtocol(transport)
super(EdenClient, self).__init__(self._protocol)
def __enter__(self):
# type: () -> EdenClient
self.open()
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
# type: (Any, Any, Any) -> Optional[bool]
self.close()
return False
def open(self):
# type: () -> None
try:
assert self._transport is not None
self._transport.open()
except TTransportException as ex:
self.close()
if ex.type == TTransportException.NOT_OPEN:
raise EdenNotRunningError(self._socket_path)
raise
def close(self):
# type: () -> None
if self._transport is not None:
self._transport.close()
self._transport = None
def shutdown(self):
# type: () -> None
self.initiateShutdown(
"EdenClient.shutdown() invoked with no reason by pid=%s uid=%s"
% (os.getpid(), os.getuid())
)
def initiateShutdown(self, reason):
# type: (str) -> None
"""Helper for stopping the server.
To swing through the transition from calling the base shutdown() method
with context to the initiateShutdown() method with a reason, we want to
try the latter method first, falling back to the old way to handle the
case where we deploy a newer client while an older server is still
running on the local system."""
try:
super().initiateShutdown(reason)
except TApplicationException as ex:
if ex.type == TApplicationException.UNKNOWN_METHOD:
# Running an older server build, fall back to the old shutdown
# method with no context
super().shutdown()
else:
raise
def set_timeout(self, timeout):
# type: (Optional[float]) -> None
if timeout is None:
timeout_ms = None
else:
timeout_ms = timeout * 1000
self.set_timeout_ms(timeout_ms)
def set_timeout_ms(self, timeout_ms):
# type: (Optional[float]) -> None
self._socket.setTimeout(timeout_ms)
class EdenClient(EdenService.Client):
"""
EdenClient is a subclass of EdenService.Client that provides
a few additional conveniences:
- Smarter constructor
- Implement the context manager __enter__ and __exit__ methods, so it can
be used in with statements.
"""
def __init__(self, eden_dir=None, socket_path=None):
# type: (Optional[str], Optional[str]) -> None
if socket_path is not None:
self._socket_path = socket_path
elif eden_dir is not None:
self._socket_path = os.path.join(eden_dir, SOCKET_PATH)
else:
raise TypeError("one of eden_dir or socket_path is required")
if sys.platform == "win32":
self._socket = WinTSocket(unix_socket=self._socket_path)
else:
self._socket = TSocket(unix_socket=self._socket_path)
# We used to set a timeout here, but picking the right duration is hard,
# and safely retrying an arbitrary thrift call may not be safe. So we
# just leave the client with no timeout.
# self.set_timeout(60)
self.set_timeout(None)
transport = THeaderTransport(self._socket)
self._transport = transport # type: Optional[THeaderTransport]
self._protocol = THeaderProtocol(transport)
super(EdenClient, self).__init__(self._protocol)
def __enter__(self):
# type: () -> EdenClient
self.open()
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
# type: (Any, Any, Any) -> Optional[bool]
self.close()
return False
def open(self):
# type: () -> None
transport = self._transport
assert transport is not None
try:
transport.open()
except TTransportException as ex:
self.close()
if ex.type == TTransportException.NOT_OPEN:
raise EdenNotRunningError(self._socket_path)
raise
except WindowsSocketException:
self.close()
raise EdenNotRunningError(self._socket_path)
def close(self):
# type: () -> None
if self._transport is not None:
# pyre-fixme[16]: `Optional` has no attribute `close`.
self._transport.close()
self._transport = None
def getDaemonInfo(self):
# type: () -> DaemonInfo
try:
info = super(EdenClient, self).getDaemonInfo()
except TApplicationException as ex:
if ex.type != TApplicationException.UNKNOWN_METHOD:
raise
# Older versions of EdenFS did not have a getDaemonInfo() method
pid = super(EdenClient, self).getPid()
info = DaemonInfo(pid=pid, status=None)
# Older versions of EdenFS did not return status information in the
# getDaemonInfo() response.
if info.status is None:
info.status = super(EdenClient, self).getStatus()
return info
def getPid(self):
# type: () -> int
try:
return self.getDaemonInfo().pid
except TApplicationException as ex:
if ex.type == TApplicationException.UNKNOWN_METHOD:
# Running on an older server build, fall back to the
# old getPid() method.
return super(EdenClient, self).getPid()
else:
raise
def set_timeout(self, timeout):
# type: (Optional[float]) -> None
if timeout is None:
timeout_ms = None
else:
timeout_ms = timeout * 1000
self.set_timeout_ms(timeout_ms)
def set_timeout_ms(self, timeout_ms):
# type: (Optional[float]) -> None
self._socket.setTimeout(timeout_ms)
class EdenClient(EdenService.Client):
"""
EdenClient is a subclass of EdenService.Client that provides
a few additional conveniences:
- Smarter constructor
- Implement the context manager __enter__ and __exit__ methods, so it can
be used in with statements.
"""
def __init__(self, eden_dir=None, socket_path=None):
# type: (Optional[str], Optional[str]) -> None
if socket_path is not None:
self._socket_path = socket_path
elif eden_dir is not None:
self._socket_path = os.path.join(eden_dir, SOCKET_PATH)
else:
raise TypeError("one of eden_dir or socket_path is required")
self._socket = TSocket(unix_socket=self._socket_path)
# We used to set a timeout here, but picking the right duration is hard,
# and safely retrying an arbitrary thrift call may not be safe. So we
# just leave the client with no timeout.
# self.set_timeout(60)
self.set_timeout(None)
transport = THeaderTransport(self._socket)
self._transport = transport # type: Optional[THeaderTransport]
self._protocol = THeaderProtocol(transport)
super(EdenClient, self).__init__(self._protocol)
def __enter__(self):
# type: () -> EdenClient
self.open()
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
# type: (Any, Any, Any) -> Optional[bool]
self.close()
return False
def open(self):
# type: () -> None
try:
assert self._transport is not None
self._transport.open()
except TTransportException as ex:
self.close()
# pyre-fixme[20]: Call `object.__eq__` expects argument `o`.
if ex.type == TTransportException.NOT_OPEN:
# pyre: Expected `str` for 1st anonymous parameter to call
# pyre: `eden.thrift.client.EdenNotRunningError.__init__` but
# pyre-fixme[6]: got `Optional[str]`.
raise EdenNotRunningError(self._socket_path)
raise
def close(self):
# type: () -> None
if self._transport is not None:
self._transport.close()
self._transport = None
def shutdown(self):
# type: () -> None
self.initiateShutdown(
"EdenClient.shutdown() invoked with no reason by pid=%s uid=%s" %
(os.getpid(), os.getuid()))
def initiateShutdown(self, reason):
# type: (str) -> None
"""Helper for stopping the server.
To swing through the transition from calling the base shutdown() method
with context to the initiateShutdown() method with a reason, we want to
try the latter method first, falling back to the old way to handle the
case where we deploy a newer client while an older server is still
running on the local system."""
try:
super().initiateShutdown(reason)
except TApplicationException as ex:
if ex.type == TApplicationException.UNKNOWN_METHOD:
# Running an older server build, fall back to the old shutdown
# method with no context
super().shutdown()
else:
raise
def set_timeout(self, timeout):
# type: (Optional[float]) -> None
if timeout is None:
timeout_ms = None
else:
timeout_ms = timeout * 1000
self.set_timeout_ms(timeout_ms)
def set_timeout_ms(self, timeout_ms):
# type: (Optional[float]) -> None
self._socket.setTimeout(timeout_ms)
class ThriftClient(object):
"""Base Class for easy interfacing with thrift services.
`ThriftClient` can be used directly, or subclassed. If subclassed,
you can override a variety of attributes in order to make instantiation
more natural:
MODULE - The generated thrift module. Should include Iface and Client
HOST - A default host to connect to. Possibly a vip.
PORT - A default port to connect to.
TRANSPORT_CLASS - By default, `TFramedTransport`
PROTOCOL_CLASS - By default, `TBinaryProtocol`
CONNECT_TIMEOUT - By default, 3.0 seconds
Instantiation should be done via the class methods, `for_hostport()`, and
`for_localhost()` as appropriate. These helpers more aggressively require
`port` and `host` arguments as appropriate. Generic construction arguments
override the class attribute defaults:
`module` - Generated thrift module.
`host` - IP address to connect to.
`port` - Port to connect to.
`connect_Timeout` - Socket connection timeout
`transport_class` - Thrift transport class
`protocol_class` - Thrift protocol class
Additional features are configurable with other arguments:
`lazy` - Default: True, connect on first RPC invocation, instead of
at construction time.
Connections are made lazily, when the first rpc invocation occurs, so you
do need to wrap client instantiation with a try-catch.
"""
MODULE = None
HOST = None
PORT = None
TRANSPORT_CLASS = TFramedTransport
PROTOCOL_CLASS = TBinaryProtocol
CONNECT_TIMEOUT = 3.0
@classmethod
def for_hostport(cls, host=None, port=None, **kwargs):
assert host or cls.HOST, "You must define a host!"
assert port or cls.PORT, "You must define a port!"
return cls(host=host, port=port, **kwargs)
@classmethod
def for_localhost(cls, port=None, **kwargs):
assert port or cls.PORT, "You must define a port!"
return cls(host='127.0.0.1', port=port, **kwargs)
def _initAttribute(self, name, value, default):
if value is None:
value = default
setattr(self, name, value)
def __init__(self,
host=None,
port=None,
module=None,
lazy=True,
connect_timeout=None,
transport_class=None,
protocol_class=None):
self._initAttribute('host', host, self.HOST)
self._initAttribute('port', port, self.PORT)
self._initAttribute('module', module, self.MODULE)
self._initAttribute('connect_timeout', connect_timeout,
self.CONNECT_TIMEOUT)
self._initAttribute('transport_class', transport_class,
self.TRANSPORT_CLASS)
self._initAttribute('protocol_class', protocol_class,
self.PROTOCOL_CLASS)
self.lazy = lazy
assert self.module is not None, "You must define a thrift module!"
if self.lazy:
self._client = None
else:
self._connect()
def _connect(self):
# TODO: Add some kind of support for HTTP or SSLSocket
self._socket = TSocket(self.host, self.port)
self._socket.setTimeout(int(self.connect_timeout * 1000))
self._transport = self.transport_class(self._socket)
self._protocol = self.protocol_class(self._transport)
self._client = self.module.Client(self._protocol)
self._transport.open()
def _lazyCall(self, name, *args, **kwargs):
if self._client is None:
self._connect()
# TODO: Automatically connect on timed out connections
return getattr(self._client, name)(*args, **kwargs)
def __getattr__(self, name):
getattr(self.module.Client, name)
return partial(self._lazyCall, name)
def connection_to_lb(self):
if self.unix_socket:
info_logger.info("Prepare open a socket to lb: %s, pid: %s",
self.unix_socket, self.pid)
else:
info_logger.info("Prepare open a socket to lb: %s:%s, pid: %s",
self.host, self.port, self.pid)
# 1. 创建一个到lb的连接，然后开始读取Frame, 并且返回数据
socket = TSocket(host=self.host,
port=self.port,
unix_socket=self.unix_socket)
try:
if not socket.isOpen():
socket.open()
socket.setTimeout(5000) # 出现异常，会自己重启
except TTransportException:
info_logger.info("Sleep %ds for another retry, pid: %s",
self.reconnect_interval, self.pid)
time.sleep(self.reconnect_interval)
print_exception(info_logger)
if self.reconnect_interval < 4:
self.reconnect_interval *= 2
return
# 2. 连接创建成功
self.reconnect_interval = 1
self.socket = socket
# 每次建立连接都重新构建
self.queue = gevent.queue.Queue()
self.connection_ok = True
info_logger.info("Begin request loop....")
# 3. 在同一个transport上进行读写数据
transport = TCyFramedTransportEx(socket)
#
# 关注 transport的接口:
# flush_frame_buff
# read_frame
#
g1 = gevent.spawn(self.loop_reader, transport, self.queue)
g2 = gevent.spawn(self.loop_writer, transport, self.queue)
g3 = gevent.spawn(self.loop_hb_detect, transport)
gevent.joinall([g1, g2, g3])
# 4. 关闭连接
try:
# 什么情况下会关闭呢? 连接断开了,
print time.strftime(ISOTIMEFORMAT, time.localtime(
)), "Trans Closed, queue size: ", self.queue.qsize(
), ", pid: ", self.pid
self.queue = None
self.socket = None
transport.close() # 关闭transport(而且transport也不会继续复用)
except:
print_exception(info_logger)
pass
class ThriftClient(object):
"""Base Class for easy interfacing with thrift services.
`ThriftClient` can be used directly, or subclassed. If subclassed,
you can override a variety of attributes in order to make instantiation
more natural:
MODULE - The generated thrift module. Should include Iface and Client
HOST - A default host to connect to. Possibly a vip.
PORT - A default port to connect to.
TRANSPORT_CLASS - By default, `TFramedTransport`
PROTOCOL_CLASS - By default, `TBinaryProtocol`
CONNECT_TIMEOUT - By default, 3.0 seconds
Instantiation should be done via the class methods, `for_hostport()`, and
`for_localhost()` as appropriate. These helpers more aggressively require
`port` and `host` arguments as appropriate. Generic construction arguments
override the class attribute defaults:
`module` - Generated thrift module.
`host` - IP address to connect to.
`port` - Port to connect to.
`connect_Timeout` - Socket connection timeout
`transport_class` - Thrift transport class
`protocol_class` - Thrift protocol class
Additional features are configurable with other arguments:
`lazy` - Default: True, connect on first RPC invocation, instead of
at construction time.
Connections are made lazily, when the first rpc invocation occurs, so you
do need to wrap client instantiation with a try-catch.
"""
MODULE = None
HOST = None
PORT = None
TRANSPORT_CLASS = TFramedTransport
PROTOCOL_CLASS = TBinaryProtocol
CONNECT_TIMEOUT = 3.0
@classmethod
def for_hostport(cls, host=None, port=None, **kwargs):
assert host or cls.HOST, "You must define a host!"
assert port or cls.PORT, "You must define a port!"
return cls(host=host, port=port, **kwargs)
@classmethod
def for_localhost(cls, port=None, **kwargs):
assert port or cls.PORT, "You must define a port!"
return cls(host='127.0.0.1', port=port, **kwargs)
def _initAttribute(self, name, value, default):
if value is None:
value = default
setattr(self, name, value)
def __init__(self, host=None, port=None, module=None, lazy=True,
connect_timeout=None, transport_class=None,
protocol_class=None):
self._initAttribute('host', host, self.HOST)
self._initAttribute('port', port, self.PORT)
self._initAttribute('module', module, self.MODULE)
self._initAttribute('connect_timeout', connect_timeout,
self.CONNECT_TIMEOUT)
self._initAttribute('transport_class', transport_class,
self.TRANSPORT_CLASS)
self._initAttribute('protocol_class', protocol_class,
self.PROTOCOL_CLASS)
self.lazy = lazy
assert self.module is not None, "You must define a thrift module!"
if self.lazy:
self._client = None
else:
self._connect()
def _connect(self):
# TODO: Add some kind of support for HTTP or SSLSocket
self._socket = TSocket(self.host, self.port)
self._socket.setTimeout(int(self.connect_timeout * 1000))
self._transport = self.transport_class(self._socket)
self._protocol = self.protocol_class(self._transport)
self._client = self.module.Client(self._protocol)
self._transport.open()
def _lazyCall(self, name, *args, **kwargs):
if self._client is None:
self._connect()
# TODO: Automatically connect on timed out connections
return getattr(self._client, name)(*args, **kwargs)
def __getattr__(self, name):
getattr(self.module.Client, name)
return partial(self._lazyCall, name)
