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
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 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 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