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 call_processor(self, input, client_type, protocol_type,
client_principal):
try:
# The input string has already had the header removed, but
# the python processor will expect it to be there. In
# order to reconstitute the message with headers, we use
# the THeaderProtocol object to write into a memory
# buffer, then pass that buffer to the python processor.
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf, client_types=[client_type])
trans.set_protocol_id(protocol_type)
trans.write(input)
trans.flush()
prot_buf = TMemoryBuffer(write_buf.getvalue())
prot = THeaderProtocol(prot_buf)
ctx = TCppConnectionContext(client_principal)
self.processor.process(prot, prot, ctx)
# And on the way out, we need to strip off the header,
# because the C++ code will expect to add it.
read_buf = TMemoryBuffer(prot_buf.getvalue())
trans = THeaderTransport(read_buf, client_types=[client_type])
trans.readFrame(0)
return trans.cstringio_buf.read()
except:
# Don't let exceptions escape back into C++
traceback.print_exc()
def __init__(self, trans, client_types=None,
fuzz_fields=None, fuzz_all_if_empty=True, verbose=False):
"""Create a TFuzzyHeaderTransport instance.
@param fuzz_fields(iterable) Collection of string names of fields to
fuzz. Each name must be a member of
TFuzzyHeaderTransport.fuzzable_fields. If the iterable is empty and
fuzz_all_if_empty is True, all fields will be fuzzed. If the iterable
is empty and fuzz_all_if_empty is False, no fields will be fuzzed.
@param verbose(bool) Whether to print fuzzed fields to stdout
@param fuzz_all_if_empty(bool) Whether to fuzz eveyr field if
the fuzz_fields iterable is None or empty.
trans and client_type are forwarded to THeaderProtocol.__init__
"""
cls = self.__class__
THeaderTransport.__init__(self, trans, client_types=client_types)
if fuzz_fields is None:
fuzz_fields = []
self._fuzz_fields = set()
for field_name in fuzz_fields:
if field_name not in cls.fuzzable_fields:
raise NameError("Invalid fuzz field: %s" % field_name)
self._fuzz_fields.add(field_name)
if not self._fuzz_fields and fuzz_all_if_empty:
# No fields were explicitly included
self._fuzz_fields = cls.fuzzable_fields
self._verbose = verbose
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 setUp(self):
self.trans = TTransport.TMemoryBuffer()
self.trans = THeaderTransport(self.trans)
self.prot = TBinaryProtocol.TBinaryProtocol(self.trans)
self.x = Xtruct()
self.x.string_thing = "Zero"
self.x.byte_thing = 0
def call_processor(self, input, headers, client_type, protocol_type,
context_data, callback):
try:
# The input string has already had the header removed, but
# the python processor will expect it to be there. In
# order to reconstitute the message with headers, we use
# the THeaderProtocol object to write into a memory
# buffer, then pass that buffer to the python processor.
should_sample = self._shouldSample()
timestamps = CallTimestamps()
timestamps.processBegin = 0
timestamps.processEnd = 0
if self.observer and should_sample:
timestamps.processBegin = int(time.time() * 10**6)
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf)
trans._THeaderTransport__client_type = client_type
trans._THeaderTransport__write_headers = headers
trans.set_protocol_id(protocol_type)
trans.write(input)
trans.flush()
prot_buf = TMemoryBuffer(write_buf.getvalue())
prot = THeaderProtocol(prot_buf, client_types=[client_type])
ctx = TCppConnectionContext(context_data)
ret = self.processor.process(prot, prot, ctx)
done_callback = partial(_ProcessorAdapter.done,
prot_buf=prot_buf,
client_type=client_type,
callback=callback)
if self.observer:
if should_sample:
timestamps.processEnd = int(time.time() * 10**6)
# This only bumps counters if `processBegin != 0` and
# `processEnd != 0` and these will only be non-zero if
# we are sampling this request.
self.observer.callCompleted(timestamps)
# This future is created by and returned from the processor's
# ThreadPoolExecutor, which keeps a reference to it. So it is
# fine for this future to end its lifecycle here.
if isinstance(ret, Future):
ret.add_done_callback(lambda x, d=done_callback: d())
else:
done_callback()
except:
# Don't let exceptions escape back into C++
traceback.print_exc()
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 __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 __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 __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 __init__(self, host, port=None, timeout=10.0):
# 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
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)
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__(self, eden_dir=None, mounted_path=None):
self._eden_dir = eden_dir
if mounted_path:
sock_path = os.readlink(
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)
# 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._socket.setTimeout(60000) # in milliseconds
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
super(EdenClient, self).__init__(self._protocol)
def __init__(self, eden_dir=None, socket_path=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._socket.setTimeout(60000) # in milliseconds
self._transport = THeaderTransport(self._socket)
self._protocol = THeaderProtocol(self._transport)
super(EdenClient, self).__init__(self._protocol)
def done(prot_buf, client_type, callback):
try:
response = prot_buf.getvalue()
if len(response) == 0:
callback.call(response)
else:
# And on the way out, we need to strip off the header,
# because the C++ code will expect to add it.
read_buf = TMemoryBuffer(response)
trans = THeaderTransport(read_buf, client_types=[client_type])
trans.readFrame(len(response))
callback.call(trans.cstringio_buf.read())
except:
traceback.print_exc()
def _get_client_by_transport(self, options, transport, socket=None):
# Create the protocol and client
if options.json:
protocol = TJSONProtocol.TJSONProtocol(transport)
elif options.compact:
protocol = TCompactProtocol.TCompactProtocol(transport)
# No explicit option about protocol is specified. Try to infer.
elif options.framed or options.unframed:
protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
elif socket is not None:
# If json, compact, framed, and unframed are not specified,
# THeaderProtocol is the default. Create a protocol using either
# fuzzy or non-fuzzy transport depending on if options.fuzz is set.
if options.fuzz is not None:
transport = TFuzzyHeaderTransport(socket,
fuzz_fields=options.fuzz,
verbose=True)
else:
transport = THeaderTransport(socket)
protocol = THeaderProtocol.THeaderProtocol(transport)
else:
self._exit(error_message=('No valid protocol '
'specified for %s' % (type(self))),
status=os.EX_USAGE)
transport.open()
self._transport = transport
client = self.service_class.Client(protocol)
return client
def create_thrift_client(
eden_dir: "Optional[str]" = None,
socket_path: "Optional[str]" = None,
timeout: "Optional[float]" = None,
) -> "EdenClient":
"""
Construct a thrift client to speak to the running eden server
instance associated with the specified mount point.
@return Returns a context manager for EdenService.Client.
"""
if socket_path is not None:
pass
elif eden_dir is not None:
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":
socket = WinTSocket(unix_socket=socket_path)
else:
socket = TSocket(unix_socket=socket_path)
# We used to set a default 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, unless one is given.
if timeout is None:
timeout_ms = None
else:
timeout_ms = timeout * 1000
socket.setTimeout(timeout_ms)
transport = THeaderTransport(socket)
protocol = THeaderProtocol(transport)
return EdenClient(socket_path, transport, protocol)
def __init__(self,
trans,
strictRead=False,
client_types=None,
client_type=None):
"""Create a THeaderProtocol instance
@param transport(TTransport) The underlying transport.
@param strictRead(bool) Turn on strictRead if using TBinaryProtocol
@param client_types([CLIENT_TYPE.HEADER, ...])
List of client types to support. Defaults to
CLIENT_TYPE.HEADER only.
"""
if isinstance(trans, THeaderTransport):
trans._THeaderTransport__supported_client_types = set(
client_types or (CLIENT_TYPE.HEADER, ))
if client_type is not None:
trans._THeaderTransport__client_type = client_type
htrans = trans
else:
htrans = THeaderTransport(trans, client_types, client_type)
TProtocolBase.__init__(self, htrans)
self.strictRead = strictRead
self.reset_protocol()
def setUp(self):
self.trans = TTransport.TMemoryBuffer()
self.trans = THeaderTransport(self.trans)
self.prot = TBinaryProtocol.TBinaryProtocol(self.trans)
self.x = Xtruct()
self.x.string_thing = "Zero"
self.x.byte_thing = 0
def schedule_timeout(self, fname, seqid):
timeout = self.timeouts[fname]
if not timeout:
return
tmo = TMemoryBuffer()
thp = THeaderTransport(tmo)
oprot = THeaderProtocol(thp)
exc = TApplicationException(TApplicationException.TIMEOUT,
"Call to {} timed out".format(fname))
oprot.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid)
exc.write(oprot)
oprot.writeMessageEnd()
thp.flush()
timeout_task = self.loop.create_task(
self.message_received(tmo.getvalue(), delay=timeout), )
self.update_pending_tasks(seqid, timeout_task)
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.readlink(
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)
# 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._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.close()
if ex.type == TTransportException.NOT_OPEN:
raise EdenNotRunningError(self._eden_dir)
raise
def close(self):
if self._transport is not None:
self._transport.close()
self._transport = None
def __init__(self, transport, allowed_client_types):
# much of the actual work for THeaderProtocol happens down in
# THeaderTransport since we need to do low-level shenanigans to detect
# if the client is sending us headers or one of the headerless formats
# we support. this wraps the real transport with the one that does all
# the magic.
if not isinstance(transport, THeaderTransport):
transport = THeaderTransport(transport, allowed_client_types)
super(THeaderProtocol, self).__init__(transport)
self._set_protocol()
def setUp(self):
"""Create two buffers, transports, and protocols.
self._h_trans uses THeaderTransport
self._f_trans uses TFuzzyHeaderTransport
"""
cls = self.__class__
# THeaderTransport attributes
self._h_buf = TMemoryBuffer()
self._h_trans = THeaderTransport(self._h_buf)
self._h_prot = THeaderProtocol(self._h_trans)
# TFuzzyHeaderTransport attributes
self._f_buf = TMemoryBuffer()
self._f_trans = TFuzzyHeaderTransport(
self._f_buf, fuzz_fields=cls.fuzz_fields,
fuzz_all_if_empty=False, verbose=False)
self._f_prot = THeaderProtocol(self._f_trans)
def schedule_timeout(self, fname, seqid):
timeout = self.timeouts[fname]
if not timeout:
return
tmo = TMemoryBuffer()
thp = THeaderTransport(tmo)
oprot = THeaderProtocol(thp)
exc = TApplicationException(
TApplicationException.TIMEOUT, "Call to {} timed out".format(fname)
)
oprot.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid)
exc.write(oprot)
oprot.writeMessageEnd()
thp.flush()
timeout_task = self.loop.create_task(
self.message_received(tmo.getvalue(), delay=timeout),
)
self.update_pending_tasks(seqid, timeout_task)
def testOversizeFrameRecv(self):
"""Test that an oversize frame on recv gets a TTransportException"""
self.trans.set_max_frame_size(200)
self.x.write(self.prot)
self.trans.flush()
inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue())
inTrans = THeaderTransport(inTrans)
inProt = TBinaryProtocol.TBinaryProtocol(inTrans)
inTrans.set_max_frame_size(2)
try:
self.x.read(inProt)
except TTransport.TTransportException:
return
self.fail("Should have gotten TTransportException")
def testOversizeFrameRecv(self):
"""Test that an oversize frame on recv gets a TTransportException"""
self.trans.set_max_frame_size(200)
self.x.write(self.prot)
self.trans.flush()
inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue())
inTrans = THeaderTransport(inTrans)
inProt = TBinaryProtocol.TBinaryProtocol(inTrans)
inTrans.set_max_frame_size(2)
try:
self.x.read(inProt)
except TTransport.TTransportException:
return
self.fail("Should have gotten TTransportException")
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 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 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 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 TestEof(unittest.TestCase):
def setUp(self):
self.trans = TTransport.TMemoryBuffer()
self.trans = THeaderTransport(self.trans)
self.prot = TBinaryProtocol.TBinaryProtocol(self.trans)
self.x = Xtruct()
self.x.string_thing = "Zero"
self.x.byte_thing = 0
def testOversizeFrameRecv(self):
"""Test that an oversize frame on recv gets a TTransportException"""
self.trans.set_max_frame_size(200)
self.x.write(self.prot)
self.trans.flush()
inTrans = TTransport.TMemoryBuffer(
self.trans.getTransport().getvalue())
inTrans = THeaderTransport(inTrans)
inProt = TBinaryProtocol.TBinaryProtocol(inTrans)
inTrans.set_max_frame_size(2)
try:
self.x.read(inProt)
except TTransport.TTransportException:
return
self.fail("Should have gotten TTransportException")
def testOversizeFrameSend(self):
"""Test that an oversize frame on send gets a TTransportException"""
self.trans.set_max_frame_size(2)
self.x.write(self.prot)
try:
self.trans.flush()
except TTransport.TTransportException:
return
self.fail("Should have gotten TTransportException")
def serialize_texception(cls, fname, seqid, exception):
"""This saves us a bit of processing time for timeout handling by
reusing the Thrift structs involved in exception serialization.
NOTE: this is not thread-safe nor it is meant to be.
"""
# the serializer is a singleton
if cls._exception_serializer is None:
buffer = TWriteOnlyBuffer()
transport = THeaderTransport(buffer)
cls._exception_serializer = THeaderProtocol(transport)
else:
transport = cls._exception_serializer.trans
buffer = transport.getTransport()
buffer.reset()
serializer = cls._exception_serializer
serializer.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid)
exception.write(serializer)
serializer.writeMessageEnd()
serializer.trans.flush()
return buffer.getvalue()
def serialize_texception(cls, fname, seqid, exception):
"""This saves us a bit of processing time for timeout handling by
reusing the Thrift structs involved in exception serialization.
NOTE: this is not thread-safe nor is it meant to be.
"""
# the serializer is a singleton
if cls._exception_serializer is None:
buffer = TWriteOnlyBuffer()
transport = THeaderTransport(buffer)
cls._exception_serializer = THeaderProtocol(transport)
else:
transport = cls._exception_serializer.trans
buffer = transport.getTransport()
buffer.reset()
serializer = cls._exception_serializer
serializer.writeMessageBegin(fname, TMessageType.EXCEPTION, seqid)
exception.write(serializer)
serializer.writeMessageEnd()
serializer.trans.flush()
return buffer.getvalue()
class TestEof(unittest.TestCase):
def setUp(self):
self.trans = TTransport.TMemoryBuffer()
self.trans = THeaderTransport(self.trans)
self.prot = TBinaryProtocol.TBinaryProtocol(self.trans)
self.x = Xtruct()
self.x.string_thing = "Zero"
self.x.byte_thing = 0
def testOversizeFrameRecv(self):
"""Test that an oversize frame on recv gets a TTransportException"""
self.trans.set_max_frame_size(200)
self.x.write(self.prot)
self.trans.flush()
inTrans = TTransport.TMemoryBuffer(self.trans.getTransport().getvalue())
inTrans = THeaderTransport(inTrans)
inProt = TBinaryProtocol.TBinaryProtocol(inTrans)
inTrans.set_max_frame_size(2)
try:
self.x.read(inProt)
except TTransport.TTransportException:
return
self.fail("Should have gotten TTransportException")
def testOversizeFrameSend(self):
"""Test that an oversize frame on send gets a TTransportException"""
self.trans.set_max_frame_size(2)
self.x.write(self.prot)
try:
self.trans.flush()
except TTransport.TTransportException:
return
self.fail("Should have gotten TTransportException")
def call_processor(self, input, headers, client_type, protocol_type,
context_data, callback):
try:
# The input string has already had the header removed, but
# the python processor will expect it to be there. In
# order to reconstitute the message with headers, we use
# the THeaderProtocol object to write into a memory
# buffer, then pass that buffer to the python processor.
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf)
trans._THeaderTransport__client_type = client_type
trans._THeaderTransport__write_headers = headers
trans.set_protocol_id(protocol_type)
trans.write(input)
trans.flush()
prot_buf = TMemoryBuffer(write_buf.getvalue())
prot = THeaderProtocol(prot_buf, client_types=[client_type])
ctx = TCppConnectionContext(context_data)
ret = self.processor.process(prot, prot, ctx)
done_callback = partial(_ProcessorAdapter.done,
prot_buf=prot_buf,
client_type=client_type,
callback=callback)
# This future is created by and returned from the processor's
# ThreadPoolExecutor, which keeps a reference to it. So it is
# fine for this future to end its lifecycle here.
if isinstance(ret, Future):
ret.add_done_callback(lambda x, d=done_callback: d())
else:
done_callback()
except:
# Don't let exceptions escape back into C++
traceback.print_exc()
def __init__(self, trans, strictRead=False, client_types=None):
"""Create a THeaderProtocol instance
@param transport(TTransport) The underlying transport.
@param strictRead(bool) Turn on strictRead if using TBinaryProtocol
@param client_types([THeaderTransport.HEADERS_CLIENT_TYPE, ...])
List of client types to support. Defaults to
HEADERS_CLIENT_TYPE only.
"""
htrans = THeaderTransport(trans, client_types)
TProtocolBase.__init__(self, htrans)
self.strictRead = strictRead
self.reset_protocol()
def setUp(self):
"""Create two buffers, transports, and protocols.
self._h_trans uses THeaderTransport
self._f_trans uses TFuzzyHeaderTransport
"""
cls = self.__class__
# THeaderTransport attributes
self._h_buf = TMemoryBuffer()
self._h_trans = THeaderTransport(self._h_buf)
self._h_prot = THeaderProtocol(self._h_trans)
# TFuzzyHeaderTransport attributes
self._f_buf = TMemoryBuffer()
self._f_trans = TFuzzyHeaderTransport(
self._f_buf, fuzz_fields=cls.fuzz_fields,
fuzz_all_if_empty=False, verbose=False)
self._f_prot = THeaderProtocol(self._f_trans)
def call_processor(self, input, client_type, protocol_type,
context_data):
try:
# The input string has already had the header removed, but
# the python processor will expect it to be there. In
# order to reconstitute the message with headers, we use
# the THeaderProtocol object to write into a memory
# buffer, then pass that buffer to the python processor.
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf)
trans._THeaderTransport__client_type = client_type
trans.set_protocol_id(protocol_type)
trans.write(input)
trans.flush()
prot_buf = TMemoryBuffer(write_buf.getvalue())
prot = THeaderProtocol(prot_buf, client_types=[client_type])
ctx = TCppConnectionContext(context_data)
self.processor.process(prot, prot, ctx)
# Check for empty result. If so, return an empty string
# here. This is probably a oneway request, but we can't
# reliably tell. The C++ code does basically the same
# thing.
response = prot_buf.getvalue()
if len(response) == 0:
return response
# And on the way out, we need to strip off the header,
# because the C++ code will expect to add it.
read_buf = TMemoryBuffer(response)
trans = THeaderTransport(read_buf, client_types=[client_type])
trans.readFrame(len(response))
return trans.cstringio_buf.read()
except:
# Don't let exceptions escape back into C++
traceback.print_exc()
class FuzzyTransportTest(object):
"""Test class that sets up a THeaderTransport and a TFuzzyHeaderTransport.
Used for writing and comparing messages using both transports.
"""
fuzz_fields = []
# Sample TestService method arguments
sampleListStruct = ttypes.ListStruct(
a=[True, False],
b=[1, 2, 3],
c=[1.2, 3.4],
d=["ab", "cd"],
e=[list(sm.xrange(n)) for n in sm.xrange(20)],
f=[{1: 2}, {3: 4, 5: 6}],
g=[{"a", "b"}, {"c"}, set()]
)
def setUp(self):
"""Create two buffers, transports, and protocols.
self._h_trans uses THeaderTransport
self._f_trans uses TFuzzyHeaderTransport
"""
cls = self.__class__
# THeaderTransport attributes
self._h_buf = TMemoryBuffer()
self._h_trans = THeaderTransport(self._h_buf)
self._h_prot = THeaderProtocol(self._h_trans)
# TFuzzyHeaderTransport attributes
self._f_buf = TMemoryBuffer()
self._f_trans = TFuzzyHeaderTransport(
self._f_buf, fuzz_fields=cls.fuzz_fields,
fuzz_all_if_empty=False, verbose=False)
self._f_prot = THeaderProtocol(self._f_trans)
def writeObject(self, obj=sampleListStruct):
"""Write an object to the test and reference protocols.
Return the contents of both buffers.
"""
obj.write(self._h_prot)
obj.write(self._f_prot)
self._h_trans.flush()
self._f_trans.flush()
h_buf = self._h_buf.getvalue()
f_buf = self._f_buf.getvalue()
return h_buf, f_buf
def differentIndices(self, header, fuzzy):
"""Return a list of byte positions at which two messages' bytes differ.
Header should be the contents of self._h_buf
Fuzzy should be the contents of self._f_buf
"""
indices = []
for i, (h, f) in enumerate(itertools.izip(header, fuzzy)):
if h != f:
indices.append(i)
return indices
def assertEqualsExceptIndices(self, header, fuzzy, indices):
"""Assert that the buffers `header` and `fuzzy` are equal,
except possibly at the byte positions included in `indices`.
This ensures that the message produced by TFuzzyHeaderProtocol (fuzzy)
is equal to the message produced by THeaderProtocol (header), except
at the byte positions that are expected to be fuzzed."""
self.assertEquals(len(header), len(fuzzy))
for diff in self.differentIndices(header, fuzzy):
self.assertIn(diff, indices)
def round_robin(self, compress=None):
original = b'A' * MAX_FRAME_SIZE
mb = TMemoryBuffer()
trans = THeaderTransport(mb, client_type=CLIENT_TYPE.HEADER)
trans.set_max_frame_size(MAX_FRAME_SIZE + MIN_HEADER_SIZE)
if compress:
trans.add_transform(compress)
trans.write(original)
trans.flush()
frame = mb.getvalue()
# Cleanup the memory buffer
mb.close()
del mb
if compress is None:
# Partial Decode the frame and see if its correct size wise
sz = struct.unpack('!I', frame[:4])[0]
self.assertEqual(sz, BIG_FRAME_MAGIC)
sz = struct.unpack('!Q', frame[4:12])[0]
self.assertEqual(len(frame), sz + 12)
# Read it back
mb = TMemoryBuffer(frame)
trans = THeaderTransport(mb, client_type=CLIENT_TYPE.HEADER)
trans.set_max_frame_size(len(frame))
trans.readFrame(0)
result = trans.read(MAX_FRAME_SIZE)
mb.close()
del mb
self.assertEqual(result, original, 'round-robin different from original')
