Fix crashes in aioredis transactions (#633)

* Fix aioredis wrappers for transactions

* Refactor aioredis test setup

* Expand aioredis tox testing

* Add no harm tests for aioredis transactions

Co-authored-by: Lalleh Rafeei <[email protected]>

Author: TimPansino

newrelic/hooks/datastore_aioredis.py

@@ -12,16 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from newrelic.api.datastore_trace import DatastoreTrace
+from newrelic.api.datastore_trace import DatastoreTrace, DatastoreTraceWrapper
 from newrelic.api.time_trace import current_trace
 from newrelic.api.transaction import current_transaction
-from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.common.object_wrapper import wrap_function_wrapper, function_wrapper, FunctionWrapper
 from newrelic.hooks.datastore_redis import (
     _redis_client_methods,
     _redis_multipart_commands,
     _redis_operation_re,
 )
 
+from newrelic.common.async_wrapper import async_wrapper
+
+import aioredis
+
+try:
+    AIOREDIS_VERSION = tuple(int(x) for x in getattr(aioredis, "__version__").split("."))
+except Exception:
+    AIOREDIS_VERSION = (0, 0, 0)
+
 
 def _conn_attrs_to_dict(connection):
     host = getattr(connection, "host", None)
@@ -45,13 +54,36 @@ def _instance_info(kwargs):
 
 
 def _wrap_AioRedis_method_wrapper(module, instance_class_name, operation):
-    async def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
+
+    @function_wrapper
+    async def _nr_wrapper_AioRedis_async_method_(wrapped, instance, args, kwargs):
         transaction = current_transaction()
         if transaction is None:
             return await wrapped(*args, **kwargs)
 
         with DatastoreTrace(product="Redis", target=None, operation=operation):
             return await wrapped(*args, **kwargs)
+    
+    def _nr_wrapper_AioRedis_method_(wrapped, instance, args, kwargs):
+        # Check for transaction and return early if found.
+        # Method will return synchronously without executing,
+        # it will be added to the command stack and run later.
+        if AIOREDIS_VERSION < (2,):
+            # AioRedis v1 uses a RedisBuffer instead of a real connection for queueing up pipeline commands
+            from aioredis.commands.transaction import _RedisBuffer
+            if isinstance(instance._pool_or_conn, _RedisBuffer):
+                # Method will return synchronously without executing,
+                # it will be added to the command stack and run later.
+                return wrapped(*args, **kwargs)
+        else:
+            # AioRedis v2 uses a Pipeline object for a client and internally queues up pipeline commands
+            from aioredis.client import Pipeline
+            if isinstance(instance, Pipeline):
+                return wrapped(*args, **kwargs)
+
+        # Method should be run when awaited, therefore we wrap in an async wrapper.
+        return _nr_wrapper_AioRedis_async_method_(wrapped)(*args, **kwargs)
+
 
     name = "%s.%s" % (instance_class_name, operation)
     wrap_function_wrapper(module, name, _nr_wrapper_AioRedis_method_)
@@ -108,6 +140,58 @@ async def wrap_Connection_send_command(wrapped, instance, args, kwargs):
         return await wrapped(*args, **kwargs)
 
 
+def wrap_RedisConnection_execute(wrapped, instance, args, kwargs):
+    # RedisConnection in aioredis v1 returns a future instead of using coroutines
+    transaction = current_transaction()
+    if not transaction:
+        return wrapped(*args, **kwargs)
+
+    host, port_path_or_id, db = (None, None, None)
+
+    try:
+        dt = transaction.settings.datastore_tracer
+        if dt.instance_reporting.enabled or dt.database_name_reporting.enabled:
+            conn_kwargs = _conn_attrs_to_dict(instance)
+            host, port_path_or_id, db = _instance_info(conn_kwargs)
+    except Exception:
+        pass
+
+    # Older Redis clients would when sending multi part commands pass
+    # them in as separate arguments to send_command(). Need to therefore
+    # detect those and grab the next argument from the set of arguments.
+
+    operation = args[0].strip().lower()
+
+    # If it's not a multi part command, there's no need to trace it, so
+    # we can return early.
+
+    if operation.split()[0] not in _redis_multipart_commands:        # Set the datastore info on the DatastoreTrace containing this function call.
+        trace = current_trace()
+
+        # Find DatastoreTrace no matter how many other traces are inbetween
+        while trace is not None and not isinstance(trace, DatastoreTrace):
+            trace = getattr(trace, "parent", None)
+
+        if trace is not None:
+            trace.host = host
+            trace.port_path_or_id = port_path_or_id
+            trace.database_name = db
+
+        return wrapped(*args, **kwargs)
+
+    # Convert multi args to single arg string
+
+    if operation in _redis_multipart_commands and len(args) > 1:
+        operation = "%s %s" % (operation, args[1].strip().lower())
+
+    operation = _redis_operation_re.sub("_", operation)
+
+    with DatastoreTrace(
+        product="Redis", target=None, operation=operation, host=host, port_path_or_id=port_path_or_id, database_name=db
+    ):
+        return wrapped(*args, **kwargs)
+
+
 def instrument_aioredis_client(module):
     # StrictRedis is just an alias of Redis, no need to wrap it as well.
     if hasattr(module, "Redis"):
@@ -124,4 +208,4 @@ def instrument_aioredis_connection(module):
 
     if hasattr(module, "RedisConnection"):
         if hasattr(module.RedisConnection, "execute"):
-            wrap_function_wrapper(module, "RedisConnection.execute", wrap_Connection_send_command)
+            wrap_function_wrapper(module, "RedisConnection.execute", wrap_RedisConnection_execute)

tests/datastore_aioredis/conftest.py

@@ -13,16 +13,22 @@
 # limitations under the License.
 
 import aioredis
-import asyncio
 import pytest
 
+from testing_support.db_settings import redis_settings
+
+from testing_support.fixture.event_loop import event_loop as loop
 from testing_support.fixtures import (  # noqa: F401
     code_coverage_fixture,
     collector_agent_registration_fixture,
     collector_available_fixture,
 )
 
 AIOREDIS_VERSION = tuple(int(x) for x in aioredis.__version__.split(".")[:2])
+SKIPIF_AIOREDIS_V1 = pytest.mark.skipif(AIOREDIS_VERSION < (2,), reason="Unsupported aioredis version.")
+SKIPIF_AIOREDIS_V2 = pytest.mark.skipif(AIOREDIS_VERSION >= (2,), reason="Unsupported aioredis version.")
+DB_SETTINGS = redis_settings()[0]
+
 
 _coverage_source = [
     "newrelic.hooks.datastore_aioredis",
@@ -45,10 +51,19 @@
 )
 
 
-event_loop = asyncio.get_event_loop()
-asyncio.set_event_loop(event_loop)
-
-
-@pytest.fixture()
-def loop():
-    yield event_loop
+@pytest.fixture(params=("Redis", "StrictRedis"))
+def client(request, loop):
+    if AIOREDIS_VERSION >= (2, 0):
+        if request.param == "Redis":
+            return aioredis.Redis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0)
+        elif request.param == "StrictRedis":
+            return aioredis.StrictRedis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0)
+        else:
+            raise NotImplementedError()
+    else:
+        if request.param == "Redis":
+            return loop.run_until_complete(aioredis.create_redis("redis://%s:%d" % (DB_SETTINGS["host"], DB_SETTINGS["port"]), db=0))
+        elif request.param == "StrictRedis":
+            pytest.skip("StrictRedis not implemented.")
+        else:
+            raise NotImplementedError()

