Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Correct Serverless Distributed Tracing Logic #870

Merged
merged 7 commits into from
Jul 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions newrelic/api/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -1030,7 +1030,9 @@ def _create_distributed_trace_data(self):

settings = self._settings
account_id = settings.account_id
trusted_account_key = settings.trusted_account_key
trusted_account_key = settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)
application_id = settings.primary_application_id

if not (account_id and application_id and trusted_account_key and settings.distributed_tracing.enabled):
Expand Down Expand Up @@ -1121,7 +1123,10 @@ def _can_accept_distributed_trace_headers(self):
return False

settings = self._settings
if not (settings.distributed_tracing.enabled and settings.trusted_account_key):
trusted_account_key = settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)
if not (settings.distributed_tracing.enabled and trusted_account_key):
return False

if self._distributed_trace_state:
Expand Down Expand Up @@ -1167,10 +1172,13 @@ def _accept_distributed_trace_payload(self, payload, transport_type="HTTP"):

settings = self._settings
account_id = data.get("ac")
trusted_account_key = settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)

# If trust key doesn't exist in the payload, use account_id
received_trust_key = data.get("tk", account_id)
if settings.trusted_account_key != received_trust_key:
if trusted_account_key != received_trust_key:
self._record_supportability("Supportability/DistributedTrace/AcceptPayload/Ignored/UntrustedAccount")
if settings.debug.log_untrusted_distributed_trace_keys:
_logger.debug(
Expand Down Expand Up @@ -1279,8 +1287,10 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
tracestate = ensure_str(tracestate)
try:
vendors = W3CTraceState.decode(tracestate)
tk = self._settings.trusted_account_key
payload = vendors.pop(tk + "@nr", "")
trusted_account_key = self._settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)
payload = vendors.pop(trusted_account_key + "@nr", "")
self.tracing_vendors = ",".join(vendors.keys())
self.tracestate = vendors.text(limit=31)
except:
Expand All @@ -1289,7 +1299,7 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
# Remove trusted new relic header if available and parse
if payload:
try:
tracestate_data = NrTraceState.decode(payload, tk)
tracestate_data = NrTraceState.decode(payload, trusted_account_key)
except:
tracestate_data = None
if tracestate_data:
Expand Down
63 changes: 63 additions & 0 deletions tests/agent_features/test_distributed_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

from newrelic.api.application import application_instance
from newrelic.api.background_task import BackgroundTask, background_task
from newrelic.api.external_trace import ExternalTrace
from newrelic.api.time_trace import current_trace
from newrelic.api.transaction import (
accept_distributed_trace_headers,
Expand Down Expand Up @@ -361,3 +362,65 @@ def test_current_span_id_inside_transaction():
def test_current_span_id_outside_transaction():
span_id = current_span_id()
assert span_id is None


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_outbound_dt_payload_generation(trusted_account_key):
@override_application_settings(
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
}
)
@background_task(name="test_outbound_dt_payload_generation")
def _test_outbound_dt_payload_generation():
transaction = current_transaction()
payload = ExternalTrace.generate_request_headers(transaction)
if trusted_account_key:
assert payload
# Ensure trusted account key present as vendor
assert dict(payload)["tracestate"].startswith("1@nr=")
else:
assert not payload

_test_outbound_dt_payload_generation()


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_inbound_dt_payload_acceptance(trusted_account_key):
@override_application_settings(
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
}
)
@background_task(name="_test_inbound_dt_payload_acceptance")
def _test_inbound_dt_payload_acceptance():
transaction = current_transaction()

payload = {
"v": [0, 1],
"d": {
"ty": "Mobile",
"ac": "1",
"tk": "1",
"ap": "2827902",
"pa": "5e5733a911cfbc73",
"id": "7d3efb1b173fecfa",
"tr": "d6b4ba0c3a712ca",
"ti": 1518469636035,
"tx": "8703ff3d88eefe9d",
},
}

result = transaction.accept_distributed_trace_payload(payload)
if trusted_account_key:
assert result
else:
assert not result

_test_inbound_dt_payload_acceptance()
144 changes: 74 additions & 70 deletions tests/agent_features/test_serverless_mode.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,16 @@
# limitations under the License.

import json

import pytest
from testing_support.fixtures import override_generic_settings
from testing_support.validators.validate_serverless_data import validate_serverless_data
from testing_support.validators.validate_serverless_metadata import (
validate_serverless_metadata,
)
from testing_support.validators.validate_serverless_payload import (
validate_serverless_payload,
)

from newrelic.api.application import application_instance
from newrelic.api.background_task import background_task
Expand All @@ -22,23 +31,14 @@
from newrelic.api.transaction import current_transaction
from newrelic.core.config import global_settings

from testing_support.fixtures import override_generic_settings
from testing_support.validators.validate_serverless_data import (
validate_serverless_data)
from testing_support.validators.validate_serverless_payload import (
validate_serverless_payload)
from testing_support.validators.validate_serverless_metadata import (
validate_serverless_metadata)


@pytest.fixture(scope='function')
@pytest.fixture(scope="function")
def serverless_application(request):
settings = global_settings()
orig = settings.serverless_mode.enabled
settings.serverless_mode.enabled = True

