Add caching for AWS metadata (#1427)

lrafeei · mergify[bot] · web-flow · commit 777e025ba924 · 2025-07-23T23:52:24.000Z
* Add caching for AWS metadata

* Reset cached metadata

* Use class attr instead of env var

* AWS utilization caching mock tests

* Reset AWS cache in cross agent test

---------

Co-authored-by: mergify[bot] &lt;37929162+mergify[bot]@users.noreply.github.com&gt;
diff --git a/newrelic/common/utilization.py b/newrelic/common/utilization.py
@@ -162,6 +162,7 @@ class AWSUtilization(CommonUtilization):
     METADATA_TOKEN_PATH = "/latest/api/token"  # noqa: S105
     HEADERS = {"X-aws-ec2-metadata-token-ttl-seconds": "21600"}
     VENDOR_NAME = "aws"
+    _utilization_data = None
 
     @classmethod
     def fetchAuthToken(cls):
@@ -184,14 +185,28 @@ def fetch(cls):
         try:
             authToken = cls.fetchAuthToken()
             if authToken is None:
-                return
+                metadata = cls._utilization_data
+                return metadata
             cls.HEADERS = {"X-aws-ec2-metadata-token": authToken}
             with cls.CLIENT_CLS(cls.METADATA_HOST, timeout=cls.FETCH_TIMEOUT) as client:
                 resp = client.send_request(
                     method="GET", path=cls.METADATA_PATH, params=cls.METADATA_QUERY, headers=cls.HEADERS
                 )
             if not 200 <= resp[0] < 300:
                 raise ValueError(resp[0])
+            # Cache this for forced agent restarts within the same
+            # environment if return value is valid.
+            try:
+                response_dict = json.loads(resp[1].decode("utf-8"))
+                availabilityZone = response_dict.get("availabilityZone", None)
+                instanceId = response_dict.get("instanceId", None)
+                instanceType = response_dict.get("instanceType", None)
+                if all((availabilityZone, instanceId, instanceType)):
+                    # Cache the utilization data for reuse
+                    cls._utilization_data = resp[1]
+            except:
+                # Exits without caching if the response is not valid
+                pass
             return resp[1]
         except Exception as e:
             _logger.debug(
diff --git a/tests/agent_unittests/aws.json b/tests/agent_unittests/aws.json
@@ -0,0 +1,105 @@
+[
+  {
+    "testname": "auth token fails, no cached data",
+    "auth_token_cls": "no_token",
+    "uri": {
+      "http://169.254.169.254/latest/dynamic/instance-identity/document": {
+        "response": {
+          "instanceId": null,      
+          "instanceType": null,        
+          "availabilityZone": null
+        },
+        "timeout": false
+      }
+    },
+    "expected_vendors_hash": null,
+    "expected_metrics": {
+      "Supportability/utilization/aws/error": {
+        "call_count": 0
+      }
+    }
+  },
+  {
+    "testname": "auth token succeeds, but utilization data is not valid, no cached data",
+    "auth_token_cls": "fake_token",
+    "uri": {
+      "http://169.254.169.254/latest/dynamic/instance-identity/document": {
+        "response": {
+          "instanceId": "i-test.19characters",
+          "instanceType": null,          
+          "availabilityZone": "us-west-2b"
+        },
+        "timeout": false
+      }
+    },
+    "expected_vendors_hash": null,
+    "expected_metrics": {
+      "Supportability/utilization/aws/error": {
+        "call_count": 1
+      }
+    }
+  },
+  {
+    "testname": "auth token succeeds, utilization data is valid, data is cached (part 1)",
+    "auth_token_cls": "fake_token",
+    "uri": {
+      "http://169.254.169.254/latest/dynamic/instance-identity/document": {
+        "response": {
+          "instanceId": "i-test.19characters",
+          "instanceType": "test.type",          
+          "availabilityZone": "us-west-2b"
+        },
+        "timeout": false
+      }
+    },
+    "expected_vendors_hash": {
+      "aws": {
+        "instanceId": "i-test.19characters",
+        "instanceType": "test.type",
+        "availabilityZone": "us-west-2b"
+      }
+    }
+  },
+  {
+    "testname": "auth token fails, but cached data exists, return cached data",
+    "auth_token_cls": "no_token",
+    "uri": {
+      "http://169.254.169.254/latest/dynamic/instance-identity/document": {
+        "response": {
+          "instanceId": null,      
+          "instanceType": null,        
+          "availabilityZone": null
+        },
+        "timeout": false
+      }
+    },
+    "expected_vendors_hash": {
+      "aws": {
+        "instanceId": "i-test.19characters",
+        "instanceType": "test.type",
+        "availabilityZone": "us-west-2b"
+      }
+    }
+  },
+  {
+    "testname": "auth token succeeds, utilization data is valid, data is cached (part 2)",
+    "auth_token_cls": "fake_token",
+    "uri": {
+      "http://169.254.169.254/latest/dynamic/instance-identity/document": {
+        "response": {
+          "instanceId": "i-test.19characters",
+          "instanceType": "test.type",          
+          "availabilityZone": "us-east-2b"
+        },
+        "timeout": false
+      }
+    },
+    "expected_vendors_hash": {
+      "aws": {
+        "instanceId": "i-test.19characters",
+        "instanceType": "test.type",
+        "availabilityZone": "us-east-2b"
+      }
+    }
+  }
+]
diff --git a/tests/agent_unittests/test_aws_utilization_caching.py b/tests/agent_unittests/test_aws_utilization_caching.py
@@ -0,0 +1,103 @@
+# 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 json
+import os
+
+import pytest
+from testing_support.mock_http_client import create_client_cls
+from testing_support.validators.validate_internal_metrics import validate_internal_metrics
+
+from newrelic.common.utilization import AWSUtilization
+
+CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
+FIXTURE = os.path.normpath(os.path.join(CURRENT_DIR, "aws.json"))
+
+_parameters_list = ["testname", "auth_token_cls", "uri", "expected_vendors_hash", "expected_metrics"]
+
+_parameters = ",".join(_parameters_list)
+
+
+@classmethod
+def fake_token(cls):
+    return "FakeToken"
+
+
+@classmethod
+def no_token(cls):
+    return None
+
+
+def _load_tests():
+    with open(FIXTURE) as fh:
+        js = fh.read()
+    return json.loads(js)
+
+
+def _parametrize_test(test):
+    return tuple([test.get(f, None) for f in _parameters_list])
+
+
+# Load the tests from the JSON fixture
+_aws_tests = [_parametrize_test(t) for t in _load_tests()]
+
+
+# Order of tests:
+# 1. auth token fails, no cached data
+# 2. Auth token succeeds, but utilization data is invalid, no cached data
+# 3. Auth token succeeds, utilization data is valid, data is cached (check both cache and returned value)
+# 4. Auth token fails, but cached data is valid, return cached data
+# 5. Auth token succeeds, utilization data is valid, data is cached (check both cache and returned value)
+
+
+@pytest.mark.parametrize(_parameters, _aws_tests)
+def test_aws_utilization_caching(monkeypatch, testname, auth_token_cls, uri, expected_vendors_hash, expected_metrics):
+    def _get_mock_return_value(api_result):
+        if api_result["timeout"]:
+            return 0, None
+        else:
+            body = json.dumps(api_result["response"])
+            return 200, body.encode("utf-8")
+
+    url, api_result = uri.popitem()
+    status, data = _get_mock_return_value(api_result)
+
+    client_cls = create_client_cls(status, data, url)
+
+    monkeypatch.setattr(AWSUtilization, "CLIENT_CLS", client_cls)
+    monkeypatch.setattr(AWSUtilization, "fetchAuthToken", fake_token if auth_token_cls == "fake_token" else no_token)
+
+    metrics = []
+    if expected_metrics:
+        metrics = [(k, v.get("call_count")) for k, v in expected_metrics.items()]
+
+    # Define function that actually runs the test
+
+    @validate_internal_metrics(metrics=metrics)
+    def _test_aws_data():
+        data = AWSUtilization.detect()
+
+        if data:
+            aws_vendor_hash = {"aws": data}
+        else:
+            aws_vendor_hash = None
+
+        assert aws_vendor_hash == expected_vendors_hash
+        if expected_vendors_hash is not None:
+            # Check that the cached data is set to the most recent valid data
+            assert json.loads(AWSUtilization._utilization_data.decode("utf-8")) == data
+
+    _test_aws_data()
+
+    assert not client_cls.FAIL
diff --git a/tests/cross_agent/test_utilization_configs.py b/tests/cross_agent/test_utilization_configs.py
@@ -27,7 +27,7 @@
 import newrelic.core.config
 from newrelic.common.object_wrapper import function_wrapper
 from newrelic.common.system_info import BootIdUtilization
-from newrelic.common.utilization import CommonUtilization
+from newrelic.common.utilization import AWSUtilization, CommonUtilization
 from newrelic.core.agent_protocol import AgentProtocol
 
 INITIAL_ENV = os.environ
@@ -187,6 +187,8 @@ def test_utilization_settings(test, monkeypatch):
     for key, val in env.items():
         monkeypatch.setenv(key, str(val))
 
+    AWSUtilization._utilization_data = None  # reset cached data before test
+
     @patch_boot_id_file(test)
     @patch_system_info(test, monkeypatch)
     def _test_utilization_data():
