def test_expected_exception_not_passed_to_server_span_finish(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocol(client_memory_trans)
client = TestService.Client(client_prot)
try:
client.example_throws(crash=False)
except TTransportException:
pass # we don't have a test response for the client
self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())
self.processor.process(self.iprot, self.oprot, self.server_context)
self.assertEqual(self.server_observer.on_start.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_count, 1)
self.assertEqual(self.server_observer.on_finish.call_args[0], (None, ))
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 __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 create_client(client_klass, host=None, port=None, client_type=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)
protocol = THeaderProtocol(
sock,
client_types=client_type, # We accept the same as our inital send_
client_type=client_type # Used for the inital send_
)
sock.open()
return client_klass(protocol)
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 test_client_proxy_flow(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocol(client_memory_trans)
class Pool(object):
@contextlib.contextmanager
def connection(self):
yield client_prot
client_factory = ThriftContextFactory(Pool(), TestService.Client)
span = mock.MagicMock()
child_span = span.make_child().__enter__()
child_span.trace_id = 1
child_span.parent_id = 1
child_span.id = 1
child_span.sampled = True
child_span.flags = None
edge_context = self.edge_context_factory.from_upstream(
SERIALIZED_EDGECONTEXT_WITH_VALID_AUTH)
edge_context.attach_context(child_span.context)
client = client_factory.make_object_for_context("test", span)
try:
client.example_simple()
except TTransportException:
pass # we don't have a test response for the client
self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())
self.processor.process(self.iprot, self.oprot, self.server_context)
context, _ = self.observer.on_server_span_created.call_args[0]
try:
self.assertEqual(context.request_context.user.id, "t2_example")
self.assertEqual(context.request_context.user.roles, set())
self.assertEqual(context.request_context.user.is_logged_in, True)
self.assertEqual(context.request_context.user.loid, "t2_deadbeef")
self.assertEqual(context.request_context.user.cookie_created_ms,
100000)
self.assertEqual(context.request_context.oauth_client.id, None)
self.assertFalse(
context.request_context.oauth_client.is_type("third_party"))
self.assertEqual(context.request_context.session.id, "beefdead")
except jwt.exceptions.InvalidAlgorithmError:
raise unittest.SkipTest("cryptography is not installed")
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 test_no_headers(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocol(client_memory_trans)
client = BaseplateService.Client(client_prot)
try:
client.is_healthy()
except:
pass # we don't have a test response for the client
self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())
self.processor.process(self.iprot, self.oprot, self.server_context)
self.assertEqual(self.observer.on_root_span_created.call_count, 1)
context, root_span = self.observer.on_root_span_created.call_args[0]
self.assertEqual(root_span.trace_id, "no-trace")
self.assertEqual(root_span.parent_id, "no-parent")
self.assertEqual(root_span.id, "no-span")
self.assertTrue(self.root_observer.on_start.called)
self.assertTrue(self.root_observer.on_stop.called)
def message_received(self, frame):
# We support the deprecated FRAMED transport for old fb303
# clients that were otherwise failing miserably.
client_types = {
THeaderTransport.HEADERS_CLIENT_TYPE,
THeaderTransport.FRAMED_DEPRECATED,
}
tm = TMemoryBuffer(frame)
prot = THeaderProtocol(tm, client_types=client_types)
try:
yield from self.processor.process(
prot, prot, self.server_context,
)
msg = tm.getvalue()
if len(msg) > 0:
self.transport.write(msg)
except Exception:
logger.exception("Exception while processing request")
self.transport.close()
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()
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 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 message_received(self, frame, delay=0):
tmi = TReadOnlyBuffer(frame)
iprot = THeaderProtocol(tmi)
(fname, mtype, rseqid) = iprot.readMessageBegin()
if delay:
yield from asyncio.sleep(delay)
else:
try:
timeout_task = self.pending_tasks.pop(rseqid)
except KeyError:
# Task doesn't have a timeout or has already been cancelled
# and pruned from `pending_tasks`.
pass
else:
timeout_task.cancel()
method = getattr(self.client, "recv_" + fname.decode(), None)
if method is None:
logger.error("Method %r is not supported", method)
self.transport.abort()
else:
method(iprot, mtype, rseqid)
def test_no_trace_headers(self, getrandbits):
getrandbits.return_value = 1234
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocol(client_memory_trans)
client = BaseplateService.Client(client_prot)
try:
client.is_healthy()
except:
pass # we don't have a test response for the client
self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())
self.processor.process(self.iprot, self.oprot, self.server_context)
self.assertEqual(self.observer.on_server_span_created.call_count, 1)
context, server_span = self.observer.on_server_span_created.call_args[0]
self.assertEqual(server_span.trace_id, 1234)
self.assertEqual(server_span.parent_id, None)
self.assertEqual(server_span.id, 1234)
self.assertTrue(self.server_observer.on_start.called)
self.assertTrue(self.server_observer.on_finish.called)
def test_edge_request_headers(self):
client_memory_trans = TMemoryBuffer()
client_prot = THeaderProtocol(client_memory_trans)
client_header_trans = client_prot.trans
client_header_trans.set_header("Authentication", self.VALID_TOKEN)
client_header_trans.set_header("Edge-Request",
self.SERIALIZED_REQUEST_HEADER)
client_header_trans.set_header("Trace", "1234")
client_header_trans.set_header("Parent", "2345")
client_header_trans.set_header("Span", "3456")
client_header_trans.set_header("Sampled", "1")
client_header_trans.set_header("Flags", "1")
client = TestService.Client(client_prot)
try:
client.example_simple()
except TTransportException:
pass # we don't have a test response for the client
self.itrans._readBuffer = StringIO(client_memory_trans.getvalue())
self.processor.process(self.iprot, self.oprot, self.server_context)
context, _ = self.observer.on_server_span_created.call_args[0]
try:
self.assertEqual(context.request_context.user.id, "test_user_id")
self.assertEqual(context.request_context.user.roles, set())
self.assertEqual(context.request_context.user.is_logged_in, True)
self.assertEqual(context.request_context.user.loid, "t2_deadbeef")
self.assertEqual(context.request_context.user.cookie_created_ms,
100000)
self.assertEqual(context.request_context.oauth_client.id, None)
self.assertFalse(
context.request_context.oauth_client.is_type("third_party"))
self.assertEqual(context.request_context.session.id, "beefdead")
except jwt.exceptions.InvalidAlgorithmError:
raise unittest.SkipTest("cryptography is not installed")
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()
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.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.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: # noqa
# Don't let exceptions escape back into C++
traceback.print_exc()
