Skip to content

Commit 777e025

Browse files
lrafeeimergify[bot]
lrafeei
and
mergify[bot]
authored
Add caching for AWS metadata (#1427)
* 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] <37929162+mergify[bot]@users.noreply.github.com>
1 parent f939956 commit 777e025

4 files changed

Lines changed: 227 additions & 2 deletions

File tree

‎newrelic/common/utilization.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,7 @@ class AWSUtilization(CommonUtilization):
162162
METADATA_TOKEN_PATH = "/latest/api/token" # noqa: S105
163163
HEADERS = {"X-aws-ec2-metadata-token-ttl-seconds": "21600"}
164164
VENDOR_NAME = "aws"
165+
_utilization_data = None
165166

166167
@classmethod
167168
def fetchAuthToken(cls):
@@ -184,14 +185,28 @@ def fetch(cls):
184185
try:
185186
authToken = cls.fetchAuthToken()
186187
if authToken is None:
187-
return
188+
metadata = cls._utilization_data
189+
return metadata
188190
cls.HEADERS = {"X-aws-ec2-metadata-token": authToken}
189191
with cls.CLIENT_CLS(cls.METADATA_HOST, timeout=cls.FETCH_TIMEOUT) as client:
190192
resp = client.send_request(
191193
method="GET", path=cls.METADATA_PATH, params=cls.METADATA_QUERY, headers=cls.HEADERS
192194
)
193195
if not 200 <= resp[0] < 300:
194196
raise ValueError(resp[0])
197+
# Cache this for forced agent restarts within the same
198+
# environment if return value is valid.
199+
try:
200+
response_dict = json.loads(resp[1].decode("utf-8"))
201+
availabilityZone = response_dict.get("availabilityZone", None)
202+
instanceId = response_dict.get("instanceId", None)
203+
instanceType = response_dict.get("instanceType", None)
204+
if all((availabilityZone, instanceId, instanceType)):
205+
# Cache the utilization data for reuse
206+
cls._utilization_data = resp[1]
207+
except:
208+
# Exits without caching if the response is not valid
209+
pass
195210
return resp[1]
196211
except Exception as e:
197212
_logger.debug(

‎tests/agent_unittests/aws.json‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
[
2+
{
3+
"testname": "auth token fails, no cached data",
4+
"auth_token_cls": "no_token",
5+
"uri": {
6+
"http://169.254.169.254/latest/dynamic/instance-identity/document": {
7+
"response": {
8+
"instanceId": null,
9+
"instanceType": null,
10+
"availabilityZone": null
11+
},
12+
"timeout": false
13+
}
14+
},
15+
"expected_vendors_hash": null,
16+
"expected_metrics": {
17+
"Supportability/utilization/aws/error": {
18+
"call_count": 0
19+
}
20+
}
21+
},
22+
{
23+
"testname": "auth token succeeds, but utilization data is not valid, no cached data",
24+
"auth_token_cls": "fake_token",
25+
"uri": {
26+
"http://169.254.169.254/latest/dynamic/instance-identity/document": {
27+
"response": {
28+
"instanceId": "i-test.19characters",
29+
"instanceType": null,
30+
"availabilityZone": "us-west-2b"
31+
},
32+
"timeout": false
33+
}
34+
},
35+
"expected_vendors_hash": null,
36+
"expected_metrics": {
37+
"Supportability/utilization/aws/error": {
38+
"call_count": 1
39+
}
40+
}
41+
},
42+
{
43+
"testname": "auth token succeeds, utilization data is valid, data is cached (part 1)",
44+
"auth_token_cls": "fake_token",
45+
"uri": {
46+
"http://169.254.169.254/latest/dynamic/instance-identity/document": {
47+
"response": {
48+
"instanceId": "i-test.19characters",
49+
"instanceType": "test.type",
50+
"availabilityZone": "us-west-2b"
51+
},
52+
"timeout": false
53+
}
54+
},
55+
"expected_vendors_hash": {
56+
"aws": {
57+
"instanceId": "i-test.19characters",
58+
"instanceType": "test.type",
59+
"availabilityZone": "us-west-2b"
60+
}
61+
}
62+
},
63+
{
64+
"testname": "auth token fails, but cached data exists, return cached data",
65+
"auth_token_cls": "no_token",
66+
"uri": {
67+
"http://169.254.169.254/latest/dynamic/instance-identity/document": {
68+
"response": {
69+
"instanceId": null,
70+
"instanceType": null,
71+
"availabilityZone": null
72+
},
73+
"timeout": false
74+
}
75+
},
76+
"expected_vendors_hash": {
77+
"aws": {
78+
"instanceId": "i-test.19characters",
79+
"instanceType": "test.type",
80+
"availabilityZone": "us-west-2b"
81+
}
82+
}
83+
},
84+
{
85+
"testname": "auth token succeeds, utilization data is valid, data is cached (part 2)",
86+
"auth_token_cls": "fake_token",
87+
"uri": {
88+
"http://169.254.169.254/latest/dynamic/instance-identity/document": {
89+
"response": {
90+
"instanceId": "i-test.19characters",
91+
"instanceType": "test.type",
92+
"availabilityZone": "us-east-2b"
93+
},
94+
"timeout": false
95+
}
96+
},
97+
"expected_vendors_hash": {
98+
"aws": {
99+
"instanceId": "i-test.19characters",
100+
"instanceType": "test.type",
101+
"availabilityZone": "us-east-2b"
102+
}
103+
}
104+
}
105+
]

‎tests/agent_unittests/test_aws_utilization_caching.py‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
# Copyright 2010 New Relic, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import json
16+
import os
17+
18+
import pytest
19+
from testing_support.mock_http_client import create_client_cls
20+
from testing_support.validators.validate_internal_metrics import validate_internal_metrics
21+
22+
from newrelic.common.utilization import AWSUtilization
23+
24+
CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
25+
FIXTURE = os.path.normpath(os.path.join(CURRENT_DIR, "aws.json"))
26+
27+
_parameters_list = ["testname", "auth_token_cls", "uri", "expected_vendors_hash", "expected_metrics"]
28+
29+
_parameters = ",".join(_parameters_list)
30+
31+
32+
@classmethod
33+
def fake_token(cls):
34+
return "FakeToken"
35+
36+
37+
@classmethod
38+
def no_token(cls):
39+
return None
40+
41+
42+
def _load_tests():
43+
with open(FIXTURE) as fh:
44+
js = fh.read()
45+
return json.loads(js)
46+
47+
48+
def _parametrize_test(test):
49+
return tuple([test.get(f, None) for f in _parameters_list])
50+
51+
52+
# Load the tests from the JSON fixture
53+
_aws_tests = [_parametrize_test(t) for t in _load_tests()]
54+
55+
56+
# Order of tests:
57+
# 1. auth token fails, no cached data
58+
# 2. Auth token succeeds, but utilization data is invalid, no cached data
59+
# 3. Auth token succeeds, utilization data is valid, data is cached (check both cache and returned value)
60+
# 4. Auth token fails, but cached data is valid, return cached data
61+
# 5. Auth token succeeds, utilization data is valid, data is cached (check both cache and returned value)
62+
63+
64+
@pytest.mark.parametrize(_parameters, _aws_tests)
65+
def test_aws_utilization_caching(monkeypatch, testname, auth_token_cls, uri, expected_vendors_hash, expected_metrics):
66+
def _get_mock_return_value(api_result):
67+
if api_result["timeout"]:
68+
return 0, None
69+
else:
70+
body = json.dumps(api_result["response"])
71+
return 200, body.encode("utf-8")
72+
73+
url, api_result = uri.popitem()
74+
status, data = _get_mock_return_value(api_result)
75+
76+
client_cls = create_client_cls(status, data, url)
77+
78+
monkeypatch.setattr(AWSUtilization, "CLIENT_CLS", client_cls)
79+
monkeypatch.setattr(AWSUtilization, "fetchAuthToken", fake_token if auth_token_cls == "fake_token" else no_token)
80+
81+
metrics = []
82+
if expected_metrics:
83+
metrics = [(k, v.get("call_count")) for k, v in expected_metrics.items()]
84+
85+
# Define function that actually runs the test
86+
87+
@validate_internal_metrics(metrics=metrics)
88+
def _test_aws_data():
89+
data = AWSUtilization.detect()
90+
91+
if data:
92+
aws_vendor_hash = {"aws": data}
93+
else:
94+
aws_vendor_hash = None
95+
96+
assert aws_vendor_hash == expected_vendors_hash
97+
if expected_vendors_hash is not None:
98+
# Check that the cached data is set to the most recent valid data
99+
assert json.loads(AWSUtilization._utilization_data.decode("utf-8")) == data
100+
101+
_test_aws_data()
102+
103+
assert not client_cls.FAIL

‎tests/cross_agent/test_utilization_configs.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import newrelic.core.config
2828
from newrelic.common.object_wrapper import function_wrapper
2929
from newrelic.common.system_info import BootIdUtilization
30-
from newrelic.common.utilization import CommonUtilization
30+
from newrelic.common.utilization import AWSUtilization, CommonUtilization
3131
from newrelic.core.agent_protocol import AgentProtocol
3232

3333
INITIAL_ENV = os.environ
@@ -187,6 +187,8 @@ def test_utilization_settings(test, monkeypatch):
187187
for key, val in env.items():
188188
monkeypatch.setenv(key, str(val))
189189

190+
AWSUtilization._utilization_data = None # reset cached data before test
191+
190192
@patch_boot_id_file(test)
191193
@patch_system_info(test, monkeypatch)
192194
def _test_utilization_data():

0 commit comments

Comments
 (0)
⚡