application_name = 'Python Agent Test (test_serverless_mode:%s)' % (
request.node.name)
application_name = "Python Agent Test (test_serverless_mode:%s)" % (request.node.name)
application = application_instance(application_name)
application.activate()

Expand All @@ -48,17 +48,18 @@ def serverless_application(request):


def test_serverless_payload(capsys, serverless_application):

@override_generic_settings(serverless_application.settings, {
'distributed_tracing.enabled': True,
})
@override_generic_settings(
serverless_application.settings,
{
"distributed_tracing.enabled": True,
},
)
@validate_serverless_data(
expected_methods=('metric_data', 'analytic_event_data'),
forgone_methods=('preconnect', 'connect', 'get_agent_commands'))
expected_methods=("metric_data", "analytic_event_data"),
forgone_methods=("preconnect", "connect", "get_agent_commands"),
)
@validate_serverless_payload()
@background_task(
application=serverless_application,
name='test_serverless_payload')
@background_task(application=serverless_application, name="test_serverless_payload")
def _test():
transaction = current_transaction()
assert transaction.settings.serverless_mode.enabled
Expand All @@ -75,17 +76,15 @@ def _test():


def test_no_cat_headers(serverless_application):
@background_task(
application=serverless_application,
name='test_cat_headers')
@background_task(application=serverless_application, name="test_cat_headers")
def _test_cat_headers():
transaction = current_transaction()

payload = ExternalTrace.generate_request_headers(transaction)
assert not payload

trace = ExternalTrace('testlib', 'http://example.com')
response_headers = [('X-NewRelic-App-Data', 'Cookies')]
trace = ExternalTrace("testlib", "http://example.com")
response_headers = [("X-NewRelic-App-Data", "Cookies")]
with trace:
trace.process_response_headers(response_headers)

Expand All @@ -94,61 +93,66 @@ def _test_cat_headers():
_test_cat_headers()


def test_dt_outbound(serverless_application):
@override_generic_settings(serverless_application.settings, {
'distributed_tracing.enabled': True,
'account_id': '1',
'trusted_account_key': '1',
'primary_application_id': '1',
})
@background_task(
application=serverless_application,
name='test_dt_outbound')
def _test_dt_outbound():
@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_outbound_dt_payload_generation(serverless_application, trusted_account_key):
@override_generic_settings(
serverless_application.settings,
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
},
)
@background_task(application=serverless_application, name="test_outbound_dt_payload_generation")
def _test_outbound_dt_payload_generation():
transaction = current_transaction()
payload = ExternalTrace.generate_request_headers(transaction)
assert payload

_test_dt_outbound()


def test_dt_inbound(serverless_application):
@override_generic_settings(serverless_application.settings, {
'distributed_tracing.enabled': True,
'account_id': '1',
'trusted_account_key': '1',
'primary_application_id': '1',
})
@background_task(
application=serverless_application,
name='test_dt_inbound')
def _test_dt_inbound():
# Ensure trusted account key or account ID present as vendor
assert dict(payload)["tracestate"].startswith("1@nr=")

_test_outbound_dt_payload_generation()


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_inbound_dt_payload_acceptance(serverless_application, trusted_account_key):
@override_generic_settings(
serverless_application.settings,
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
},
)
@background_task(application=serverless_application, name="test_inbound_dt_payload_acceptance")
def _test_inbound_dt_payload_acceptance():
transaction = current_transaction()

payload = {
'v': [0, 1],
'd': {
'ty': 'Mobile',
'ac': '1',
'tk': '1',
'ap': '2827902',
'pa': '5e5733a911cfbc73',
'id': '7d3efb1b173fecfa',
'tr': 'd6b4ba0c3a712ca',
'ti': 1518469636035,
'tx': '8703ff3d88eefe9d',
}
"v": [0, 1],
"d": {
"ty": "Mobile",
"ac": "1",
"tk": "1",
"ap": "2827902",
"pa": "5e5733a911cfbc73",
"id": "7d3efb1b173fecfa",
"tr": "d6b4ba0c3a712ca",
"ti": 1518469636035,
"tx": "8703ff3d88eefe9d",
},
}

result = transaction.accept_distributed_trace_payload(payload)
assert result

_test_dt_inbound()
_test_inbound_dt_payload_acceptance()


@pytest.mark.parametrize('arn_set', (True, False))
@pytest.mark.parametrize("arn_set", (True, False))
def test_payload_metadata_arn(serverless_application, arn_set):

# If the session object gathers the arn from the settings object before the
# lambda handler records it there, then this test will fail.

Expand All @@ -157,17 +161,17 @@ def test_payload_metadata_arn(serverless_application, arn_set):

arn = None
if arn_set:
arn = 'arrrrrrrrrrRrrrrrrrn'
arn = "arrrrrrrrrrRrrrrrrrn"

settings.aws_lambda_metadata.update({'arn': arn, 'function_version': '$LATEST'})
settings.aws_lambda_metadata.update({"arn": arn, "function_version": "$LATEST"})

class Context(object):
invoked_function_arn = arn

@validate_serverless_metadata(exact_metadata={'arn': arn})
@validate_serverless_metadata(exact_metadata={"arn": arn})
@lambda_handler(application=serverless_application)
def handler(event, context):
assert settings.aws_lambda_metadata['arn'] == arn
assert settings.aws_lambda_metadata["arn"] == arn
return {}

try:
Expand Down
Loading