def invoke(self, comp, arg):
comp.type_signature.check_function()
# Save the type signature before compiling. Compilation currently loses
# container types, so we must remember them here so that they can be
# restored in the output.
result_type = comp.type_signature.result
if self._compiler_pipeline is not None:
with tracing.span('ExecutionContext', 'Compile', span=True):
comp = self._compiler_pipeline.compile(comp)
with tracing.span('ExecutionContext', 'Invoke', span=True):
@contextlib.contextmanager
def executor_closer(ex_factory, cardinalities):
"""Wraps an Executor into a closeable resource."""
# TODO(b/168744510): The lifecycles embedded here are confusing; unify
# or clarify the need for them.
ex = ex_factory.create_executor(cardinalities)
try:
yield ex
except Exception as e:
ex_factory.clean_up_executors()
raise e
finally:
ex.close()
if arg is not None:
py_typecheck.check_type(arg, ExecutionContextValue)
unwrapped_arg = _unwrap_execution_context_value(arg)
cardinalities = cardinalities_utils.infer_cardinalities(
unwrapped_arg, arg.type_signature)
else:
cardinalities = {}
with executor_closer(self._executor_factory, cardinalities) as executor:
py_typecheck.check_type(executor, executor_base.Executor)
def get_event_loop():
new_loop = asyncio.new_event_loop()
new_loop.set_task_factory(
tracing.propagate_trace_context_task_factory)
return new_loop
event_loop = get_event_loop()
if arg is not None:
arg = event_loop.run_until_complete(
tracing.wrap_coroutine_in_current_trace_context(
_ingest(executor, unwrapped_arg, arg.type_signature)))
return event_loop.run_until_complete(
tracing.wrap_coroutine_in_current_trace_context(
_invoke(executor, comp, arg, result_type)))
def test_parenting_non_async_to_async_to_nested_async(self):
mock = set_mock_trace()
loop = asyncio.new_event_loop()
loop.set_task_factory(tracing.propagate_trace_context_task_factory)
def run_loop():
loop.run_forever()
loop.close()
thread = threading.Thread(target=functools.partial(run_loop),
daemon=True)
thread.start()
@tracing.trace
async def middle():
with tracing.span('inner', ''):
pass
with tracing.span('outer', ''):
# This sends the coroutine over to another thread,
# keeping the current trace context.
coro_with_trace_ctx = tracing.wrap_coroutine_in_current_trace_context(
middle())
asyncio.run_coroutine_threadsafe(coro_with_trace_ctx,
loop).result()
loop.call_soon_threadsafe(loop.stop)
thread.join()
self.assertEqual(mock.parent_span_yields, [None, 0, 1])
self.assertEqual(mock.scopes, ['outer', '<locals>', 'inner'])
self.assertEqual(mock.sub_scopes, ['', 'middle', ''])
def invoke(self, comp, arg):
comp.type_signature.check_function()
# Save the type signature before compiling. Compilation currently loses
# container types, so we must remember them here so that they can be
# restored in the output.
result_type = comp.type_signature.result
if self._compiler_pipeline is not None:
with tracing.span('ExecutionContext', 'Compile', span=True):
comp = self._compiler_pipeline.compile(comp)
with tracing.span('ExecutionContext', 'Invoke', span=True):
@contextlib.contextmanager
def executor_closer(ex_factory, cardinalities):
"""Wraps an Executor into a closeable resource."""
ex = ex_factory.create_executor(cardinalities)
try:
yield ex
except Exception as e:
ex_factory.clean_up_executors()
raise e
if arg is not None:
py_typecheck.check_type(arg, ExecutionContextValue)
unwrapped_arg = _unwrap_execution_context_value(arg)
cardinalities = cardinalities_utils.infer_cardinalities(
unwrapped_arg, arg.type_signature)
else:
cardinalities = {}
with executor_closer(self._executor_factory,
cardinalities) as executor:
py_typecheck.check_type(executor, executor_base.Executor)
if arg is not None:
arg = self._event_loop.run_until_complete(
tracing.wrap_coroutine_in_current_trace_context(
_ingest(executor, unwrapped_arg,
arg.type_signature)))
return self._event_loop.run_until_complete(
tracing.wrap_coroutine_in_current_trace_context(
_invoke(executor, comp, arg, result_type)))
def invoke(self, comp, arg):
with tracing.span('ExecutionContext', 'Invoke', span=True):
@contextlib.contextmanager
def executor_closer(wrapped_executor):
"""Wraps an Executor into a closeable resource."""
try:
yield wrapped_executor
finally:
wrapped_executor.close()
if arg is not None:
py_typecheck.check_type(arg, ExecutionContextValue)
unwrapped_arg = _unwrap_execution_context_value(arg)
cardinalities = cardinalities_utils.infer_cardinalities(
unwrapped_arg, arg.type_signature)
else:
cardinalities = {}
with executor_closer(
self._executor_factory.create_executor(
cardinalities)) as executor:
py_typecheck.check_type(executor, executor_base.Executor)
def get_event_loop():
new_loop = asyncio.new_event_loop()
new_loop.set_task_factory(
tracing.propagate_trace_context_task_factory)
return new_loop
event_loop = get_event_loop()
if arg is not None:
arg = event_loop.run_until_complete(
tracing.wrap_coroutine_in_current_trace_context(
_ingest(executor, unwrapped_arg,
arg.type_signature)))
return event_loop.run_until_complete(
tracing.wrap_coroutine_in_current_trace_context(
_invoke(executor, comp, arg)))
def _run_coro_threadsafe_with_tracing(self, coro):
"""Runs `coro` on `self._event_loop` inside the current trace spans."""
with tracing.with_trace_context_from_rpc():
return asyncio.run_coroutine_threadsafe(
tracing.wrap_coroutine_in_current_trace_context(coro),
self._event_loop)
def _delegate(self, coro):
coro_with_trace_ctx = tracing.wrap_coroutine_in_current_trace_context(
coro)
return asyncio.wrap_future(
asyncio.run_coroutine_threadsafe(coro_with_trace_ctx,
self._event_loop))
def _delegate_with_trace_ctx(coro, async_runner): coro_with_trace_ctx = tracing.wrap_coroutine_in_current_trace_context(coro) return async_runner.await_coro_and_return_result(coro_with_trace_ctx)
