Support capturing spans from code run in an event loop executor

elasticapm/instrumentation/packages/asyncio/base.py

@@ -29,8 +29,33 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
+from elasticapm.traces import execution_context
+from elasticapm.utils.wrapt import function_wrapper
 
 
 class AsyncAbstractInstrumentedModule(AbstractInstrumentedModule):
  async def call(self, module, method, wrapped, instance, args, kwargs):
  raise NotImplementedError()
+
+
+class RunInExecutorInstrumentation(AbstractInstrumentedModule):
+ name = "executor"
+ instrument_list = [("asyncio.base_events", "BaseEventLoop.run_in_executor")]
+
+ def call(self, module, method, wrapped, instance, args, kwargs):
+ func = args[1]
+ transaction = execution_context.get_transaction()
+ span = execution_context.get_span()
+
+ @function_wrapper
+ def wrapper(wrapped, instance, args, kwargs):
+ execution_context.set_transaction(transaction)
+ execution_context.set_span(span)
+ try:
+ return wrapped(*args, **kwargs)
+ finally:
+ execution_context.get_transaction(clear=True)
+
+ func = wrapper(func)
+ args = (args[0], func) + args[2:]
+ return wrapped(*args, **kwargs)

elasticapm/instrumentation/register.py

@@ -80,6 +80,7 @@
  "elasticapm.instrumentation.packages.tornado.TornadoRequestExecuteInstrumentation",
  "elasticapm.instrumentation.packages.tornado.TornadoHandleRequestExceptionInstrumentation",
  "elasticapm.instrumentation.packages.tornado.TornadoRenderInstrumentation",
+ "elasticapm.instrumentation.packages.asyncio.base.RunInExecutorInstrumentation",
  "elasticapm.instrumentation.packages.asyncio.httpcore.HTTPCoreAsyncInstrumentation",
  "elasticapm.instrumentation.packages.asyncio.aioredis.RedisConnectionPoolInstrumentation",
  "elasticapm.instrumentation.packages.asyncio.aioredis.RedisPipelineInstrumentation",

tests/instrumentation/asyncio_tests/base_tests.py

@@ -28,7 +28,9 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import threading
 from asyncio import tasks
+from concurrent import futures
 
 import pytest
 
@@ -64,3 +66,20 @@ async def do_some_work():
  assert d["subtype"] == "custom"
  assert not d["sync"]
  assert d["transaction_id"] == transaction["id"]
+
+
+async def test_threadpool_executor(instrument, event_loop, elasticapm_client):
+ executor = futures.ThreadPoolExecutor(max_workers=2)
+ elasticapm_client.begin_transaction("foo")
+
+ def myfunc(a, b):
+ with elasticapm.capture_span("nothing", extra={"thread_id": threading.get_ident()}):
+ pass
+
+ await event_loop.run_in_executor(executor, myfunc, 1, 2)
+ elasticapm_client.end_transaction("bar", "OK")
+ transaction = elasticapm_client.events[constants.TRANSACTION][0]
+ spans = elasticapm_client.spans_for_transaction(transaction)
+ assert len(spans) == 1
+ assert spans[0]["name"] == "nothing"
+ assert spans[0]["context"]["thread_id"] != threading.get_ident()