diff --git a/aws_xray_sdk/core/models/mimic_segment.py b/aws_xray_sdk/core/models/mimic_segment.py index 4cbb0bc9..350681f9 100644 --- a/aws_xray_sdk/core/models/mimic_segment.py +++ b/aws_xray_sdk/core/models/mimic_segment.py @@ -18,6 +18,9 @@ class MimicSegment(Segment): as a node on the service graph. For all purposes, the MimicSegment can be interacted as if it's a real segment, meaning that all methods that exist only in a Segment but not a subsegment is available to be used. + + The following methods are no-ops and will not be sent to the service: + set_rule_name, set_service, and set_user """ def __init__(self, facade_segment, original_segment): @@ -28,10 +31,19 @@ def __init__(self, facade_segment, original_segment): traceid=facade_segment.trace_id, parent_id=facade_segment.id, sampled=facade_segment.sampled) + def set_rule_name(self, rule_name): + pass + + def set_service(self, service_info): + pass + + def set_user(self, user): + pass + def __getstate__(self): """ - Used during serialization. We mark the subsegment properties to let the dataplane know - that we want the mimic segment to be transformed as a subsegment. + Used during serialization. We mark the mimic segment as a subsegment to + let the X-Ray service know to create this mimic segment as a subsegment. """ properties = super(MimicSegment, self).__getstate__() properties['type'] = 'subsegment' diff --git a/aws_xray_sdk/core/serverless_lambda_context.py b/aws_xray_sdk/core/serverless_lambda_context.py index 01f75e3f..5910cc39 100644 --- a/aws_xray_sdk/core/serverless_lambda_context.py +++ b/aws_xray_sdk/core/serverless_lambda_context.py @@ -18,7 +18,7 @@ class ServerlessLambdaContext(LambdaContext): creates a Segment masked as a Subsegment known as a MimicSegment underneath the Lambda-generated Facade Segment. This ensures that middleware->recorder's consequent calls to "put_segment()" will not throw exceptions but instead create - subsegments underneath the lambda-generated segment. This context also + subsegments underneath the lambda-generated Facade Segment. This context also ensures that FacadeSegments exist through underlying calls to _refresh_context(). """ def __init__(self, context_missing='RUNTIME_ERROR'): @@ -31,7 +31,6 @@ def put_segment(self, segment): """ Convert the segment into a mimic segment and append it to FacadeSegment's subsegment list. :param Segment segment: - :return: """ # When putting a segment, convert it to a mimic segment and make it a child of the Facade Segment. parent_facade_segment = self.__get_facade_entity() # type: FacadeSegment @@ -102,8 +101,9 @@ def get_trace_entity(self): def set_trace_entity(self, trace_entity): """ - Store the input trace_entity to local context. It will overwrite all - existing ones if there is any. + Stores the input trace_entity to local context. It will overwrite all + existing ones if there is any. If the entity passed in is a segment, + it will automatically be converted to a mimic segment. """ if type(trace_entity) == Segment: # Convert to a mimic segment. diff --git a/tests/test_mimic_segment.py b/tests/test_mimic_segment.py index 2c7311a7..f235a669 100644 --- a/tests/test_mimic_segment.py +++ b/tests/test_mimic_segment.py @@ -71,6 +71,8 @@ def test_facade_segment_properties(): def test_segment_methods_on_mimic(): # Test to make sure that segment methods exist and function for the Mimic Segment + # And ensure that the methods (other than get/set origin_trace_header) don't modify + # the segment. mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) # type: MimicSegment assert not getattr(mimic_segment, "service", None) assert not getattr(mimic_segment, "user", None) @@ -82,10 +84,6 @@ def test_segment_methods_on_mimic(): assert getattr(original_segment, "ref_counter", None) assert getattr(original_segment, "_subsegments_counter", None) - mimic_segment.set_service("SomeService") - original_segment.set_service("SomeService") - assert original_segment.service == original_segment.service - assert original_segment.get_origin_trace_header() == mimic_segment.get_origin_trace_header() mimic_segment.save_origin_trace_header("someheader") original_segment.save_origin_trace_header("someheader") @@ -93,9 +91,37 @@ def test_segment_methods_on_mimic(): # No exception is thrown test_dict = {"akey": "avalue"} - original_segment.set_aws(test_dict) - original_segment.set_rule_name(test_dict) + test_rule_name = {"arule": "name"} + original_segment.set_aws(test_dict.copy()) + original_segment.set_rule_name(test_rule_name.copy()) original_segment.set_user("SomeUser") - mimic_segment.set_aws(test_dict) - mimic_segment.set_rule_name(test_dict) + original_segment.set_service("SomeService") + mimic_segment.set_aws(test_dict.copy()) + mimic_segment.set_rule_name(test_rule_name.copy()) mimic_segment.set_user("SomeUser") + mimic_segment.set_service("SomeService") + + # Original segment should contain these properties + # but not the mimic segment. + assert getattr(original_segment, "service", None) + assert getattr(original_segment, "user", None) + assert 'xray' in getattr(original_segment, "aws", None) + assert 'sampling_rule_name' in getattr(original_segment, "aws", None)['xray'] + + assert not getattr(mimic_segment, "service", None) + assert not getattr(mimic_segment, "user", None) + # Originally set by rule_name, but no-op so nothing is set. + assert 'xray' not in getattr(mimic_segment, "aws", None) + + # Ensure serialization also doesn't have those properties in mimic but do in original + original_segment_serialized = original_segment.__getstate__() + mimic_segment_serialized = mimic_segment.__getstate__() + + assert 'service' in original_segment_serialized + assert 'user' in original_segment_serialized + assert 'xray' in original_segment_serialized['aws'] + assert 'sampling_rule_name' in original_segment_serialized['aws']['xray'] + + assert 'service' not in mimic_segment_serialized + assert 'user' not in mimic_segment_serialized + assert 'xray' not in mimic_segment_serialized['aws']