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)