Cache Package Version Lookups

newrelic/common/package_version_utils.py

@@ -14,6 +14,44 @@
 
 import sys
 
+try:
+    from functools import cache as _cache_package_versions
+except ImportError:
+    from functools import wraps
+    from threading import Lock
+
+    _package_version_cache = {}
+    _package_version_cache_lock = Lock()
+
+    def _cache_package_versions(wrapped):
+        """
+        Threadsafe implementation of caching for _get_package_version.
+
+        Python 2.7 does not have the @functools.cache decorator, and
+        must be reimplemented with support for clearing the cache.
+        """
+
+        @wraps(wrapped)
+        def _wrapper(name):
+            if name in _package_version_cache:
+                return _package_version_cache[name]
+
+            with _package_version_cache_lock:
+                if name in _package_version_cache:
+                    return _package_version_cache[name]
+
+                version = _package_version_cache[name] = wrapped(name)
+                return version
+
+        def cache_clear():
+            """Cache clear function to mimic @functools.cache"""
+            with _package_version_cache_lock:
+                _package_version_cache.clear()
+
+        _wrapper.cache_clear = cache_clear
+        return _wrapper
+
+
 # Need to account for 4 possible variations of version declaration specified in (rejected) PEP 396
 VERSION_ATTRS = ("__version__", "version", "__version_tuple__", "version_tuple")  # nosec
 NULL_VERSIONS = frozenset((None, "", "0", "0.0", "0.0.0", "0.0.0.0", (0,), (0, 0), (0, 0, 0), (0, 0, 0, 0)))  # nosec
@@ -67,6 +105,7 @@ def int_or_str(value):
     return version
 
 
+@_cache_package_versions
 def _get_package_version(name):
     module = sys.modules.get(name, None)
     version = None
@@ -75,7 +114,7 @@ def _get_package_version(name):
     if "importlib" in sys.modules and hasattr(sys.modules["importlib"], "metadata"):
         try:
             # In Python3.10+ packages_distribution can be checked for as well
-            if hasattr(sys.modules["importlib"].metadata, "packages_distributions"):    # pylint: disable=E1101
+            if hasattr(sys.modules["importlib"].metadata, "packages_distributions"):  # pylint: disable=E1101
                 distributions = sys.modules["importlib"].metadata.packages_distributions()  # pylint: disable=E1101
                 distribution_name = distributions.get(name, name)
             else:

tests/agent_unittests/test_package_version_utils.py

@@ -20,6 +20,7 @@
 from newrelic.common.package_version_utils import (
     NULL_VERSIONS,
     VERSION_ATTRS,
+    _get_package_version,
     get_package_version,
     get_package_version_tuple,
 )
@@ -31,7 +32,7 @@
 # such as distribution_packages and removed pkg_resources.
 
 IS_PY38_PLUS = sys.version_info[:2] >= (3, 8)
-IS_PY310_PLUS = sys.version_info[:2] >= (3,10)
+IS_PY310_PLUS = sys.version_info[:2] >= (3, 10)
 SKIP_IF_NOT_IMPORTLIB_METADATA = pytest.mark.skipif(not IS_PY38_PLUS, reason="importlib.metadata is not supported.")
 SKIP_IF_IMPORTLIB_METADATA = pytest.mark.skipif(
     IS_PY38_PLUS, reason="importlib.metadata is preferred over pkg_resources."
@@ -46,7 +47,13 @@ def patched_pytest_module(monkeypatch):
             monkeypatch.delattr(pytest, attr)
 
     yield pytest
-
+
+
+@pytest.fixture(scope="function", autouse=True)
+def cleared_package_version_cache():
+    """Ensure cache is empty before every test to exercise code paths."""
+    _get_package_version.cache_clear()
+
 
 # This test only works on Python 3.7
 @SKIP_IF_IMPORTLIB_METADATA
@@ -123,3 +130,16 @@ def test_mapping_import_to_distribution_packages():
 def test_pkg_resources_metadata():
     version = get_package_version("pytest")
     assert version not in NULL_VERSIONS, version
+
+
+def test_version_caching(monkeypatch):
+    # Add fake module to be deleted later
+    sys.modules["mymodule"] = sys.modules["pytest"]
+    setattr(pytest, "__version__", "1.0.0")
+    version = get_package_version("mymodule")
+    assert version not in NULL_VERSIONS, version
+
+    # Ensure after deleting that the call to _get_package_version still completes because of caching
+    del sys.modules["mymodule"]
+    version = get_package_version("mymodule")
+    assert version not in NULL_VERSIONS, version