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 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")
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 _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 __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 call_processor(self, input, headers, 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._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)
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()
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, 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 __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=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, 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)
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 __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 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 _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)
if options.headers is not None:
try:
parsed_headers = eval(options.headers)
except Exception:
self._exit(
error_message='Request headers (--headers) argument'
' failed eval')
if not isinstance(parsed_headers, dict):
self._exit(
error_message='Request headers (--headers) argument'
' must evaluate to a dict')
for header_name, header_value in parsed_headers.items():
transport.set_header(header_name, header_value)
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 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.set_max_frame_size(MAX_BIG_FRAME_SIZE)
trans.readFrame(len(response))
callback.call(trans.cstringio_buf.read())
except: # noqa
# Don't let exceptions escape back into C++
traceback.print_exc()
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 _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.
transport = THeaderTransport(socket)
if options.headers is not None:
try:
parsed_headers = eval(options.headers)
except Exception:
self._exit(
error_message="Request headers (--headers) argument"
" failed eval")
if not isinstance(parsed_headers, dict):
self._exit(
error_message="Request headers (--headers) argument"
" must evaluate to a dict")
for header_name, header_value in parsed_headers.items():
transport.set_header(header_name, header_value)
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 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 __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 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 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 call_processor(
self, input, headers, client_type, protocol_type, context_data, callback
):
try:
# TCppServer threads are not created by Python so they are
# missing settrace() hooks. We need to manually set the
# hook here for things to work (e.g. coverage and pdb).
if sys.gettrace() is None and threading._trace_hook is not None:
sys.settrace(threading._trace_hook)
# 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()
if self.observer and should_sample:
timestamps.setProcessBeginNow()
write_buf = TMemoryBuffer()
trans = THeaderTransport(write_buf)
trans.set_max_frame_size(MAX_BIG_FRAME_SIZE)
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])
prot.trans.set_max_frame_size(MAX_BIG_FRAME_SIZE)
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.setProcessEndNow()
# 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: # noqa
# Don't let exceptions escape back into C++
traceback.print_exc()
