properly await in asyncio redis instrumentation

Author: beniwohli

elasticapm/instrumentation/packages/asyncio/redis_asyncio.py

@@ -31,49 +31,49 @@
 from __future__ import absolute_import
 
 from elasticapm.contrib.asyncio.traces import async_capture_span
-from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
+from elasticapm.instrumentation.packages.asyncio.base import AsyncAbstractInstrumentedModule
 
 
-class RedisAsyncioInstrumentation(AbstractInstrumentedModule):
+class RedisAsyncioInstrumentation(AsyncAbstractInstrumentedModule):
  name = "redis"
 
  instrument_list = [
  ("redis.asyncio.client", "Redis.execute_command"),
  ("redis.asyncio.client", "PubSub.execute_command"),
  ]
 
- def call(self, module, method, wrapped, instance, args, kwargs):
+ async def call(self, module, method, wrapped, instance, args, kwargs):
  if len(args) > 0:
  wrapped_name = args[0]
  if isinstance(wrapped_name, bytes):
  wrapped_name = wrapped_name.decode()
  else:
  wrapped_name = self.get_wrapped_name(wrapped, instance, method)
 
- with async_capture_span(
+ async with async_capture_span(
  wrapped_name, span_type="db", span_subtype="redis", span_action="query", leaf=True
  ) as span:
  if span.context is not None:
  span.context["destination"] = _get_destination_info(instance)
 
- return wrapped(*args, **kwargs)
+ return await wrapped(*args, **kwargs)
 
 
-class RedisPipelineInstrumentation(AbstractInstrumentedModule):
+class RedisPipelineInstrumentation(AsyncAbstractInstrumentedModule):
  name = "redis"
 
  instrument_list = [("redis.asyncio.client", "Pipeline.execute")]
 
- def call(self, module, method, wrapped, instance, args, kwargs):
+ async def call(self, module, method, wrapped, instance, args, kwargs):
  wrapped_name = self.get_wrapped_name(wrapped, instance, method)
 
- with async_capture_span(
+ async with async_capture_span(
  wrapped_name, span_type="db", span_subtype="redis", span_action="query", leaf=True
  ) as span:
  if span.context is not None:
  span.context["destination"] = _get_destination_info(instance)
 
- return wrapped(*args, **kwargs)
+ return await wrapped(*args, **kwargs)
 
 
 def _get_destination_info(connection):

tests/instrumentation/asyncio_tests/redis_asyncio_tests.py

@@ -68,6 +68,7 @@ async def test_ping(instrument, elasticapm_client, redis_async_conn):
  transaction = elasticapm_client.events[TRANSACTION][0]
  span = elasticapm_client.spans_for_transaction(transaction)[0]
  assert span["name"] == "PING"
+ assert span["duration"] > 0.2 # sanity test to ensure we measure the actual call
 
 
 @pytest.mark.integrationtest
@@ -92,6 +93,7 @@ async def test_pipeline(instrument, elasticapm_client, redis_async_conn):
  "port": int(os.environ.get("REDIS_PORT", 6379)),
  "service": {"name": "", "resource": "redis", "type": ""},
  }
+ assert spans[0]["duration"] > 0.2 # sanity test to ensure we measure the actual call
 
  assert spans[1]["name"] == "test_pipeline"
  assert spans[1]["type"] == "test"
@@ -123,6 +125,7 @@ async def test_redis_client(instrument, elasticapm_client, redis_async_conn):
  "port": int(os.environ.get("REDIS_PORT", 6379)),
  "service": {"name": "", "resource": "redis", "type": ""},
  }
+ assert spans[0]["duration"] > 0.2 # sanity test to ensure we measure the actual call
 
  assert spans[1]["name"] == "RPUSH"
  assert spans[1]["type"] == "db"
@@ -133,6 +136,7 @@ async def test_redis_client(instrument, elasticapm_client, redis_async_conn):
  "port": int(os.environ.get("REDIS_PORT", 6379)),
  "service": {"name": "", "resource": "redis", "type": ""},
  }
+ assert spans[1]["duration"] > 0.2 # sanity test to ensure we measure the actual call
 
  assert spans[2]["name"] == "test_redis_client"
  assert spans[2]["type"] == "test"
@@ -169,6 +173,7 @@ async def test_publish_subscribe_async(instrument, elasticapm_client, redis_asyn
  "port": int(os.environ.get("REDIS_PORT", 6379)),
  "service": {"name": "", "resource": "redis", "type": ""},
  }
+ assert spans[0]["duration"] > 0.2 # sanity test to ensure we measure the actual call
 
  assert spans[1]["name"] == "SUBSCRIBE"
  assert spans[1]["type"] == "db"
@@ -179,6 +184,7 @@ async def test_publish_subscribe_async(instrument, elasticapm_client, redis_asyn
  "port": int(os.environ.get("REDIS_PORT", 6379)),
  "service": {"name": "", "resource": "redis", "type": ""},
  }
+ assert spans[1]["duration"] > 0.2 # sanity test to ensure we measure the actual call
 
  assert spans[2]["name"] == "test_publish_subscribe"
  assert spans[2]["type"] == "test"