From 5996de662544f77310b5dd5da0ed425989f96357 Mon Sep 17 00:00:00 2001 From: Timothy Pansino <11214426+TimPansino@users.noreply.github.com> Date: Thu, 19 Oct 2023 10:46:27 -0700 Subject: [PATCH] Fix Redis Generator Methods (#947) * Fix scan_iter for redis * Replace generator methods * Update instance info instrumentation * Remove mistake from uninstrumented methods * Add skip condition to asyncio generator tests * Add skip condition to asyncio generator tests --------- Co-authored-by: Lalleh Rafeei Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- newrelic/hooks/datastore_redis.py | 74 +++-- tests/datastore_redis/conftest.py | 1 + tests/datastore_redis/test_generators.py | 258 ++++++++++++++++++ .../test_uninstrumented_methods.py | 1 - 4 files changed, 294 insertions(+), 40 deletions(-) create mode 100644 tests/datastore_redis/test_generators.py diff --git a/newrelic/hooks/datastore_redis.py b/newrelic/hooks/datastore_redis.py index 6ba192002..bbc517586 100644 --- a/newrelic/hooks/datastore_redis.py +++ b/newrelic/hooks/datastore_redis.py @@ -14,10 +14,11 @@ import re -from newrelic.api.datastore_trace import DatastoreTrace +from newrelic.api.datastore_trace import DatastoreTrace, DatastoreTraceWrapper, wrap_datastore_trace from newrelic.api.time_trace import current_trace from newrelic.api.transaction import current_transaction -from newrelic.common.object_wrapper import function_wrapper, wrap_function_wrapper +from newrelic.common.object_wrapper import wrap_function_wrapper +from newrelic.common.async_wrapper import coroutine_wrapper, async_generator_wrapper, generator_wrapper _redis_client_sync_methods = { "acl_dryrun", @@ -136,6 +137,7 @@ "client_no_evict", "client_pause", "client_reply", + "client_setinfo", "client_setname", "client_tracking", "client_trackinginfo", @@ -162,7 +164,6 @@ "cluster_reset", "cluster_save_config", "cluster_set_config_epoch", - "client_setinfo", "cluster_setslot", "cluster_slaves", "cluster_slots", @@ -248,7 +249,7 @@ "hmset_dict", "hmset", "hrandfield", - "hscan_inter", + "hscan_iter", "hscan", "hset", "hsetnx", @@ -399,8 +400,8 @@ "syndump", "synupdate", "tagvals", - "tfcall", "tfcall_async", + "tfcall", "tfunction_delete", "tfunction_list", "tfunction_load", @@ -473,6 +474,13 @@ "zunionstore", } +_redis_client_gen_methods = { + "scan_iter", + "hscan_iter", + "sscan_iter", + "zscan_iter", +} + _redis_client_methods = _redis_client_sync_methods.union(_redis_client_async_methods) _redis_multipart_commands = set(["client", "cluster", "command", "config", "debug", "sentinel", "slowlog", "script"]) @@ -498,50 +506,31 @@ def _instance_info(kwargs): def _wrap_Redis_method_wrapper_(module, instance_class_name, operation): - def _nr_wrapper_Redis_method_(wrapped, instance, args, kwargs): - transaction = current_transaction() - - if transaction is None: - return wrapped(*args, **kwargs) - - dt = DatastoreTrace(product="Redis", target=None, operation=operation, source=wrapped) - - transaction._nr_datastore_instance_info = (None, None, None) - - with dt: - result = wrapped(*args, **kwargs) - - host, port_path_or_id, db = transaction._nr_datastore_instance_info - dt.host = host - dt.port_path_or_id = port_path_or_id - dt.database_name = db - - return result - name = "%s.%s" % (instance_class_name, operation) - wrap_function_wrapper(module, name, _nr_wrapper_Redis_method_) + if operation in _redis_client_gen_methods: + async_wrapper = generator_wrapper + else: + async_wrapper = None + wrap_datastore_trace(module, name, product="Redis", target=None, operation=operation, async_wrapper=async_wrapper) -def _wrap_asyncio_Redis_method_wrapper(module, instance_class_name, operation): - @function_wrapper - async def _nr_wrapper_asyncio_Redis_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 _wrap_asyncio_Redis_method_wrapper(module, instance_class_name, operation): def _nr_wrapper_asyncio_Redis_method_(wrapped, instance, args, kwargs): from redis.asyncio.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_asyncio_Redis_async_method_(wrapped)(*args, **kwargs) + # Method should be run when awaited or iterated, therefore we wrap in an async wrapper. + return DatastoreTraceWrapper(wrapped, product="Redis", target=None, operation=operation, async_wrapper=async_wrapper)(*args, **kwargs) name = "%s.%s" % (instance_class_name, operation) + if operation in _redis_client_gen_methods: + async_wrapper = async_generator_wrapper + else: + async_wrapper = coroutine_wrapper + wrap_function_wrapper(module, name, _nr_wrapper_asyncio_Redis_method_) @@ -614,7 +603,15 @@ def _nr_Connection_send_command_wrapper_(wrapped, instance, args, kwargs): except: pass - transaction._nr_datastore_instance_info = (host, port_path_or_id, db) + # Find DatastoreTrace no matter how many other traces are inbetween + trace = current_trace() + 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 # Older Redis clients would when sending multi part commands pass # them in as separate arguments to send_command(). Need to therefore @@ -666,7 +663,6 @@ def instrument_asyncio_redis_client(module): if hasattr(class_, operation): _wrap_asyncio_Redis_method_wrapper(module, "Redis", operation) - def instrument_redis_commands_core(module): _instrument_redis_commands_module(module, "CoreCommands") diff --git a/tests/datastore_redis/conftest.py b/tests/datastore_redis/conftest.py index 53ff2658d..6747039b4 100644 --- a/tests/datastore_redis/conftest.py +++ b/tests/datastore_redis/conftest.py @@ -15,6 +15,7 @@ import pytest from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture # noqa: F401; pylint: disable=W0611 +from testing_support.fixture.event_loop import event_loop as loop # noqa: F401; pylint: disable=W0611 _default_settings = { diff --git a/tests/datastore_redis/test_generators.py b/tests/datastore_redis/test_generators.py new file mode 100644 index 000000000..f747838e1 --- /dev/null +++ b/tests/datastore_redis/test_generators.py @@ -0,0 +1,258 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import redis +from testing_support.db_settings import redis_settings +from testing_support.fixtures import override_application_settings +from testing_support.util import instance_hostname +from testing_support.validators.validate_transaction_metrics import ( + validate_transaction_metrics, +) + +from newrelic.api.background_task import background_task +from newrelic.api.datastore_trace import DatastoreTrace +from newrelic.api.time_trace import current_trace +from newrelic.common.package_version_utils import get_package_version_tuple + +DB_SETTINGS = redis_settings()[0] +REDIS_PY_VERSION = get_package_version_tuple("redis") + +# Settings + +_enable_instance_settings = { + "datastore_tracer.instance_reporting.enabled": True, +} +_disable_instance_settings = { + "datastore_tracer.instance_reporting.enabled": False, +} + +# Metrics + +_base_scoped_metrics = ( + ("Datastore/operation/Redis/scan_iter", 1), + ("Datastore/operation/Redis/sscan_iter", 1), + ("Datastore/operation/Redis/zscan_iter", 1), + ("Datastore/operation/Redis/hscan_iter", 1), + ("Datastore/operation/Redis/set", 1), + ("Datastore/operation/Redis/sadd", 1), + ("Datastore/operation/Redis/zadd", 1), + ("Datastore/operation/Redis/hset", 1), +) + +_base_rollup_metrics = ( + ("Datastore/all", 8), + ("Datastore/allOther", 8), + ("Datastore/Redis/all", 8), + ("Datastore/Redis/allOther", 8), + ("Datastore/operation/Redis/scan_iter", 1), + ("Datastore/operation/Redis/sscan_iter", 1), + ("Datastore/operation/Redis/zscan_iter", 1), + ("Datastore/operation/Redis/hscan_iter", 1), + ("Datastore/operation/Redis/set", 1), + ("Datastore/operation/Redis/sadd", 1), + ("Datastore/operation/Redis/zadd", 1), + ("Datastore/operation/Redis/hset", 1), +) + +_disable_rollup_metrics = list(_base_rollup_metrics) +_enable_rollup_metrics = list(_base_rollup_metrics) + +_host = instance_hostname(DB_SETTINGS["host"]) +_port = DB_SETTINGS["port"] + +_instance_metric_name = "Datastore/instance/Redis/%s/%s" % (_host, _port) + +_enable_rollup_metrics.append((_instance_metric_name, 8)) + +_disable_rollup_metrics.append((_instance_metric_name, None)) + +# Operations + + +def exercise_redis(client): + """ + Exercise client generators by iterating on various methods and ensuring they are + non-empty, and that traces are started and stopped with the generator. + """ + + # Set existing values + client.set("scan-key", "value") + client.sadd("sscan-key", "value") + client.zadd("zscan-key", {"value": 1}) + client.hset("hscan-key", "field", "value") + + # Check generators + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + for k in client.scan_iter("scan-*"): + assert k == b"scan-key" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + for k in client.sscan_iter("sscan-key"): + assert k == b"value" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + for k, _ in client.zscan_iter("zscan-key"): + assert k == b"value" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + for f, v in client.hscan_iter("hscan-key"): + assert f == b"field" + assert v == b"value" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + +async def exercise_redis_async(client): + """ + Exercise client generators by iterating on various methods and ensuring they are + non-empty, and that traces are started and stopped with the generator. + """ + + # Set existing values + await client.set("scan-key", "value") + await client.sadd("sscan-key", "value") + await client.zadd("zscan-key", {"value": 1}) + await client.hset("hscan-key", "field", "value") + + # Check generators + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + async for k in client.scan_iter("scan-*"): + assert k == b"scan-key" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + async for k in client.sscan_iter("sscan-key"): + assert k == b"value" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + async for k, _ in client.zscan_iter("zscan-key"): + assert k == b"value" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + flag = False + assert not isinstance(current_trace(), DatastoreTrace) # Assert no active DatastoreTrace + async for f, v in client.hscan_iter("hscan-key"): + assert f == b"field" + assert v == b"value" + assert isinstance(current_trace(), DatastoreTrace) # Assert DatastoreTrace now active + flag = True + assert flag + + +# Tests + + +@override_application_settings(_enable_instance_settings) +@validate_transaction_metrics( + "test_generators:test_strict_redis_generator_enable_instance", + scoped_metrics=_base_scoped_metrics, + rollup_metrics=_enable_rollup_metrics, + background_task=True, +) +@background_task() +def test_strict_redis_generator_enable_instance(): + client = redis.StrictRedis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0) + exercise_redis(client) + + +@override_application_settings(_disable_instance_settings) +@validate_transaction_metrics( + "test_generators:test_strict_redis_generator_disable_instance", + scoped_metrics=_base_scoped_metrics, + rollup_metrics=_disable_rollup_metrics, + background_task=True, +) +@background_task() +def test_strict_redis_generator_disable_instance(): + client = redis.StrictRedis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0) + exercise_redis(client) + + +@override_application_settings(_enable_instance_settings) +@validate_transaction_metrics( + "test_generators:test_redis_generator_enable_instance", + scoped_metrics=_base_scoped_metrics, + rollup_metrics=_enable_rollup_metrics, + background_task=True, +) +@background_task() +def test_redis_generator_enable_instance(): + client = redis.Redis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0) + exercise_redis(client) + + +@override_application_settings(_disable_instance_settings) +@validate_transaction_metrics( + "test_generators:test_redis_generator_disable_instance", + scoped_metrics=_base_scoped_metrics, + rollup_metrics=_disable_rollup_metrics, + background_task=True, +) +@background_task() +def test_redis_generator_disable_instance(): + client = redis.Redis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0) + exercise_redis(client) + + +@pytest.mark.skipif(REDIS_PY_VERSION < (4, 2), reason="Redis.asyncio was not added until v4.2") +@override_application_settings(_enable_instance_settings) +@validate_transaction_metrics( + "test_generators:test_redis_async_generator_enable_instance", + scoped_metrics=_base_scoped_metrics, + rollup_metrics=_enable_rollup_metrics, + background_task=True, +) +@background_task() +def test_redis_async_generator_enable_instance(loop): + client = redis.asyncio.Redis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0) + loop.run_until_complete(exercise_redis_async(client)) + + +@pytest.mark.skipif(REDIS_PY_VERSION < (4, 2), reason="Redis.asyncio was not added until v4.2") +@override_application_settings(_disable_instance_settings) +@validate_transaction_metrics( + "test_generators:test_redis_async_generator_disable_instance", + scoped_metrics=_base_scoped_metrics, + rollup_metrics=_disable_rollup_metrics, + background_task=True, +) +@background_task() +def test_redis_async_generator_disable_instance(loop): + client = redis.asyncio.Redis(host=DB_SETTINGS["host"], port=DB_SETTINGS["port"], db=0) + loop.run_until_complete(exercise_redis_async(client)) diff --git a/tests/datastore_redis/test_uninstrumented_methods.py b/tests/datastore_redis/test_uninstrumented_methods.py index d86f4de95..c0be684b2 100644 --- a/tests/datastore_redis/test_uninstrumented_methods.py +++ b/tests/datastore_redis/test_uninstrumented_methods.py @@ -65,7 +65,6 @@ "get_property", "get_relation", "get_retry", - "hscan_iter", "index_name", "labels", "list_keys",