Add Redis pub/sub instrumentation (#1129)

* Add instrumentation of redis pub/sub commands

* Add instrumentation of aioredis pub/sub commands

Co-authored-by: Benjamin Wohlwend <[email protected]>

Author: Sparkycz

elasticapm/instrumentation/packages/asyncio/aioredis.py

@@ -38,7 +38,8 @@
 class RedisConnectionPoolInstrumentation(AbstractInstrumentedModule):
     name = "aioredis"
 
-    instrument_list = [("aioredis.pool", "ConnectionsPool.execute")]
+    instrument_list = [("aioredis.pool", "ConnectionsPool.execute"),
+                       ("aioredis.pool", "ConnectionsPool.execute_pubsub")]
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         if len(args) > 0:
@@ -73,7 +74,8 @@ def call(self, module, method, wrapped, instance, args, kwargs):
 class RedisConnectionInstrumentation(AbstractInstrumentedModule):
     name = "aioredis"
 
-    instrument_list = (("aioredis.connection", "RedisConnection.execute"),)
+    instrument_list = (("aioredis.connection", "RedisConnection.execute"),
+                       ("aioredis.pool", "ConnectionsPool.execute_pubsub"))
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         span = execution_context.get_span()

elasticapm/instrumentation/packages/redis.py

@@ -53,7 +53,7 @@ class RedisInstrumentation(Redis3CheckMixin, AbstractInstrumentedModule):
     name = "redis"
 
     # no need to instrument StrictRedis in redis-py >= 3.0
-    instrument_list_3 = [("redis.client", "Redis.execute_command")]
+    instrument_list_3 = [("redis.client", "Redis.execute_command"), ("redis.client", "PubSub.execute_command")]
     instrument_list = [("redis.client", "Redis.execute_command"), ("redis.client", "StrictRedis.execute_command")]
 
     def call(self, module, method, wrapped, instance, args, kwargs):

tests/instrumentation/asyncio_tests/aioredis_tests.py

@@ -123,3 +123,48 @@ async def test_redis_client(instrument, elasticapm_client, redis_conn):
     assert spans[2]["type"] == "test"
 
     assert len(spans) == 3
+
+
+@pytest.mark.integrationtest
+async def test_publish_subscribe(instrument, elasticapm_client, redis_conn):
+    elasticapm_client.begin_transaction("transaction.test")
+    with capture_span("test_publish_subscribe", "test"):
+        # publish
+        await redis_conn.publish("mykey", "a")
+
+        #subscribe
+        await redis_conn.subscribe("mykey")
+
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+    expected_signatures = {"test_publish_subscribe", "PUBLISH", "SUBSCRIBE"}
+
+    assert {t["name"] for t in spans} == expected_signatures
+
+    assert spans[0]["name"] == "PUBLISH"
+    assert spans[0]["type"] == "db"
+    assert spans[0]["subtype"] == "redis"
+    assert spans[0]["action"] == "query"
+    assert spans[0]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[1]["name"] == "SUBSCRIBE"
+    assert spans[1]["type"] == "db"
+    assert spans[1]["subtype"] == "redis"
+    assert spans[1]["action"] == "query"
+    assert spans[1]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "aioredis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[2]["name"] == "test_publish_subscribe"
+    assert spans[2]["type"] == "test"
+
+    assert len(spans) == 3

tests/instrumentation/redis_tests.py

@@ -163,3 +163,49 @@ def test_unix_domain_socket_connection_destination_info():
     destination_info = get_destination_info(conn)
     assert destination_info["port"] is None
     assert destination_info["address"] == "unix:///some/path"
+
+
+@pytest.mark.integrationtest
+def test_publish_subscribe(instrument, elasticapm_client, redis_conn):
+    elasticapm_client.begin_transaction("transaction.test")
+    with capture_span("test_publish_subscribe", "test"):
+        # publish
+        redis_conn.publish("mykey", "a")
+
+        # subscribe
+        with redis_conn.pubsub() as channel:
+            channel.subscribe("mykey")
+
+    elasticapm_client.end_transaction("MyView")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    spans = elasticapm_client.spans_for_transaction(transactions[0])
+
+    expected_signatures = {"test_publish_subscribe", "PUBLISH", "SUBSCRIBE"}
+
+    assert {t["name"] for t in spans} == expected_signatures
+
+    assert spans[0]["name"] == "PUBLISH"
+    assert spans[0]["type"] == "db"
+    assert spans[0]["subtype"] == "redis"
+    assert spans[0]["action"] == "query"
+    assert spans[0]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "redis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[1]["name"] == "SUBSCRIBE"
+    assert spans[1]["type"] == "db"
+    assert spans[1]["subtype"] == "redis"
+    assert spans[1]["action"] == "query"
+    assert spans[1]["context"]["destination"] == {
+        "address": os.environ.get("REDIS_HOST", "localhost"),
+        "port": int(os.environ.get("REDIS_PORT", 6379)),
+        "service": {"name": "redis", "resource": "redis", "type": "db"},
+    }
+
+    assert spans[2]["name"] == "test_publish_subscribe"
+    assert spans[2]["type"] == "test"
+
+    assert len(spans) == 3