From d5cd83beb3fe9736846259151ea9304d37a16efa Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Mon, 7 Nov 2022 16:43:32 -0800 Subject: [PATCH 1/9] implemented oversampling mitigation --- aws_xray_sdk/core/models/dummy_entities.py | 2 +- aws_xray_sdk/core/models/entity.py | 1 + aws_xray_sdk/core/recorder.py | 12 ++-- aws_xray_sdk/core/utils/sqs_message_helper.py | 10 +++ aws_xray_sdk/ext/util.py | 2 +- tests/test_facade_segment.py | 33 +++++++++ tests/test_recorder.py | 35 ++++++++++ tests/test_sqs_message_helper.py | 68 +++++++++++++++++++ tests/test_trace_entities.py | 44 ++++++++++++ tests/test_utils.py | 43 +++++++++++- 10 files changed, 243 insertions(+), 7 deletions(-) create mode 100644 aws_xray_sdk/core/utils/sqs_message_helper.py create mode 100644 tests/test_sqs_message_helper.py diff --git a/aws_xray_sdk/core/models/dummy_entities.py b/aws_xray_sdk/core/models/dummy_entities.py index 9e4a0379..6d962d71 100644 --- a/aws_xray_sdk/core/models/dummy_entities.py +++ b/aws_xray_sdk/core/models/dummy_entities.py @@ -11,7 +11,7 @@ class DummySegment(Segment): the segment based on sampling rules. Adding data to a dummy segment becomes a no-op except for subsegments. This is to reduce the memory footprint of the SDK. - A dummy segment will not be sent to the X-Ray daemon. Manually create + A dummy segment will not be sent to the X-Ray daemon. Manually creating dummy segments is not recommended. """ diff --git a/aws_xray_sdk/core/models/entity.py b/aws_xray_sdk/core/models/entity.py index 41ef3893..427a368a 100644 --- a/aws_xray_sdk/core/models/entity.py +++ b/aws_xray_sdk/core/models/entity.py @@ -81,6 +81,7 @@ def add_subsegment(self, subsegment): """ self._check_ended() subsegment.parent_id = self.id + subsegment.sampled = self.sampled if subsegment.sampled == True else False self.subsegments.append(subsegment) def remove_subsegment(self, subsegment): diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index 66b5731c..a850491b 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -275,7 +275,7 @@ def current_segment(self): else: return entity - def begin_subsegment(self, name, namespace='local'): + def begin_subsegment(self, name, namespace='local', sampling=True): """ Begin a new subsegment. If there is open subsegment, the newly created subsegment will be the @@ -296,8 +296,11 @@ def begin_subsegment(self, name, namespace='local'): log.warning("No segment found, cannot begin subsegment %s." % name) return None - if not segment.sampled: + current_entity = self.get_trace_entity() + + if not current_entity.sampled or not sampling: subsegment = DummySubsegment(segment, name) + subsegment.sampled = False else: subsegment = Subsegment(name, namespace, segment) @@ -335,7 +338,7 @@ def end_subsegment(self, end_time=None): # if segment is already close, we check if we can send entire segment # otherwise we check if we need to stream some subsegments - if self.current_segment().ready_to_send(): + if self.current_segment().ready_to_send(): self._send_segment() else: self.stream_subsegments() @@ -487,7 +490,8 @@ def _send_segment(self): def _stream_subsegment_out(self, subsegment): log.debug("streaming subsegments...") - self.emitter.send_entity(subsegment) + if subsegment.sampled: + self.emitter.send_entity(subsegment) def _load_sampling_rules(self, sampling_rules): diff --git a/aws_xray_sdk/core/utils/sqs_message_helper.py b/aws_xray_sdk/core/utils/sqs_message_helper.py new file mode 100644 index 00000000..947176d3 --- /dev/null +++ b/aws_xray_sdk/core/utils/sqs_message_helper.py @@ -0,0 +1,10 @@ +class SqsMessageHelper: + + @staticmethod + def isSampled(sqs_message): + attributes = sqs_message['attributes'] + + if 'AWSTraceHeader' not in attributes: + return False + + return 'Sampled=1' in attributes['AWSTraceHeader'] \ No newline at end of file diff --git a/aws_xray_sdk/ext/util.py b/aws_xray_sdk/ext/util.py index 8390f9ee..69c3fdf7 100644 --- a/aws_xray_sdk/ext/util.py +++ b/aws_xray_sdk/ext/util.py @@ -35,7 +35,7 @@ def inject_trace_header(headers, entity): else: header = entity.get_origin_trace_header() data = header.data if header else None - + to_insert = TraceHeader( root=entity.trace_id, parent=entity.id, diff --git a/tests/test_facade_segment.py b/tests/test_facade_segment.py index 30842019..57e12071 100644 --- a/tests/test_facade_segment.py +++ b/tests/test_facade_segment.py @@ -55,3 +55,36 @@ def test_structure_intact(): assert segment.subsegments[0] is subsegment assert subsegment.subsegments[0] is subsegment2 + +def test_adding_unsampled_subsegment(): + + segment = FacadeSegment('name', 'id', 'id', True) + subsegment = Subsegment('sampled', 'local', segment) + subsegment2 = Subsegment('unsampled', 'local', segment) + subsegment2.sampled = False + + segment.add_subsegment(subsegment) + subsegment.add_subsegment(subsegment2) + + + assert segment.subsegments[0] is subsegment + assert subsegment.subsegments[0] is subsegment2 + assert subsegment2.sampled == False + +def test_adding_subsegment_respects_parent_sampling_decision(): + + segment = FacadeSegment('name', 'id', 'id', True) + subsegment = Subsegment('sampled', 'local', segment) + subsegment2 = Subsegment('unsampled', 'local', segment) + subsegment2.sampled = False + + subsegment3 = Subsegment('unsampled-child', 'local', segment) + segment.add_subsegment(subsegment) + subsegment.add_subsegment(subsegment2) + + subsegment2.add_subsegment(subsegment3) + + assert segment.subsegments[0] is subsegment + assert subsegment.subsegments[0] is subsegment2 + assert subsegment2.sampled == False + assert subsegment3.sampled == False \ No newline at end of file diff --git a/tests/test_recorder.py b/tests/test_recorder.py index ee60e5a9..c99db141 100644 --- a/tests/test_recorder.py +++ b/tests/test_recorder.py @@ -141,6 +141,24 @@ def test_first_begin_segment_sampled(): assert segment.sampled +def test_unsampled_subsegment_of_sampled_parent(): + xray_recorder = get_new_stubbed_recorder() + xray_recorder.configure(sampling=True) + segment = xray_recorder.begin_segment('name', sampling=True) + subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + + assert segment.sampled == True + assert subsegment.sampled == False + +def test_begin_subsegment_unsampled(): + xray_recorder = get_new_stubbed_recorder() + xray_recorder.configure(sampling=False) + segment = xray_recorder.begin_segment('name', sampling=False) + subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + + assert segment.sampled == False + assert subsegment.sampled == False + def test_in_segment_closing(): xray_recorder = get_new_stubbed_recorder() @@ -201,6 +219,23 @@ def test_disable_is_dummy(): assert type(xray_recorder.current_segment()) is DummySegment assert type(xray_recorder.current_subsegment()) is DummySubsegment +def test_unsampled_subsegment_is_dummy(): + assert global_sdk_config.sdk_enabled() + segment = xray_recorder.begin_segment('name') + subsegment = xray_recorder.begin_subsegment('name', sampling=False) + + assert type(xray_recorder.current_subsegment()) is DummySubsegment + +def test_subsegment_respects_parent_sampling_decision(): + assert global_sdk_config.sdk_enabled() + segment = xray_recorder.begin_segment('name') + subsegment = xray_recorder.begin_subsegment('name2', sampling=False) + subsegment2 = xray_recorder.begin_subsegment('unsampled-subsegment') + + assert type(xray_recorder.current_subsegment()) is DummySubsegment + assert subsegment.sampled == False + assert subsegment2.sampled == False + def test_disabled_empty_context_current_calls(): global_sdk_config.set_sdk_enabled(False) diff --git a/tests/test_sqs_message_helper.py b/tests/test_sqs_message_helper.py new file mode 100644 index 00000000..6ee44b8d --- /dev/null +++ b/tests/test_sqs_message_helper.py @@ -0,0 +1,68 @@ +from aws_xray_sdk.core.utils.sqs_message_helper import SqsMessageHelper + +import pytest + +sampleSqsMessageEvent = { + "Records": [ + { + "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d", + "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a...", + "body": "Test message.", + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1545082649183", + "SenderId": "AIDAIENQZJOLO23YVJ4VO", + "ApproximateFirstReceiveTimestamp": "1545082649185", + "AWSTraceHeader":"Root=1-632BB806-bd862e3fe1be46a994272793;Sampled=1" + }, + "messageAttributes": {}, + "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue", + "awsRegion": "us-east-2" + }, + { + "messageId": "2e1424d4-f796-459a-8184-9c92662be6da", + "receiptHandle": "AQEBzWwaftRI0KuVm4tP+/7q1rGgNqicHq...", + "body": "Test message.", + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1545082650636", + "SenderId": "AIDAIENQZJOLO23YVJ4VO", + "ApproximateFirstReceiveTimestamp": "1545082650649", + "AWSTraceHeader":"Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=0" + }, + "messageAttributes": {}, + "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue", + "awsRegion": "us-east-2" + }, + { + "messageId": "2e1424d4-f796-459a-8184-9c92662be6da", + "receiptHandle": "AQEBzWwaftRI0KuVm4tP+/7q1rGgNqicHq...", + "body": "Test message.", + "attributes": { + "ApproximateReceiveCount": "1", + "SentTimestamp": "1545082650636", + "SenderId": "AIDAIENQZJOLO23YVJ4VO", + "ApproximateFirstReceiveTimestamp": "1545082650649", + "AWSTraceHeader":"Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8" + }, + "messageAttributes": {}, + "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue", + "awsRegion": "us-east-2" + } + ] + } + +def test_return_true_when_sampling_1(): + assert SqsMessageHelper.isSampled(sampleSqsMessageEvent['Records'][0]) == True + +def test_return_false_when_sampling_0(): + assert SqsMessageHelper.isSampled(sampleSqsMessageEvent['Records'][1]) == False + +def test_return_false_with_no_sampling_flag(): + assert SqsMessageHelper.isSampled(sampleSqsMessageEvent['Records'][2]) == False \ No newline at end of file diff --git a/tests/test_trace_entities.py b/tests/test_trace_entities.py index e42cee0c..29300aaa 100644 --- a/tests/test_trace_entities.py +++ b/tests/test_trace_entities.py @@ -263,3 +263,47 @@ def test_add_exception_appending_exceptions(): assert isinstance(segment.cause, dict) assert len(segment.cause['exceptions']) == 2 + +def test_adding_unsampled_subsegment(): + segment = Segment('seg') + subsegment = Subsegment('sub', 'local', segment) + subsegment.sampled = False + segment.add_subsegment(subsegment) + + assert subsegment.sampled == False + assert segment.sampled == True + +def test_adding_unsampled_subsegment_to_subsegment(): + segment = Segment('seg') + subsegment = Subsegment('sub', 'local', segment) + segment.add_subsegment(subsegment) + subsubsegment = Subsegment('subsub', 'local', segment) + subsubsegment.sampled = False + subsegment.add_subsegment(subsubsegment) + + assert segment.sampled == True + assert subsegment.sampled == True + assert subsubsegment.sampled == False + +def test_adding_subsegment_respects_parent_sampling(): + segment = Segment('seg') + subsegment = Subsegment('sub', 'local', segment) + subsegment.sampled = False + segment.add_subsegment(subsegment) + subsubsegment = Subsegment('subsub', 'local', segment) + subsegment.add_subsegment(subsubsegment) + + assert segment.sampled == True + assert subsegment.sampled == False + assert subsubsegment.sampled == False + +def test_adding_subsegment_respects_parent_sampling2(): + segment = Segment('seg') + subsegment = Subsegment('sub', 'local', segment) + segment.add_subsegment(subsegment) + subsubsegment = Subsegment('subsub', 'local', segment) + subsegment.add_subsegment(subsubsegment) + + assert segment.sampled == True + assert subsegment.sampled == True + assert subsubsegment.sampled == True \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py index 939fde42..74001ded 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,4 +1,10 @@ -from aws_xray_sdk.ext.util import to_snake_case, get_hostname, strip_url +from aws_xray_sdk.ext.util import to_snake_case, get_hostname, strip_url, inject_trace_header +from aws_xray_sdk.core.models.segment import Segment +from aws_xray_sdk.core.models.subsegment import Subsegment +from aws_xray_sdk.core.models.dummy_entities import DummySegment, DummySubsegment +from .util import get_new_stubbed_recorder + +xray_recorder = get_new_stubbed_recorder() UNKNOWN_HOST = "UNKNOWN HOST" @@ -52,3 +58,38 @@ def test_strip_url(): assert strip_url("") == "" assert not strip_url(None) + + +def test_inject_trace_header_unsampled(): + headers = {'host': 'test', 'accept': '*/*', 'connection': 'keep-alive', 'X-Amzn-Trace-Id': 'Root=1-6369739a-7d8bb07e519b795eb24d382d;Parent=089e3de743fb9e79;Sampled=1'} + xray_recorder = get_new_stubbed_recorder() + xray_recorder.configure(sampling=True) + segment = xray_recorder.begin_segment('name', sampling=True) + subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + + inject_trace_header(headers, subsegment) + + assert 'Sampled=0' in headers['X-Amzn-Trace-Id'] + +def test_inject_trace_header_respects_parent_subsegment(): + headers = {'host': 'test', 'accept': '*/*', 'connection': 'keep-alive', 'X-Amzn-Trace-Id': 'Root=1-6369739a-7d8bb07e519b795eb24d382d;Parent=089e3de743fb9e79;Sampled=1'} + + xray_recorder = get_new_stubbed_recorder() + xray_recorder.configure(sampling=True) + segment = xray_recorder.begin_segment('name', sampling=True) + subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + subsegment2 = xray_recorder.begin_subsegment('unsampled2') + inject_trace_header(headers, subsegment2) + + assert 'Sampled=0' in headers['X-Amzn-Trace-Id'] + +def test_inject_trace_header_sampled(): + headers = {'host': 'test', 'accept': '*/*', 'connection': 'keep-alive', 'X-Amzn-Trace-Id': 'Root=1-6369739a-7d8bb07e519b795eb24d382d;Parent=089e3de743fb9e79;Sampled=1'} + xray_recorder = get_new_stubbed_recorder() + xray_recorder.configure(sampling=True) + segment = xray_recorder.begin_segment('name') + subsegment = xray_recorder.begin_subsegment('unsampled') + + inject_trace_header(headers, subsegment) + + assert 'Sampled=1' in headers['X-Amzn-Trace-Id'] \ No newline at end of file From 3042958ff25f9026eb5f5c0a12eb0bc9ce5349db Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Mon, 7 Nov 2022 16:59:46 -0800 Subject: [PATCH 2/9] updated begin_subsegment docstring to include sampling parameter --- aws_xray_sdk/core/recorder.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index a850491b..3cdbaa2f 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -284,6 +284,8 @@ def begin_subsegment(self, name, namespace='local', sampling=True): :param str name: the name of the subsegment. :param str namespace: currently can only be 'local', 'remote', 'aws'. + :param bool sampling: sampling decision for the subsegment being created, + defaults to True """ # Generating the parent dummy segment is necessary. # We don't need to store anything in context. Assumption here From 4aca4473dda841e43cfe0a559caac46727fb9aa9 Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Tue, 8 Nov 2022 14:14:26 -0800 Subject: [PATCH 3/9] created separate API for adding subsegments without sampling --- aws_xray_sdk/core/recorder.py | 36 ++++++++++++++++--- aws_xray_sdk/core/utils/sqs_message_helper.py | 5 +-- aws_xray_sdk/ext/util.py | 1 - tests/test_recorder.py | 8 ++--- tests/test_trace_entities.py | 21 ++++++++++- tests/test_utils.py | 4 +-- 6 files changed, 60 insertions(+), 15 deletions(-) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index 3cdbaa2f..1a3ea75a 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -275,7 +275,7 @@ def current_segment(self): else: return entity - def begin_subsegment(self, name, namespace='local', sampling=True): + def begin_subsegment(self, name, namespace='local'): """ Begin a new subsegment. If there is open subsegment, the newly created subsegment will be the @@ -284,8 +284,6 @@ def begin_subsegment(self, name, namespace='local', sampling=True): :param str name: the name of the subsegment. :param str namespace: currently can only be 'local', 'remote', 'aws'. - :param bool sampling: sampling decision for the subsegment being created, - defaults to True """ # Generating the parent dummy segment is necessary. # We don't need to store anything in context. Assumption here @@ -300,7 +298,7 @@ def begin_subsegment(self, name, namespace='local', sampling=True): current_entity = self.get_trace_entity() - if not current_entity.sampled or not sampling: + if not current_entity.sampled: subsegment = DummySubsegment(segment, name) subsegment.sampled = False else: @@ -310,6 +308,34 @@ def begin_subsegment(self, name, namespace='local', sampling=True): return subsegment + def begin_subsegment_without_sampling(self, name): + """ + Begin a new subsegment. + If there is open subsegment, the newly created subsegment will be the + child of latest opened subsegment. + If not, it will be the child of the current open segment. + + :param str name: the name of the subsegment. + :param str namespace: currently can only be 'local', 'remote', 'aws'. + """ + # Generating the parent dummy segment is necessary. + # We don't need to store anything in context. Assumption here + # is that we only work with recorder-level APIs. + if not global_sdk_config.sdk_enabled(): + return DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)) + + segment = self.current_segment() + if not segment: + log.warning("No segment found, cannot begin subsegment %s." % name) + return None + + subsegment = DummySubsegment(segment, name) + subsegment.sampled = False + + self.context.put_subsegment(subsegment) + + return subsegment + def current_subsegment(self): """ Return the latest opened subsegment. In a multithreading environment, @@ -340,7 +366,7 @@ def end_subsegment(self, end_time=None): # if segment is already close, we check if we can send entire segment # otherwise we check if we need to stream some subsegments - if self.current_segment().ready_to_send(): + if self.current_segment().ready_to_send(): self._send_segment() else: self.stream_subsegments() diff --git a/aws_xray_sdk/core/utils/sqs_message_helper.py b/aws_xray_sdk/core/utils/sqs_message_helper.py index 947176d3..f2a1a1c8 100644 --- a/aws_xray_sdk/core/utils/sqs_message_helper.py +++ b/aws_xray_sdk/core/utils/sqs_message_helper.py @@ -1,10 +1,11 @@ +SQS_XRAY_HEADER = "AWSTraceHeader" class SqsMessageHelper: @staticmethod def isSampled(sqs_message): attributes = sqs_message['attributes'] - if 'AWSTraceHeader' not in attributes: + if SQS_XRAY_HEADER not in attributes: return False - return 'Sampled=1' in attributes['AWSTraceHeader'] \ No newline at end of file + return 'Sampled=1' in attributes[SQS_XRAY_HEADER] \ No newline at end of file diff --git a/aws_xray_sdk/ext/util.py b/aws_xray_sdk/ext/util.py index 69c3fdf7..ad9d5207 100644 --- a/aws_xray_sdk/ext/util.py +++ b/aws_xray_sdk/ext/util.py @@ -35,7 +35,6 @@ def inject_trace_header(headers, entity): else: header = entity.get_origin_trace_header() data = header.data if header else None - to_insert = TraceHeader( root=entity.trace_id, parent=entity.id, diff --git a/tests/test_recorder.py b/tests/test_recorder.py index c99db141..614de01b 100644 --- a/tests/test_recorder.py +++ b/tests/test_recorder.py @@ -145,7 +145,7 @@ def test_unsampled_subsegment_of_sampled_parent(): xray_recorder = get_new_stubbed_recorder() xray_recorder.configure(sampling=True) segment = xray_recorder.begin_segment('name', sampling=True) - subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled') assert segment.sampled == True assert subsegment.sampled == False @@ -154,7 +154,7 @@ def test_begin_subsegment_unsampled(): xray_recorder = get_new_stubbed_recorder() xray_recorder.configure(sampling=False) segment = xray_recorder.begin_segment('name', sampling=False) - subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled') assert segment.sampled == False assert subsegment.sampled == False @@ -222,14 +222,14 @@ def test_disable_is_dummy(): def test_unsampled_subsegment_is_dummy(): assert global_sdk_config.sdk_enabled() segment = xray_recorder.begin_segment('name') - subsegment = xray_recorder.begin_subsegment('name', sampling=False) + subsegment = xray_recorder.begin_subsegment_without_sampling('name') assert type(xray_recorder.current_subsegment()) is DummySubsegment def test_subsegment_respects_parent_sampling_decision(): assert global_sdk_config.sdk_enabled() segment = xray_recorder.begin_segment('name') - subsegment = xray_recorder.begin_subsegment('name2', sampling=False) + subsegment = xray_recorder.begin_subsegment_without_sampling('name2') subsegment2 = xray_recorder.begin_subsegment('unsampled-subsegment') assert type(xray_recorder.current_subsegment()) is DummySubsegment diff --git a/tests/test_trace_entities.py b/tests/test_trace_entities.py index 29300aaa..c0f17c99 100644 --- a/tests/test_trace_entities.py +++ b/tests/test_trace_entities.py @@ -11,6 +11,9 @@ from aws_xray_sdk.core.exceptions.exceptions import AlreadyEndedException from .util import entity_to_dict +from .util import get_new_stubbed_recorder + +xray_recorder = get_new_stubbed_recorder() def test_unicode_entity_name(): @@ -306,4 +309,20 @@ def test_adding_subsegment_respects_parent_sampling2(): assert segment.sampled == True assert subsegment.sampled == True - assert subsubsegment.sampled == True \ No newline at end of file + assert subsubsegment.sampled == True + +def test_adding_subsegments_with_recorder(): + xray_recorder.configure(sampling=False) + xray_recorder.clear_trace_entities() + + segment = xray_recorder.begin_segment('parent'); + subsegment = xray_recorder.begin_subsegment('sampled-child') + unsampled_subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled-child1') + unsampled_child_subsegment = xray_recorder.begin_subsegment('unsampled-child2') + + assert segment.sampled == True + assert subsegment.sampled == True + assert unsampled_subsegment.sampled == False + assert unsampled_child_subsegment.sampled == False + + xray_recorder.clear_trace_entities() \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py index 74001ded..9c35ad84 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -65,7 +65,7 @@ def test_inject_trace_header_unsampled(): xray_recorder = get_new_stubbed_recorder() xray_recorder.configure(sampling=True) segment = xray_recorder.begin_segment('name', sampling=True) - subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled') inject_trace_header(headers, subsegment) @@ -77,7 +77,7 @@ def test_inject_trace_header_respects_parent_subsegment(): xray_recorder = get_new_stubbed_recorder() xray_recorder.configure(sampling=True) segment = xray_recorder.begin_segment('name', sampling=True) - subsegment = xray_recorder.begin_subsegment('unsampled', sampling=False) + subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled') subsegment2 = xray_recorder.begin_subsegment('unsampled2') inject_trace_header(headers, subsegment2) From 4d4b9b54206957d1806034599f942dfd53b7466a Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Wed, 9 Nov 2022 10:12:13 -0800 Subject: [PATCH 4/9] Modified add_subsegment method to log warning for orphaned subsegments --- aws_xray_sdk/core/models/entity.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/aws_xray_sdk/core/models/entity.py b/aws_xray_sdk/core/models/entity.py index 427a368a..3583bf28 100644 --- a/aws_xray_sdk/core/models/entity.py +++ b/aws_xray_sdk/core/models/entity.py @@ -81,7 +81,10 @@ def add_subsegment(self, subsegment): """ self._check_ended() subsegment.parent_id = self.id - subsegment.sampled = self.sampled if subsegment.sampled == True else False + + if not self.sampled and subsegment.sampled: + log.warning("This sampled subsegment is being added to an unsampled parent segment/subsegment and will be orphaned.") + self.subsegments.append(subsegment) def remove_subsegment(self, subsegment): From 8e4ed46ceec60fe9479203c049db406044edbc9a Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Wed, 9 Nov 2022 10:21:43 -0800 Subject: [PATCH 5/9] updated unit tests --- tests/test_facade_segment.py | 18 ------------------ tests/test_trace_entities.py | 23 ----------------------- 2 files changed, 41 deletions(-) diff --git a/tests/test_facade_segment.py b/tests/test_facade_segment.py index 57e12071..5b95115a 100644 --- a/tests/test_facade_segment.py +++ b/tests/test_facade_segment.py @@ -70,21 +70,3 @@ def test_adding_unsampled_subsegment(): assert segment.subsegments[0] is subsegment assert subsegment.subsegments[0] is subsegment2 assert subsegment2.sampled == False - -def test_adding_subsegment_respects_parent_sampling_decision(): - - segment = FacadeSegment('name', 'id', 'id', True) - subsegment = Subsegment('sampled', 'local', segment) - subsegment2 = Subsegment('unsampled', 'local', segment) - subsegment2.sampled = False - - subsegment3 = Subsegment('unsampled-child', 'local', segment) - segment.add_subsegment(subsegment) - subsegment.add_subsegment(subsegment2) - - subsegment2.add_subsegment(subsegment3) - - assert segment.subsegments[0] is subsegment - assert subsegment.subsegments[0] is subsegment2 - assert subsegment2.sampled == False - assert subsegment3.sampled == False \ No newline at end of file diff --git a/tests/test_trace_entities.py b/tests/test_trace_entities.py index c0f17c99..210f09f4 100644 --- a/tests/test_trace_entities.py +++ b/tests/test_trace_entities.py @@ -288,29 +288,6 @@ def test_adding_unsampled_subsegment_to_subsegment(): assert subsegment.sampled == True assert subsubsegment.sampled == False -def test_adding_subsegment_respects_parent_sampling(): - segment = Segment('seg') - subsegment = Subsegment('sub', 'local', segment) - subsegment.sampled = False - segment.add_subsegment(subsegment) - subsubsegment = Subsegment('subsub', 'local', segment) - subsegment.add_subsegment(subsubsegment) - - assert segment.sampled == True - assert subsegment.sampled == False - assert subsubsegment.sampled == False - -def test_adding_subsegment_respects_parent_sampling2(): - segment = Segment('seg') - subsegment = Subsegment('sub', 'local', segment) - segment.add_subsegment(subsegment) - subsubsegment = Subsegment('subsub', 'local', segment) - subsegment.add_subsegment(subsubsegment) - - assert segment.sampled == True - assert subsegment.sampled == True - assert subsubsegment.sampled == True - def test_adding_subsegments_with_recorder(): xray_recorder.configure(sampling=False) xray_recorder.clear_trace_entities() From a99fc6de4f7bbfeb9bcc364ec98ad1151d3698d3 Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Wed, 9 Nov 2022 16:50:07 -0800 Subject: [PATCH 6/9] addressing feedback --- aws_xray_sdk/core/recorder.py | 45 ++++++++++++++++------------------- 1 file changed, 21 insertions(+), 24 deletions(-) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index 1a3ea75a..d9eac595 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -274,6 +274,20 @@ def current_segment(self): return entity.parent_segment else: return entity + + def _begin_subsegment_helper(self, name): + ''' + Helper method to begin_subsegment and begin_subsegment_without_sampling + ''' + result = True + if not global_sdk_config.sdk_enabled(): + result = DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)) + + segment = self.current_segment() + if not segment: + log.warning("No segment found, cannot begin subsegment %s." % name) + result = None + return result, segment def begin_subsegment(self, name, namespace='local'): """ @@ -288,52 +302,35 @@ def begin_subsegment(self, name, namespace='local'): # Generating the parent dummy segment is necessary. # We don't need to store anything in context. Assumption here # is that we only work with recorder-level APIs. - if not global_sdk_config.sdk_enabled(): - return DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)) - - segment = self.current_segment() - if not segment: - log.warning("No segment found, cannot begin subsegment %s." % name) - return None + result, segment = self._begin_subsegment_helper(name) + if type(result) == DummySubsegment or result == None: + return result current_entity = self.get_trace_entity() if not current_entity.sampled: subsegment = DummySubsegment(segment, name) - subsegment.sampled = False else: subsegment = Subsegment(name, namespace, segment) self.context.put_subsegment(subsegment) - return subsegment def begin_subsegment_without_sampling(self, name): """ - Begin a new subsegment. + Begin a new unsampled subsegment. If there is open subsegment, the newly created subsegment will be the child of latest opened subsegment. If not, it will be the child of the current open segment. :param str name: the name of the subsegment. - :param str namespace: currently can only be 'local', 'remote', 'aws'. """ - # Generating the parent dummy segment is necessary. - # We don't need to store anything in context. Assumption here - # is that we only work with recorder-level APIs. - if not global_sdk_config.sdk_enabled(): - return DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)) - - segment = self.current_segment() - if not segment: - log.warning("No segment found, cannot begin subsegment %s." % name) - return None + result, segment = self._begin_subsegment_helper(name) + if type(result) == DummySubsegment or result == None: + return result subsegment = DummySubsegment(segment, name) - subsegment.sampled = False - self.context.put_subsegment(subsegment) - return subsegment def current_subsegment(self): From 05503a833c0d6bfa31f0bd624ed15c99961ae4ff Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Thu, 10 Nov 2022 09:21:49 -0800 Subject: [PATCH 7/9] final design changes --- aws_xray_sdk/core/recorder.py | 48 +++++++++++++++-------------------- tests/test_trace_entities.py | 21 --------------- 2 files changed, 20 insertions(+), 49 deletions(-) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index d9eac595..b5d83426 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -274,20 +274,32 @@ def current_segment(self): return entity.parent_segment else: return entity - - def _begin_subsegment_helper(self, name): + + def _begin_subsegment_helper(self, name, namespace='local', beginWithoutSampling=False): ''' Helper method to begin_subsegment and begin_subsegment_without_sampling ''' - result = True + # Generating the parent dummy segment is necessary. + # We don't need to store anything in context. Assumption here + # is that we only work with recorder-level APIs. if not global_sdk_config.sdk_enabled(): - result = DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)) + return DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)) segment = self.current_segment() if not segment: log.warning("No segment found, cannot begin subsegment %s." % name) - result = None - return result, segment + return None + + current_entity = self.get_trace_entity() + if not current_entity.sampled or beginWithoutSampling: + subsegment = DummySubsegment(segment, name) + else: + subsegment = Subsegment(name, namespace, segment) + + self.context.put_subsegment(subsegment) + return subsegment + + def begin_subsegment(self, name, namespace='local'): """ @@ -299,22 +311,8 @@ def begin_subsegment(self, name, namespace='local'): :param str name: the name of the subsegment. :param str namespace: currently can only be 'local', 'remote', 'aws'. """ - # Generating the parent dummy segment is necessary. - # We don't need to store anything in context. Assumption here - # is that we only work with recorder-level APIs. - result, segment = self._begin_subsegment_helper(name) - if type(result) == DummySubsegment or result == None: - return result - - current_entity = self.get_trace_entity() - - if not current_entity.sampled: - subsegment = DummySubsegment(segment, name) - else: - subsegment = Subsegment(name, namespace, segment) + return self._begin_subsegment_helper(name) - self.context.put_subsegment(subsegment) - return subsegment def begin_subsegment_without_sampling(self, name): """ @@ -325,13 +323,7 @@ def begin_subsegment_without_sampling(self, name): :param str name: the name of the subsegment. """ - result, segment = self._begin_subsegment_helper(name) - if type(result) == DummySubsegment or result == None: - return result - - subsegment = DummySubsegment(segment, name) - self.context.put_subsegment(subsegment) - return subsegment + return self._begin_subsegment_helper(name, beginWithoutSampling=True) def current_subsegment(self): """ diff --git a/tests/test_trace_entities.py b/tests/test_trace_entities.py index 210f09f4..7d987ed0 100644 --- a/tests/test_trace_entities.py +++ b/tests/test_trace_entities.py @@ -267,27 +267,6 @@ def test_add_exception_appending_exceptions(): assert isinstance(segment.cause, dict) assert len(segment.cause['exceptions']) == 2 -def test_adding_unsampled_subsegment(): - segment = Segment('seg') - subsegment = Subsegment('sub', 'local', segment) - subsegment.sampled = False - segment.add_subsegment(subsegment) - - assert subsegment.sampled == False - assert segment.sampled == True - -def test_adding_unsampled_subsegment_to_subsegment(): - segment = Segment('seg') - subsegment = Subsegment('sub', 'local', segment) - segment.add_subsegment(subsegment) - subsubsegment = Subsegment('subsub', 'local', segment) - subsubsegment.sampled = False - subsegment.add_subsegment(subsubsegment) - - assert segment.sampled == True - assert subsegment.sampled == True - assert subsubsegment.sampled == False - def test_adding_subsegments_with_recorder(): xray_recorder.configure(sampling=False) xray_recorder.clear_trace_entities() From 645fb06c93f21120f5e0e5802c68ad72523ee7c2 Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Thu, 10 Nov 2022 09:26:45 -0800 Subject: [PATCH 8/9] remove default namespace value --- aws_xray_sdk/core/recorder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index b5d83426..cd4aa730 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -275,7 +275,7 @@ def current_segment(self): else: return entity - def _begin_subsegment_helper(self, name, namespace='local', beginWithoutSampling=False): + def _begin_subsegment_helper(self, name, namespace, beginWithoutSampling=False): ''' Helper method to begin_subsegment and begin_subsegment_without_sampling ''' @@ -311,7 +311,7 @@ def begin_subsegment(self, name, namespace='local'): :param str name: the name of the subsegment. :param str namespace: currently can only be 'local', 'remote', 'aws'. """ - return self._begin_subsegment_helper(name) + return self._begin_subsegment_helper(name, namespace) def begin_subsegment_without_sampling(self, name): From 91fa8c551d5d845fe6aee64c18ac91f3c2dbed0a Mon Sep 17 00:00:00 2001 From: Carol Abadeer Date: Thu, 10 Nov 2022 09:28:29 -0800 Subject: [PATCH 9/9] minor fix --- aws_xray_sdk/core/recorder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index cd4aa730..ff4f20b5 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -275,7 +275,7 @@ def current_segment(self): else: return entity - def _begin_subsegment_helper(self, name, namespace, beginWithoutSampling=False): + def _begin_subsegment_helper(self, name, namespace='local', beginWithoutSampling=False): ''' Helper method to begin_subsegment and begin_subsegment_without_sampling '''