tests/datastore_aioredis/test_custom_conn_pool.py

@@ -17,12 +17,6 @@
 will not result in an error.
 """
 
-import asyncio
-import pytest
-import aioredis
-
-from conftest import event_loop, loop, AIOREDIS_VERSION
-
 from newrelic.api.background_task import background_task
 
 # from testing_support.fixture.event_loop import event_loop as loop
@@ -43,7 +37,7 @@ async def get_connection(self, name=None, *keys, **options):
         return self.connection
 
     async def release(self, connection):
-        self.connection.disconnect()
+        await self.connection.disconnect()
 
     async def execute(self, *args, **kwargs):
         return await self.connection.execute(*args, **kwargs)
@@ -105,18 +99,6 @@ async def exercise_redis(client):
         await client.execute("CLIENT", "LIST")
 
 
-if AIOREDIS_VERSION >= (2, 0):
-    clients = [
-        aioredis.Redis(host=DB_SETTINGS["host"], port=_port, db=0),
-        aioredis.StrictRedis(host=DB_SETTINGS["host"], port=_port, db=0),
-    ]
-else:
-    clients = [
-        event_loop.run_until_complete(aioredis.create_redis("redis://%s:%d" % (DB_SETTINGS["host"], _port), db=0)),
-    ]
-
-
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_enable_instance_settings)
 @validate_transaction_metrics(
     "test_custom_conn_pool:test_fake_conn_pool_enable_instance",
@@ -125,7 +107,7 @@ async def exercise_redis(client):
     background_task=True,
 )
 @background_task()
-def test_fake_conn_pool_enable_instance(client, loop):
+def test_fake_conn_pool_enable_instance(client, loop, monkeypatch):
     # Get a real connection
     conn = getattr(client, "_pool_or_conn", None)
     if conn is None:
@@ -135,14 +117,13 @@ def test_fake_conn_pool_enable_instance(client, loop):
     # have the `connection_kwargs` attribute.
 
     fake_pool = FakeConnectionPool(conn)
-    client.connection_pool = fake_pool
-    client._pool_or_conn = fake_pool
+    monkeypatch.setattr(client, "connection_pool", fake_pool, raising=False)
+    monkeypatch.setattr(client, "_pool_or_conn", fake_pool, raising=False)
     assert not hasattr(client.connection_pool, "connection_kwargs")
 
     loop.run_until_complete(exercise_redis(client))
 
 
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_disable_instance_settings)
 @validate_transaction_metrics(
     "test_custom_conn_pool:test_fake_conn_pool_disable_instance",
@@ -151,15 +132,18 @@ def test_fake_conn_pool_enable_instance(client, loop):
     background_task=True,
 )
 @background_task()
-def test_fake_conn_pool_disable_instance(client, loop):
+def test_fake_conn_pool_disable_instance(client, loop, monkeypatch):
     # Get a real connection
-    conn = loop.run_until_complete(client.connection_pool.get_connection("GET"))
+    conn = getattr(client, "_pool_or_conn", None)
+    if conn is None:
+        conn = loop.run_until_complete(client.connection_pool.get_connection("GET"))
 
     # Replace the original connection pool with one that doesn't
     # have the `connection_kwargs` attribute.
 
     fake_pool = FakeConnectionPool(conn)
-    client.connection_pool = fake_pool
+    monkeypatch.setattr(client, "connection_pool", fake_pool, raising=False)
+    monkeypatch.setattr(client, "_pool_or_conn", fake_pool, raising=False)
     assert not hasattr(client.connection_pool, "connection_kwargs")
 
     loop.run_until_complete(exercise_redis(client))

tests/datastore_aioredis/test_execute_command.py

@@ -13,11 +13,10 @@
 # limitations under the License.
 
 import pytest
-import aioredis
 from newrelic.api.background_task import background_task
 
 from testing_support.fixtures import validate_transaction_metrics, override_application_settings
-from conftest import event_loop, loop, AIOREDIS_VERSION
+from conftest import AIOREDIS_VERSION
 from testing_support.db_settings import redis_settings
 from testing_support.util import instance_hostname
 
@@ -70,19 +69,7 @@ async def exercise_redis_single_arg(client):
     await client.execute_command("CLIENT LIST")
 
 
-if AIOREDIS_VERSION >= (2, 0):
-    clients = [
-        aioredis.Redis(host=DB_SETTINGS["host"], port=_port, db=0),
-        aioredis.StrictRedis(host=DB_SETTINGS["host"], port=_port, db=0),
-    ]
-else:
-    clients = [
-        event_loop.run_until_complete(aioredis.create_redis("redis://%s:%d" % (DB_SETTINGS["host"], _port), db=0)),
-    ]
-
-
 @SKIP_IF_AIOREDIS_V1
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_enable_instance_settings)
 @validate_transaction_metrics(
     "test_execute_command:test_redis_execute_command_as_one_arg_enable",
@@ -96,7 +83,6 @@ def test_redis_execute_command_as_one_arg_enable(client, loop):
 
 
 @SKIP_IF_AIOREDIS_V1
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_disable_instance_settings)
 @validate_transaction_metrics(
     "test_execute_command:test_redis_execute_command_as_one_arg_disable",
@@ -109,7 +95,6 @@ def test_redis_execute_command_as_one_arg_disable(client, loop):
     loop.run_until_complete(exercise_redis_single_arg(client))
 
 
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_enable_instance_settings)
 @validate_transaction_metrics(
     "test_execute_command:test_redis_execute_command_as_two_args_enable",
@@ -122,7 +107,6 @@ def test_redis_execute_command_as_two_args_enable(client, loop):
     loop.run_until_complete(exercise_redis_multi_args(client))
 
 
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_disable_instance_settings)
 @validate_transaction_metrics(
     "test_execute_command:test_redis_execute_command_as_two_args_disable",

tests/datastore_aioredis/test_get_and_set.py

@@ -12,12 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-import aioredis
-
 from newrelic.api.background_task import background_task
 
-from conftest import event_loop, loop, AIOREDIS_VERSION
 from testing_support.fixtures import validate_transaction_metrics, override_application_settings
 from testing_support.db_settings import redis_settings
 from testing_support.util import instance_hostname
@@ -64,23 +60,11 @@
 _disable_rollup_metrics.append((_instance_metric_name, None))
 
 
-if AIOREDIS_VERSION >= (2, 0):
-    clients = [
-        aioredis.Redis(host=DB_SETTINGS["host"], port=_port, db=0),
-        aioredis.StrictRedis(host=DB_SETTINGS["host"], port=_port, db=0),
-    ]
-else:
-    clients = [
-        event_loop.run_until_complete(aioredis.create_redis("redis://%s:%d" % (DB_SETTINGS["host"], _port), db=0)),
-    ]
-
-
 async def exercise_redis(client):
     await client.set("key", "value")
     await client.get("key")
 
 
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_enable_instance_settings)
 @validate_transaction_metrics(
     "test_get_and_set:test_redis_client_operation_enable_instance",
@@ -93,7 +77,6 @@ def test_redis_client_operation_enable_instance(client, loop):
     loop.run_until_complete(exercise_redis(client))
 
 
-@pytest.mark.parametrize("client", clients)
 @override_application_settings(_disable_instance_settings)
 @validate_transaction_metrics(
     "test_get_and_set:test_redis_client_operation_disable_instance",