From c3f333fcd9321b92fdaaa52a7cc3302a148144dc Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Thu, 7 Feb 2019 15:39:33 -0800 Subject: [PATCH 1/6] Add Serverless Framework Support (#2) * Serverless architecture in this case includes one that utilizes Lambda and API Gateway. * A new "Serverless" context is created to give the abstraction of Segments being the toplevel entities but is then converted to a subsegment upon transmission to the data plane. These segments are called MimicSegments. All generated segments have a parent segment that is the FacadeSegment. * Currently supports Flask and Django as middlewares; this has been confirmed to be natively working with Zappa if the application is running under Flask/Django. --- aws_xray_sdk/core/exceptions/exceptions.py | 4 + aws_xray_sdk/core/models/mimic_segment.py | 35 ++++ aws_xray_sdk/core/recorder.py | 3 +- aws_xray_sdk/core/serverless_context.py | 131 ++++++++++++++ aws_xray_sdk/ext/django/middleware.py | 7 + aws_xray_sdk/ext/flask/middleware.py | 7 + tests/test_mimic_segment.py | 93 ++++++++++ tests/test_serverless_context.py | 197 +++++++++++++++++++++ 8 files changed, 476 insertions(+), 1 deletion(-) create mode 100644 aws_xray_sdk/core/models/mimic_segment.py create mode 100644 aws_xray_sdk/core/serverless_context.py create mode 100644 tests/test_mimic_segment.py create mode 100644 tests/test_serverless_context.py diff --git a/aws_xray_sdk/core/exceptions/exceptions.py b/aws_xray_sdk/core/exceptions/exceptions.py index cbf4e703..5970a6d6 100644 --- a/aws_xray_sdk/core/exceptions/exceptions.py +++ b/aws_xray_sdk/core/exceptions/exceptions.py @@ -22,6 +22,10 @@ class FacadeSegmentMutationException(Exception): pass +class MimicSegmentInvalidException(Exception): + pass + + class MissingPluginNames(Exception): pass diff --git a/aws_xray_sdk/core/models/mimic_segment.py b/aws_xray_sdk/core/models/mimic_segment.py new file mode 100644 index 00000000..4c58ff69 --- /dev/null +++ b/aws_xray_sdk/core/models/mimic_segment.py @@ -0,0 +1,35 @@ +from .segment import Segment +from ..exceptions.exceptions import MimicSegmentInvalidException + + +class MimicSegment(Segment): + """ + The MimicSegment is an entity that mimics a segment for the use of the serverless context. + When the MimicSegment is generated, its parent segment is assigned to be the FacadeSegment + generated by the Lambda Environment. Upon serialization and transmission of the MimicSegment, + it is converted to a locally-namespaced, subsegment. This is only done during serialization. + All Segment-related method calls done on this object are valid. + + Subsegments are automatically created with the namespace "local" to prevent it from appearing + 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. + """ + + def __init__(self, facade_segment=None, original_segment=None): + if not original_segment or not facade_segment: + raise MimicSegmentInvalidException("Invalid MimicSegment construction. " + "Please put in the original segment and the facade segment.") + super(MimicSegment, self).__init__(name=original_segment.name, entityid=original_segment.id, + traceid=facade_segment.trace_id, parent_id=facade_segment.id, + sampled=facade_segment.sampled) + + def __getstate__(self): + """ + Used during serialization. We mark the subsegment properties to let the dataplane know + that we want the mimic segment to be represented as a subsegment. + """ + properties = super(MimicSegment, self).__getstate__() + properties['type'] = 'subsegment' + properties['namespace'] = 'local' + return properties diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index b953cece..5d423ab9 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -243,7 +243,8 @@ def begin_segment(self, name=None, traceid=None, self._populate_runtime_context(segment, decision) self.context.put_segment(segment) - return segment + current_segment = self.get_trace_entity() + return current_segment def end_segment(self, end_time=None): """ diff --git a/aws_xray_sdk/core/serverless_context.py b/aws_xray_sdk/core/serverless_context.py new file mode 100644 index 00000000..2fff9fe5 --- /dev/null +++ b/aws_xray_sdk/core/serverless_context.py @@ -0,0 +1,131 @@ +import os +import logging + +from .models.facade_segment import FacadeSegment +from .models.segment import Segment +from .models.mimic_segment import MimicSegment +from .context import CXT_MISSING_STRATEGY_KEY +from .lambda_launcher import LambdaContext +from .context import Context + + +log = logging.getLogger(__name__) + + +class ServerlessContext(LambdaContext): + """ + Context used specifically for running middlewares on Lambda through the + Serverless design. This context is built on top of the LambdaContext, but + 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 + ensures that FacadeSegments exist through underlying calls to _refresh_context(). + """ + def __init__(self, context_missing='RUNTIME_ERROR'): + super(ServerlessContext, self).__init__() + + strategy = os.getenv(CXT_MISSING_STRATEGY_KEY, context_missing) + self._context_missing = strategy + + 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 + mimic_segment = MimicSegment(parent_facade_segment, segment) + parent_facade_segment.add_subsegment(mimic_segment) + Context.put_segment(self, mimic_segment) + + def end_segment(self, end_time=None): + """ + Close the MimicSegment + """ + # Close the last mimic segment opened then remove it from our facade segment. + mimic_segment = self.get_trace_entity() + Context.end_segment(self, end_time) + if type(mimic_segment) == MimicSegment: + # The facade segment can only hold mimic segments. + facade_segment = self.__get_facade_entity() + facade_segment.remove_subsegment(mimic_segment) + + def put_subsegment(self, subsegment): + """ + Appends the subsegment as a subsegment of either the mimic segment or + another subsegment if they are the last opened entity. + :param subsegment: The subsegment to to be added as a subsegment. + """ + Context.put_subsegment(self, subsegment) + + def end_subsegment(self, end_time=None): + """ + End the current subsegment. In our case, subsegments + will either be a subsegment of a mimic segment or another + subsegment. + :param int end_time: epoch in seconds. If not specified the current + system time will be used. + :return: True on success, false if no parent mimic segment/subsegment is found. + """ + return Context.end_subsegment(self, end_time) + + def __get_facade_entity(self): + """ + Retrieves the Facade segment from thread local. This facade segment should always be present + because it was generated by the Lambda Container. + :return: FacadeSegment + """ + self._refresh_context() + facade_segment = self._local.segment # type: FacadeSegment + return facade_segment + + def get_trace_entity(self): + """ + Return the latest entity added. In this case, it'll either be a Mimic Segment or + a subsegment. Facade Segments are never returned. + If no mimic segments or subsegments were ever passed in, throw the default + context missing error. + :return: Entity + """ + # Call to Context.get_trace_entity() returns the latest mimic segment/subsegment if they exist. + # Otherwise, returns None through the following way: + # No mimic segment/subsegment exists so Context calls LambdaContext's handle_context_missing(). + # By default, Lambda's method returns no-op, so it will return None to ServerlessContext. + # Take that None as an indication to return the rightful handle_context_missing(), otherwise + # return the entity. + entity = Context.get_trace_entity(self) + if entity is None: + return Context.handle_context_missing(self) + else: + return entity + + 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. + """ + if type(trace_entity) == Segment: + # Convert to a mimic segment. + parent_facade_segment = self.__get_facade_entity() # type: FacadeSegment + converted_segment = MimicSegment(parent_facade_segment, trace_entity) + mimic_segment = converted_segment + else: + # Should be a Mimic Segment. If it's a subsegment, grandparent Context's + # behavior would be invoked. + mimic_segment = trace_entity + + Context.set_trace_entity(self, mimic_segment) + self.__get_facade_entity().subsegments = [mimic_segment] + + def _is_subsegment(self, entity): + return super(ServerlessContext, self)._is_subsegment(entity) and type(entity) != MimicSegment + + @property + def context_missing(self): + return self._context_missing + + @context_missing.setter + def context_missing(self, value): + self._context_missing = value diff --git a/aws_xray_sdk/ext/django/middleware.py b/aws_xray_sdk/ext/django/middleware.py index 9f071de9..39ffffb0 100644 --- a/aws_xray_sdk/ext/django/middleware.py +++ b/aws_xray_sdk/ext/django/middleware.py @@ -1,7 +1,9 @@ import logging from aws_xray_sdk.core import xray_recorder +from aws_xray_sdk.core.lambda_launcher import check_in_lambda from aws_xray_sdk.core.models import http +from aws_xray_sdk.core.serverless_context import ServerlessContext from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header @@ -25,6 +27,11 @@ def __init__(self, get_response): self.get_response = get_response + # The case when the middleware is initialized in a Lambda Context, we make sure + # to use the ServerlessContext so that the middleware properly functions. + if check_in_lambda() is not None: + xray_recorder.context = ServerlessContext() + # hooks for django version >= 1.10 def __call__(self, request): diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index 9e0b4877..65661076 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -1,7 +1,9 @@ import flask.templating from flask import request +from aws_xray_sdk.core.lambda_launcher import check_in_lambda from aws_xray_sdk.core.models import http +from aws_xray_sdk.core.serverless_context import ServerlessContext from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header @@ -18,6 +20,11 @@ def __init__(self, app, recorder): self.app.after_request(self._after_request) self.app.teardown_request(self._handle_exception) + # The case when the middleware is initialized in a Lambda Context, we make sure + # to use the ServerlessContext so that the middleware properly functions. + if check_in_lambda() is not None: + self._recorder.context = ServerlessContext() + _patch_render(recorder) def _before_request(self): diff --git a/tests/test_mimic_segment.py b/tests/test_mimic_segment.py new file mode 100644 index 00000000..f3320b74 --- /dev/null +++ b/tests/test_mimic_segment.py @@ -0,0 +1,93 @@ +import pytest + +from aws_xray_sdk.core.models.facade_segment import FacadeSegment +from aws_xray_sdk.core.models.segment import Segment +from aws_xray_sdk.core.models.subsegment import Subsegment +from aws_xray_sdk.core.models.mimic_segment import MimicSegment +from aws_xray_sdk.core.exceptions.exceptions import MimicSegmentInvalidException + + +original_segment = Segment("RealSegment") +facade_segment = FacadeSegment("FacadeSegment", "entityid", "traceid", True) + + +@pytest.fixture(autouse=True) +def cleanup_ctx(): + global original_segment, facade_segment + original_segment = Segment("RealSegment") + facade_segment = FacadeSegment("FacadeSegment", "entityid", "traceid", True) + yield + original_segment = Segment("RealSegment") + facade_segment = FacadeSegment("FacadeSegment", "entityid", "traceid", True) + + +def test_ready(): + mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) + mimic_segment.in_progress = False + assert mimic_segment.ready_to_send() + + +def test_invalid_init(): + with pytest.raises(MimicSegmentInvalidException): + MimicSegment(facade_segment=None, original_segment=original_segment) + MimicSegment(facade_segment=facade_segment, original_segment=None) + MimicSegment(facade_segment=Subsegment("Test", "local", original_segment), original_segment=None) + MimicSegment(facade_segment=None, original_segment=Subsegment("Test", "local", original_segment)) + + +def test_init_similar(): + mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) # type: MimicSegment + + assert mimic_segment.id == original_segment.id + assert mimic_segment.name == original_segment.name + assert mimic_segment.in_progress == original_segment.in_progress + + assert mimic_segment.trace_id == facade_segment.trace_id + assert mimic_segment.parent_id == facade_segment.id + assert mimic_segment.sampled == facade_segment.sampled + + mimic_segment_serialized = mimic_segment.__getstate__() + assert mimic_segment_serialized['namespace'] == "local" + assert mimic_segment_serialized['type'] == "subsegment" + + +def test_facade_segment_properties(): + # Sampling decision is made by Facade Segment + original_segment.sampled = False + facade_segment.sampled = True + mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) # type: MimicSegment + + assert mimic_segment.sampled == facade_segment.sampled + assert mimic_segment.sampled != original_segment.sampled + + +def test_segment_methods_on_mimic(): + # Test to make sure that segment methods exist and function for the Mimic 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) + assert getattr(mimic_segment, "ref_counter", None) + assert getattr(mimic_segment, "_subsegments_counter", None) + + assert not getattr(original_segment, "service", None) + assert not getattr(original_segment, "user", None) + 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") + assert original_segment.get_origin_trace_header() == mimic_segment.get_origin_trace_header() + + # No exception is thrown + test_dict = {"akey": "avalue"} + original_segment.set_aws(test_dict) + original_segment.set_rule_name(test_dict) + original_segment.set_user("SomeUser") + mimic_segment.set_aws(test_dict) + mimic_segment.set_rule_name(test_dict) + mimic_segment.set_user("SomeUser") diff --git a/tests/test_serverless_context.py b/tests/test_serverless_context.py new file mode 100644 index 00000000..e591cb25 --- /dev/null +++ b/tests/test_serverless_context.py @@ -0,0 +1,197 @@ +import os +import pytest + +from aws_xray_sdk.core import serverless_context +from aws_xray_sdk.core import context +from aws_xray_sdk.core.lambda_launcher import LAMBDA_TRACE_HEADER_KEY +from aws_xray_sdk.core.exceptions.exceptions import AlreadyEndedException, SegmentNotFoundException +from aws_xray_sdk.core.models.segment import Segment +from aws_xray_sdk.core.models.subsegment import Subsegment +from aws_xray_sdk.core.models.mimic_segment import MimicSegment +from aws_xray_sdk.core.models.facade_segment import FacadeSegment + + +TRACE_ID = '1-5759e988-bd862e3fe1be46a994272793' +PARENT_ID = '53995c3f42cd8ad8' +HEADER_VAR = "Root=%s;Parent=%s;Sampled=1" % (TRACE_ID, PARENT_ID) + +os.environ[LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR +context = serverless_context.ServerlessContext() + +service_name = "Test Flask Server" + + +@pytest.fixture(autouse=True) +def cleanup_ctx(): + context.clear_trace_entities() + yield + context.clear_trace_entities() + + +def test_segment_generation(): + # Ensure we create Mimic Segments, and that parents of Mimic segments are Facade Segments. + segment = Segment(service_name) + context.put_segment(segment) + + mimic_segment = context.get_trace_entity() + assert type(mimic_segment) == MimicSegment + + facade_segment = getattr(context._local, 'segment', None) + assert type(facade_segment) == FacadeSegment + assert mimic_segment.parent_id == facade_segment.id + + assert facade_segment.id == PARENT_ID + assert facade_segment.trace_id == TRACE_ID + assert facade_segment.sampled + + +def test_facade_in_threadlocal(): + # Ensure that facade segments are stored in threadlocal.segment + assert not getattr(context._local, 'segment', None) + + # Refresh context to generate the facade segment. + context._refresh_context() + facade_segment = getattr(context._local, 'segment', None) + assert facade_segment + assert type(facade_segment) == FacadeSegment + assert facade_segment.id == PARENT_ID + assert facade_segment.trace_id == TRACE_ID + + +def test_put_subsegment(): + segment = Segment(service_name) + context.put_segment(segment) + + segment = context.get_trace_entity() + subsegment = Subsegment('name', 'local', segment) + context.put_subsegment(subsegment) + assert context.get_trace_entity().id == subsegment.id + + subsegment2 = Subsegment('name', 'local', segment) + context.put_subsegment(subsegment2) + assert context.get_trace_entity().id == subsegment2.id + + assert subsegment.subsegments[0] is subsegment2 + assert subsegment2.parent_id == subsegment.id + assert subsegment.parent_id == segment.id + assert subsegment2.parent_segment is segment + + assert context.get_trace_entity().id == subsegment2.id + + context.end_subsegment() + assert context.get_trace_entity().id == subsegment.id + + context.end_subsegment() + assert context.get_trace_entity().id == segment.id + assert context.get_trace_entity().in_progress + + context.end_segment() + assert context.get_trace_entity().id == segment.id + assert not context.get_trace_entity().in_progress + + +def test_remote_mimic_segment(): + # Ensure that the mimic-generated segment is set as a subsegment + # when being serialized through jsonpickle (which uses .__getstate__()) + segment = Segment(service_name) + context.put_segment(segment) + + mimic_segment = context.get_trace_entity() + + segment_properties = mimic_segment.__getstate__() + assert segment_properties['namespace'] == "local" + assert segment_properties['type'] == "subsegment" + + +def test_segment_methods_on_mimic_segment(): + # Ensure that segment operations made on the mimic segment all works. + comparison_segment = Segment(service_name) # type: Segment + context.put_segment(comparison_segment) + mimic_segment = context.get_trace_entity() # type: MimicSegment + + trace_header = "Someheader" + comparison_segment.save_origin_trace_header(trace_header) + mimic_segment.save_origin_trace_header(trace_header) + assert mimic_segment.get_origin_trace_header() == comparison_segment.get_origin_trace_header() + + assert mimic_segment.get_total_subsegments_size() == comparison_segment.get_total_subsegments_size() + comparison_segment.increment() + mimic_segment.increment() + assert mimic_segment.get_total_subsegments_size() == comparison_segment.get_total_subsegments_size() + + comparison_segment.decrement_subsegments_size() + assert mimic_segment.get_total_subsegments_size() != comparison_segment.get_total_subsegments_size() + mimic_segment.decrement_subsegments_size() + assert mimic_segment.get_total_subsegments_size() == comparison_segment.get_total_subsegments_size() + + assert mimic_segment.ready_to_send() == comparison_segment.ready_to_send() + + +def test_empty_context(): + # Test to make sure an touched context is absolutely clear. + # Call to get_trace_entity should produce a facade segment. + assert not context._local.__dict__ + + # Empty context should throw an exception + with pytest.raises(SegmentNotFoundException): + context.get_trace_entity() + + # Induce the creation of a facade segment by putting in a segment. + assert len(context._local.__dict__) == 0 + segment = Segment(service_name) + context.put_segment(segment) + + assert type(context._local.segment) == FacadeSegment + assert len(context._local.entities) == 1 + + +def test_set_trace_entity(): + segment_one = Segment(service_name) + context.put_segment(segment_one) + first_mimic_segment = context.get_trace_entity() + facade_segment = getattr(context._local, 'segment', None) + + segment_two = Segment("WOOH") + context.set_trace_entity(segment_two) + second_mimic_segment = context.get_trace_entity() + + assert first_mimic_segment.id == segment_one.id + assert first_mimic_segment.name == segment_one.name + assert first_mimic_segment.parent_id == facade_segment.id + assert first_mimic_segment.trace_id == facade_segment.trace_id + assert second_mimic_segment.id == segment_two.id + assert second_mimic_segment.name == segment_two.name + assert second_mimic_segment.parent_id == facade_segment.id + assert second_mimic_segment.trace_id == facade_segment.trace_id + + +def test_segment_close_subsegment_open(): + # Tests to make sure that when the parent, mimic segment is closed, + # and the last entity is a subsegment, the segment itself closes. + segment = Segment(service_name) + context.put_segment(segment) + mimic_segment = context.get_trace_entity() + assert mimic_segment.in_progress + subsegment = Subsegment("test", "local", mimic_segment) + context.put_subsegment(subsegment) + context.end_segment() + assert not mimic_segment.in_progress + + +def test_begin_close_twice(): + segment_one = Segment(service_name) + context.put_segment(segment_one) + context.end_segment() + entity_one = context.get_trace_entity() + context.put_segment(segment_one) + context.end_segment() + entity_two = context.get_trace_entity() + assert entity_one != entity_two + + +def test_cant_end_segment_twice(): + segment_one = Segment(service_name) + context.put_segment(segment_one) + context.end_segment() + with pytest.raises(AlreadyEndedException): + context.end_segment() From 88bf381393c795bbc212a3ea1373ad3da921712a Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Wed, 13 Feb 2019 11:10:16 -0800 Subject: [PATCH 2/6] Add Serverless Framework Support Revision 2 (#127) * Renamed ServerlessContext to ServerlessLambdaContext to be more precise about its purpose. * Instead of explicitly using Context class's put_subsegment/put_segment methods, use LambdaContext's grandparent class to store. LambdaContext's methods are really only used to acquire the FacadeSegment. * Parameters when initializing the mimic segment are now required. * Unit tests added to test these changes. --- aws_xray_sdk/core/models/mimic_segment.py | 12 ++++++---- ...ontext.py => serverless_lambda_context.py} | 23 +++++++++---------- aws_xray_sdk/ext/django/middleware.py | 12 ++++++---- aws_xray_sdk/ext/flask/middleware.py | 12 ++++++---- tests/test_mimic_segment.py | 8 +++++++ ...t.py => test_serverless_lambda_context.py} | 5 ++-- 6 files changed, 43 insertions(+), 29 deletions(-) rename aws_xray_sdk/core/{serverless_context.py => serverless_lambda_context.py} (87%) rename tests/{test_serverless_context.py => test_serverless_lambda_context.py} (98%) diff --git a/aws_xray_sdk/core/models/mimic_segment.py b/aws_xray_sdk/core/models/mimic_segment.py index 4c58ff69..4cbb0bc9 100644 --- a/aws_xray_sdk/core/models/mimic_segment.py +++ b/aws_xray_sdk/core/models/mimic_segment.py @@ -1,10 +1,14 @@ from .segment import Segment +from .facade_segment import FacadeSegment from ..exceptions.exceptions import MimicSegmentInvalidException class MimicSegment(Segment): """ - The MimicSegment is an entity that mimics a segment for the use of the serverless context. + The MimicSegment is an entity that mimics a segment. It's primary use is for a special-case + in Lambda; specifically, for the Serverless Design Pattern. It is not recommended to use + this MimicSegment for any other purpose. + When the MimicSegment is generated, its parent segment is assigned to be the FacadeSegment generated by the Lambda Environment. Upon serialization and transmission of the MimicSegment, it is converted to a locally-namespaced, subsegment. This is only done during serialization. @@ -16,8 +20,8 @@ class MimicSegment(Segment): is available to be used. """ - def __init__(self, facade_segment=None, original_segment=None): - if not original_segment or not facade_segment: + def __init__(self, facade_segment, original_segment): + if not issubclass(type(original_segment), Segment) or type(facade_segment) != FacadeSegment: raise MimicSegmentInvalidException("Invalid MimicSegment construction. " "Please put in the original segment and the facade segment.") super(MimicSegment, self).__init__(name=original_segment.name, entityid=original_segment.id, @@ -27,7 +31,7 @@ def __init__(self, facade_segment=None, original_segment=None): def __getstate__(self): """ Used during serialization. We mark the subsegment properties to let the dataplane know - that we want the mimic segment to be represented as a subsegment. + that we want the mimic segment to be transformed as a subsegment. """ properties = super(MimicSegment, self).__getstate__() properties['type'] = 'subsegment' diff --git a/aws_xray_sdk/core/serverless_context.py b/aws_xray_sdk/core/serverless_lambda_context.py similarity index 87% rename from aws_xray_sdk/core/serverless_context.py rename to aws_xray_sdk/core/serverless_lambda_context.py index 2fff9fe5..01f75e3f 100644 --- a/aws_xray_sdk/core/serverless_context.py +++ b/aws_xray_sdk/core/serverless_lambda_context.py @@ -6,13 +6,12 @@ from .models.mimic_segment import MimicSegment from .context import CXT_MISSING_STRATEGY_KEY from .lambda_launcher import LambdaContext -from .context import Context log = logging.getLogger(__name__) -class ServerlessContext(LambdaContext): +class ServerlessLambdaContext(LambdaContext): """ Context used specifically for running middlewares on Lambda through the Serverless design. This context is built on top of the LambdaContext, but @@ -23,7 +22,7 @@ class ServerlessContext(LambdaContext): ensures that FacadeSegments exist through underlying calls to _refresh_context(). """ def __init__(self, context_missing='RUNTIME_ERROR'): - super(ServerlessContext, self).__init__() + super(ServerlessLambdaContext, self).__init__() strategy = os.getenv(CXT_MISSING_STRATEGY_KEY, context_missing) self._context_missing = strategy @@ -38,7 +37,7 @@ def put_segment(self, segment): parent_facade_segment = self.__get_facade_entity() # type: FacadeSegment mimic_segment = MimicSegment(parent_facade_segment, segment) parent_facade_segment.add_subsegment(mimic_segment) - Context.put_segment(self, mimic_segment) + super(LambdaContext, self).put_segment(mimic_segment) def end_segment(self, end_time=None): """ @@ -46,7 +45,7 @@ def end_segment(self, end_time=None): """ # Close the last mimic segment opened then remove it from our facade segment. mimic_segment = self.get_trace_entity() - Context.end_segment(self, end_time) + super(LambdaContext, self).end_segment(end_time) if type(mimic_segment) == MimicSegment: # The facade segment can only hold mimic segments. facade_segment = self.__get_facade_entity() @@ -58,7 +57,7 @@ def put_subsegment(self, subsegment): another subsegment if they are the last opened entity. :param subsegment: The subsegment to to be added as a subsegment. """ - Context.put_subsegment(self, subsegment) + super(LambdaContext, self).put_subsegment(subsegment) def end_subsegment(self, end_time=None): """ @@ -69,7 +68,7 @@ def end_subsegment(self, end_time=None): system time will be used. :return: True on success, false if no parent mimic segment/subsegment is found. """ - return Context.end_subsegment(self, end_time) + return super(LambdaContext, self).end_subsegment(end_time) def __get_facade_entity(self): """ @@ -92,12 +91,12 @@ def get_trace_entity(self): # Call to Context.get_trace_entity() returns the latest mimic segment/subsegment if they exist. # Otherwise, returns None through the following way: # No mimic segment/subsegment exists so Context calls LambdaContext's handle_context_missing(). - # By default, Lambda's method returns no-op, so it will return None to ServerlessContext. + # By default, Lambda's method returns no-op, so it will return None to ServerlessLambdaContext. # Take that None as an indication to return the rightful handle_context_missing(), otherwise # return the entity. - entity = Context.get_trace_entity(self) + entity = super(LambdaContext, self).get_trace_entity() if entity is None: - return Context.handle_context_missing(self) + return super(LambdaContext, self).handle_context_missing() else: return entity @@ -116,11 +115,11 @@ def set_trace_entity(self, trace_entity): # behavior would be invoked. mimic_segment = trace_entity - Context.set_trace_entity(self, mimic_segment) + super(LambdaContext, self).set_trace_entity(mimic_segment) self.__get_facade_entity().subsegments = [mimic_segment] def _is_subsegment(self, entity): - return super(ServerlessContext, self)._is_subsegment(entity) and type(entity) != MimicSegment + return super(ServerlessLambdaContext, self)._is_subsegment(entity) and type(entity) != MimicSegment @property def context_missing(self): diff --git a/aws_xray_sdk/ext/django/middleware.py b/aws_xray_sdk/ext/django/middleware.py index 39ffffb0..07b732d9 100644 --- a/aws_xray_sdk/ext/django/middleware.py +++ b/aws_xray_sdk/ext/django/middleware.py @@ -1,9 +1,9 @@ import logging from aws_xray_sdk.core import xray_recorder -from aws_xray_sdk.core.lambda_launcher import check_in_lambda +from aws_xray_sdk.core.lambda_launcher import check_in_lambda, LambdaContext from aws_xray_sdk.core.models import http -from aws_xray_sdk.core.serverless_context import ServerlessContext +from aws_xray_sdk.core.serverless_lambda_context import ServerlessLambdaContext from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header @@ -28,9 +28,11 @@ def __init__(self, get_response): self.get_response = get_response # The case when the middleware is initialized in a Lambda Context, we make sure - # to use the ServerlessContext so that the middleware properly functions. - if check_in_lambda() is not None: - xray_recorder.context = ServerlessContext() + # to use the ServerlessLambdaContext so that the middleware properly functions. + # We also check if the current context is a LambdaContext to not override customer + # provided contexts. + if check_in_lambda() is not None and type(xray_recorder.context) == LambdaContext: + xray_recorder.context = ServerlessLambdaContext() # hooks for django version >= 1.10 def __call__(self, request): diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index 65661076..74ec0d9e 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -1,9 +1,9 @@ import flask.templating from flask import request -from aws_xray_sdk.core.lambda_launcher import check_in_lambda +from aws_xray_sdk.core.lambda_launcher import check_in_lambda, LambdaContext from aws_xray_sdk.core.models import http -from aws_xray_sdk.core.serverless_context import ServerlessContext +from aws_xray_sdk.core.serverless_lambda_context import ServerlessLambdaContext from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header @@ -21,9 +21,11 @@ def __init__(self, app, recorder): self.app.teardown_request(self._handle_exception) # The case when the middleware is initialized in a Lambda Context, we make sure - # to use the ServerlessContext so that the middleware properly functions. - if check_in_lambda() is not None: - self._recorder.context = ServerlessContext() + # to use the ServerlessLambdaContext so that the middleware properly functions. + # We also check if the current context is a LambdaContext to not override customer + # provided contexts. + if check_in_lambda() is not None and type(self._recorder.context) == LambdaContext: + self._recorder.context = ServerlessLambdaContext() _patch_render(recorder) diff --git a/tests/test_mimic_segment.py b/tests/test_mimic_segment.py index f3320b74..2c7311a7 100644 --- a/tests/test_mimic_segment.py +++ b/tests/test_mimic_segment.py @@ -30,9 +30,17 @@ def test_ready(): def test_invalid_init(): with pytest.raises(MimicSegmentInvalidException): MimicSegment(facade_segment=None, original_segment=original_segment) + with pytest.raises(MimicSegmentInvalidException): MimicSegment(facade_segment=facade_segment, original_segment=None) + with pytest.raises(MimicSegmentInvalidException): MimicSegment(facade_segment=Subsegment("Test", "local", original_segment), original_segment=None) + with pytest.raises(MimicSegmentInvalidException): MimicSegment(facade_segment=None, original_segment=Subsegment("Test", "local", original_segment)) + with pytest.raises(MimicSegmentInvalidException): + MimicSegment(facade_segment=facade_segment, original_segment=Subsegment("Test", "local", original_segment)) + with pytest.raises(MimicSegmentInvalidException): + MimicSegment(facade_segment=original_segment, original_segment=facade_segment) + MimicSegment(facade_segment=facade_segment, original_segment=original_segment) def test_init_similar(): diff --git a/tests/test_serverless_context.py b/tests/test_serverless_lambda_context.py similarity index 98% rename from tests/test_serverless_context.py rename to tests/test_serverless_lambda_context.py index e591cb25..54a65b3a 100644 --- a/tests/test_serverless_context.py +++ b/tests/test_serverless_lambda_context.py @@ -1,8 +1,7 @@ import os import pytest -from aws_xray_sdk.core import serverless_context -from aws_xray_sdk.core import context +from aws_xray_sdk.core.serverless_lambda_context import ServerlessLambdaContext from aws_xray_sdk.core.lambda_launcher import LAMBDA_TRACE_HEADER_KEY from aws_xray_sdk.core.exceptions.exceptions import AlreadyEndedException, SegmentNotFoundException from aws_xray_sdk.core.models.segment import Segment @@ -16,7 +15,7 @@ HEADER_VAR = "Root=%s;Parent=%s;Sampled=1" % (TRACE_ID, PARENT_ID) os.environ[LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR -context = serverless_context.ServerlessContext() +context = ServerlessLambdaContext() service_name = "Test Flask Server" From c8e8012f8f2f494e2b68590d2e99dc8ab919cc70 Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Wed, 20 Feb 2019 16:14:33 -0800 Subject: [PATCH 3/6] Add Serverless Framework Support Revision 3 (#127) * Mimic Segments now perform no-op on segment-only methods. * Added comments to address this contract. * Unit tests to ensure segment-only methods don't add data to the mimic segment. Serialization methods are also tested to ensure this. --- aws_xray_sdk/core/models/mimic_segment.py | 16 ++++++- .../core/serverless_lambda_context.py | 8 ++-- tests/test_mimic_segment.py | 42 +++++++++++++++---- 3 files changed, 52 insertions(+), 14 deletions(-) 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'] From 033e03fcefaaac0087a04d585b654697973e109b Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Fri, 22 Feb 2019 11:34:44 -0800 Subject: [PATCH 4/6] Revert back to master before serverless implementation. --- aws_xray_sdk/core/exceptions/exceptions.py | 4 - aws_xray_sdk/core/models/mimic_segment.py | 51 ----- aws_xray_sdk/core/recorder.py | 3 +- .../core/serverless_lambda_context.py | 130 ------------ aws_xray_sdk/ext/django/middleware.py | 9 - aws_xray_sdk/ext/flask/middleware.py | 9 - tests/test_mimic_segment.py | 127 ------------ tests/test_serverless_lambda_context.py | 196 ------------------ 8 files changed, 1 insertion(+), 528 deletions(-) delete mode 100644 aws_xray_sdk/core/models/mimic_segment.py delete mode 100644 aws_xray_sdk/core/serverless_lambda_context.py delete mode 100644 tests/test_mimic_segment.py delete mode 100644 tests/test_serverless_lambda_context.py diff --git a/aws_xray_sdk/core/exceptions/exceptions.py b/aws_xray_sdk/core/exceptions/exceptions.py index 5970a6d6..cbf4e703 100644 --- a/aws_xray_sdk/core/exceptions/exceptions.py +++ b/aws_xray_sdk/core/exceptions/exceptions.py @@ -22,10 +22,6 @@ class FacadeSegmentMutationException(Exception): pass -class MimicSegmentInvalidException(Exception): - pass - - class MissingPluginNames(Exception): pass diff --git a/aws_xray_sdk/core/models/mimic_segment.py b/aws_xray_sdk/core/models/mimic_segment.py deleted file mode 100644 index 350681f9..00000000 --- a/aws_xray_sdk/core/models/mimic_segment.py +++ /dev/null @@ -1,51 +0,0 @@ -from .segment import Segment -from .facade_segment import FacadeSegment -from ..exceptions.exceptions import MimicSegmentInvalidException - - -class MimicSegment(Segment): - """ - The MimicSegment is an entity that mimics a segment. It's primary use is for a special-case - in Lambda; specifically, for the Serverless Design Pattern. It is not recommended to use - this MimicSegment for any other purpose. - - When the MimicSegment is generated, its parent segment is assigned to be the FacadeSegment - generated by the Lambda Environment. Upon serialization and transmission of the MimicSegment, - it is converted to a locally-namespaced, subsegment. This is only done during serialization. - All Segment-related method calls done on this object are valid. - - Subsegments are automatically created with the namespace "local" to prevent it from appearing - 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): - if not issubclass(type(original_segment), Segment) or type(facade_segment) != FacadeSegment: - raise MimicSegmentInvalidException("Invalid MimicSegment construction. " - "Please put in the original segment and the facade segment.") - super(MimicSegment, self).__init__(name=original_segment.name, entityid=original_segment.id, - 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 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' - properties['namespace'] = 'local' - return properties diff --git a/aws_xray_sdk/core/recorder.py b/aws_xray_sdk/core/recorder.py index 5d423ab9..b953cece 100644 --- a/aws_xray_sdk/core/recorder.py +++ b/aws_xray_sdk/core/recorder.py @@ -243,8 +243,7 @@ def begin_segment(self, name=None, traceid=None, self._populate_runtime_context(segment, decision) self.context.put_segment(segment) - current_segment = self.get_trace_entity() - return current_segment + return segment def end_segment(self, end_time=None): """ diff --git a/aws_xray_sdk/core/serverless_lambda_context.py b/aws_xray_sdk/core/serverless_lambda_context.py deleted file mode 100644 index 5910cc39..00000000 --- a/aws_xray_sdk/core/serverless_lambda_context.py +++ /dev/null @@ -1,130 +0,0 @@ -import os -import logging - -from .models.facade_segment import FacadeSegment -from .models.segment import Segment -from .models.mimic_segment import MimicSegment -from .context import CXT_MISSING_STRATEGY_KEY -from .lambda_launcher import LambdaContext - - -log = logging.getLogger(__name__) - - -class ServerlessLambdaContext(LambdaContext): - """ - Context used specifically for running middlewares on Lambda through the - Serverless design. This context is built on top of the LambdaContext, but - 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 Facade Segment. This context also - ensures that FacadeSegments exist through underlying calls to _refresh_context(). - """ - def __init__(self, context_missing='RUNTIME_ERROR'): - super(ServerlessLambdaContext, self).__init__() - - strategy = os.getenv(CXT_MISSING_STRATEGY_KEY, context_missing) - self._context_missing = strategy - - def put_segment(self, segment): - """ - Convert the segment into a mimic segment and append it to FacadeSegment's subsegment list. - :param Segment segment: - """ - # 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 - mimic_segment = MimicSegment(parent_facade_segment, segment) - parent_facade_segment.add_subsegment(mimic_segment) - super(LambdaContext, self).put_segment(mimic_segment) - - def end_segment(self, end_time=None): - """ - Close the MimicSegment - """ - # Close the last mimic segment opened then remove it from our facade segment. - mimic_segment = self.get_trace_entity() - super(LambdaContext, self).end_segment(end_time) - if type(mimic_segment) == MimicSegment: - # The facade segment can only hold mimic segments. - facade_segment = self.__get_facade_entity() - facade_segment.remove_subsegment(mimic_segment) - - def put_subsegment(self, subsegment): - """ - Appends the subsegment as a subsegment of either the mimic segment or - another subsegment if they are the last opened entity. - :param subsegment: The subsegment to to be added as a subsegment. - """ - super(LambdaContext, self).put_subsegment(subsegment) - - def end_subsegment(self, end_time=None): - """ - End the current subsegment. In our case, subsegments - will either be a subsegment of a mimic segment or another - subsegment. - :param int end_time: epoch in seconds. If not specified the current - system time will be used. - :return: True on success, false if no parent mimic segment/subsegment is found. - """ - return super(LambdaContext, self).end_subsegment(end_time) - - def __get_facade_entity(self): - """ - Retrieves the Facade segment from thread local. This facade segment should always be present - because it was generated by the Lambda Container. - :return: FacadeSegment - """ - self._refresh_context() - facade_segment = self._local.segment # type: FacadeSegment - return facade_segment - - def get_trace_entity(self): - """ - Return the latest entity added. In this case, it'll either be a Mimic Segment or - a subsegment. Facade Segments are never returned. - If no mimic segments or subsegments were ever passed in, throw the default - context missing error. - :return: Entity - """ - # Call to Context.get_trace_entity() returns the latest mimic segment/subsegment if they exist. - # Otherwise, returns None through the following way: - # No mimic segment/subsegment exists so Context calls LambdaContext's handle_context_missing(). - # By default, Lambda's method returns no-op, so it will return None to ServerlessLambdaContext. - # Take that None as an indication to return the rightful handle_context_missing(), otherwise - # return the entity. - entity = super(LambdaContext, self).get_trace_entity() - if entity is None: - return super(LambdaContext, self).handle_context_missing() - else: - return entity - - def set_trace_entity(self, trace_entity): - """ - 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. - parent_facade_segment = self.__get_facade_entity() # type: FacadeSegment - converted_segment = MimicSegment(parent_facade_segment, trace_entity) - mimic_segment = converted_segment - else: - # Should be a Mimic Segment. If it's a subsegment, grandparent Context's - # behavior would be invoked. - mimic_segment = trace_entity - - super(LambdaContext, self).set_trace_entity(mimic_segment) - self.__get_facade_entity().subsegments = [mimic_segment] - - def _is_subsegment(self, entity): - return super(ServerlessLambdaContext, self)._is_subsegment(entity) and type(entity) != MimicSegment - - @property - def context_missing(self): - return self._context_missing - - @context_missing.setter - def context_missing(self, value): - self._context_missing = value diff --git a/aws_xray_sdk/ext/django/middleware.py b/aws_xray_sdk/ext/django/middleware.py index 07b732d9..9f071de9 100644 --- a/aws_xray_sdk/ext/django/middleware.py +++ b/aws_xray_sdk/ext/django/middleware.py @@ -1,9 +1,7 @@ import logging from aws_xray_sdk.core import xray_recorder -from aws_xray_sdk.core.lambda_launcher import check_in_lambda, LambdaContext from aws_xray_sdk.core.models import http -from aws_xray_sdk.core.serverless_lambda_context import ServerlessLambdaContext from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header @@ -27,13 +25,6 @@ def __init__(self, get_response): self.get_response = get_response - # The case when the middleware is initialized in a Lambda Context, we make sure - # to use the ServerlessLambdaContext so that the middleware properly functions. - # We also check if the current context is a LambdaContext to not override customer - # provided contexts. - if check_in_lambda() is not None and type(xray_recorder.context) == LambdaContext: - xray_recorder.context = ServerlessLambdaContext() - # hooks for django version >= 1.10 def __call__(self, request): diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index 74ec0d9e..9e0b4877 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -1,9 +1,7 @@ import flask.templating from flask import request -from aws_xray_sdk.core.lambda_launcher import check_in_lambda, LambdaContext from aws_xray_sdk.core.models import http -from aws_xray_sdk.core.serverless_lambda_context import ServerlessLambdaContext from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header @@ -20,13 +18,6 @@ def __init__(self, app, recorder): self.app.after_request(self._after_request) self.app.teardown_request(self._handle_exception) - # The case when the middleware is initialized in a Lambda Context, we make sure - # to use the ServerlessLambdaContext so that the middleware properly functions. - # We also check if the current context is a LambdaContext to not override customer - # provided contexts. - if check_in_lambda() is not None and type(self._recorder.context) == LambdaContext: - self._recorder.context = ServerlessLambdaContext() - _patch_render(recorder) def _before_request(self): diff --git a/tests/test_mimic_segment.py b/tests/test_mimic_segment.py deleted file mode 100644 index f235a669..00000000 --- a/tests/test_mimic_segment.py +++ /dev/null @@ -1,127 +0,0 @@ -import pytest - -from aws_xray_sdk.core.models.facade_segment import FacadeSegment -from aws_xray_sdk.core.models.segment import Segment -from aws_xray_sdk.core.models.subsegment import Subsegment -from aws_xray_sdk.core.models.mimic_segment import MimicSegment -from aws_xray_sdk.core.exceptions.exceptions import MimicSegmentInvalidException - - -original_segment = Segment("RealSegment") -facade_segment = FacadeSegment("FacadeSegment", "entityid", "traceid", True) - - -@pytest.fixture(autouse=True) -def cleanup_ctx(): - global original_segment, facade_segment - original_segment = Segment("RealSegment") - facade_segment = FacadeSegment("FacadeSegment", "entityid", "traceid", True) - yield - original_segment = Segment("RealSegment") - facade_segment = FacadeSegment("FacadeSegment", "entityid", "traceid", True) - - -def test_ready(): - mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) - mimic_segment.in_progress = False - assert mimic_segment.ready_to_send() - - -def test_invalid_init(): - with pytest.raises(MimicSegmentInvalidException): - MimicSegment(facade_segment=None, original_segment=original_segment) - with pytest.raises(MimicSegmentInvalidException): - MimicSegment(facade_segment=facade_segment, original_segment=None) - with pytest.raises(MimicSegmentInvalidException): - MimicSegment(facade_segment=Subsegment("Test", "local", original_segment), original_segment=None) - with pytest.raises(MimicSegmentInvalidException): - MimicSegment(facade_segment=None, original_segment=Subsegment("Test", "local", original_segment)) - with pytest.raises(MimicSegmentInvalidException): - MimicSegment(facade_segment=facade_segment, original_segment=Subsegment("Test", "local", original_segment)) - with pytest.raises(MimicSegmentInvalidException): - MimicSegment(facade_segment=original_segment, original_segment=facade_segment) - MimicSegment(facade_segment=facade_segment, original_segment=original_segment) - - -def test_init_similar(): - mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) # type: MimicSegment - - assert mimic_segment.id == original_segment.id - assert mimic_segment.name == original_segment.name - assert mimic_segment.in_progress == original_segment.in_progress - - assert mimic_segment.trace_id == facade_segment.trace_id - assert mimic_segment.parent_id == facade_segment.id - assert mimic_segment.sampled == facade_segment.sampled - - mimic_segment_serialized = mimic_segment.__getstate__() - assert mimic_segment_serialized['namespace'] == "local" - assert mimic_segment_serialized['type'] == "subsegment" - - -def test_facade_segment_properties(): - # Sampling decision is made by Facade Segment - original_segment.sampled = False - facade_segment.sampled = True - mimic_segment = MimicSegment(facade_segment=facade_segment, original_segment=original_segment) # type: MimicSegment - - assert mimic_segment.sampled == facade_segment.sampled - assert mimic_segment.sampled != original_segment.sampled - - -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) - assert getattr(mimic_segment, "ref_counter", None) - assert getattr(mimic_segment, "_subsegments_counter", None) - - assert not getattr(original_segment, "service", None) - assert not getattr(original_segment, "user", None) - assert getattr(original_segment, "ref_counter", None) - assert getattr(original_segment, "_subsegments_counter", None) - - 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") - assert original_segment.get_origin_trace_header() == mimic_segment.get_origin_trace_header() - - # No exception is thrown - test_dict = {"akey": "avalue"} - 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") - 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'] diff --git a/tests/test_serverless_lambda_context.py b/tests/test_serverless_lambda_context.py deleted file mode 100644 index 54a65b3a..00000000 --- a/tests/test_serverless_lambda_context.py +++ /dev/null @@ -1,196 +0,0 @@ -import os -import pytest - -from aws_xray_sdk.core.serverless_lambda_context import ServerlessLambdaContext -from aws_xray_sdk.core.lambda_launcher import LAMBDA_TRACE_HEADER_KEY -from aws_xray_sdk.core.exceptions.exceptions import AlreadyEndedException, SegmentNotFoundException -from aws_xray_sdk.core.models.segment import Segment -from aws_xray_sdk.core.models.subsegment import Subsegment -from aws_xray_sdk.core.models.mimic_segment import MimicSegment -from aws_xray_sdk.core.models.facade_segment import FacadeSegment - - -TRACE_ID = '1-5759e988-bd862e3fe1be46a994272793' -PARENT_ID = '53995c3f42cd8ad8' -HEADER_VAR = "Root=%s;Parent=%s;Sampled=1" % (TRACE_ID, PARENT_ID) - -os.environ[LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR -context = ServerlessLambdaContext() - -service_name = "Test Flask Server" - - -@pytest.fixture(autouse=True) -def cleanup_ctx(): - context.clear_trace_entities() - yield - context.clear_trace_entities() - - -def test_segment_generation(): - # Ensure we create Mimic Segments, and that parents of Mimic segments are Facade Segments. - segment = Segment(service_name) - context.put_segment(segment) - - mimic_segment = context.get_trace_entity() - assert type(mimic_segment) == MimicSegment - - facade_segment = getattr(context._local, 'segment', None) - assert type(facade_segment) == FacadeSegment - assert mimic_segment.parent_id == facade_segment.id - - assert facade_segment.id == PARENT_ID - assert facade_segment.trace_id == TRACE_ID - assert facade_segment.sampled - - -def test_facade_in_threadlocal(): - # Ensure that facade segments are stored in threadlocal.segment - assert not getattr(context._local, 'segment', None) - - # Refresh context to generate the facade segment. - context._refresh_context() - facade_segment = getattr(context._local, 'segment', None) - assert facade_segment - assert type(facade_segment) == FacadeSegment - assert facade_segment.id == PARENT_ID - assert facade_segment.trace_id == TRACE_ID - - -def test_put_subsegment(): - segment = Segment(service_name) - context.put_segment(segment) - - segment = context.get_trace_entity() - subsegment = Subsegment('name', 'local', segment) - context.put_subsegment(subsegment) - assert context.get_trace_entity().id == subsegment.id - - subsegment2 = Subsegment('name', 'local', segment) - context.put_subsegment(subsegment2) - assert context.get_trace_entity().id == subsegment2.id - - assert subsegment.subsegments[0] is subsegment2 - assert subsegment2.parent_id == subsegment.id - assert subsegment.parent_id == segment.id - assert subsegment2.parent_segment is segment - - assert context.get_trace_entity().id == subsegment2.id - - context.end_subsegment() - assert context.get_trace_entity().id == subsegment.id - - context.end_subsegment() - assert context.get_trace_entity().id == segment.id - assert context.get_trace_entity().in_progress - - context.end_segment() - assert context.get_trace_entity().id == segment.id - assert not context.get_trace_entity().in_progress - - -def test_remote_mimic_segment(): - # Ensure that the mimic-generated segment is set as a subsegment - # when being serialized through jsonpickle (which uses .__getstate__()) - segment = Segment(service_name) - context.put_segment(segment) - - mimic_segment = context.get_trace_entity() - - segment_properties = mimic_segment.__getstate__() - assert segment_properties['namespace'] == "local" - assert segment_properties['type'] == "subsegment" - - -def test_segment_methods_on_mimic_segment(): - # Ensure that segment operations made on the mimic segment all works. - comparison_segment = Segment(service_name) # type: Segment - context.put_segment(comparison_segment) - mimic_segment = context.get_trace_entity() # type: MimicSegment - - trace_header = "Someheader" - comparison_segment.save_origin_trace_header(trace_header) - mimic_segment.save_origin_trace_header(trace_header) - assert mimic_segment.get_origin_trace_header() == comparison_segment.get_origin_trace_header() - - assert mimic_segment.get_total_subsegments_size() == comparison_segment.get_total_subsegments_size() - comparison_segment.increment() - mimic_segment.increment() - assert mimic_segment.get_total_subsegments_size() == comparison_segment.get_total_subsegments_size() - - comparison_segment.decrement_subsegments_size() - assert mimic_segment.get_total_subsegments_size() != comparison_segment.get_total_subsegments_size() - mimic_segment.decrement_subsegments_size() - assert mimic_segment.get_total_subsegments_size() == comparison_segment.get_total_subsegments_size() - - assert mimic_segment.ready_to_send() == comparison_segment.ready_to_send() - - -def test_empty_context(): - # Test to make sure an touched context is absolutely clear. - # Call to get_trace_entity should produce a facade segment. - assert not context._local.__dict__ - - # Empty context should throw an exception - with pytest.raises(SegmentNotFoundException): - context.get_trace_entity() - - # Induce the creation of a facade segment by putting in a segment. - assert len(context._local.__dict__) == 0 - segment = Segment(service_name) - context.put_segment(segment) - - assert type(context._local.segment) == FacadeSegment - assert len(context._local.entities) == 1 - - -def test_set_trace_entity(): - segment_one = Segment(service_name) - context.put_segment(segment_one) - first_mimic_segment = context.get_trace_entity() - facade_segment = getattr(context._local, 'segment', None) - - segment_two = Segment("WOOH") - context.set_trace_entity(segment_two) - second_mimic_segment = context.get_trace_entity() - - assert first_mimic_segment.id == segment_one.id - assert first_mimic_segment.name == segment_one.name - assert first_mimic_segment.parent_id == facade_segment.id - assert first_mimic_segment.trace_id == facade_segment.trace_id - assert second_mimic_segment.id == segment_two.id - assert second_mimic_segment.name == segment_two.name - assert second_mimic_segment.parent_id == facade_segment.id - assert second_mimic_segment.trace_id == facade_segment.trace_id - - -def test_segment_close_subsegment_open(): - # Tests to make sure that when the parent, mimic segment is closed, - # and the last entity is a subsegment, the segment itself closes. - segment = Segment(service_name) - context.put_segment(segment) - mimic_segment = context.get_trace_entity() - assert mimic_segment.in_progress - subsegment = Subsegment("test", "local", mimic_segment) - context.put_subsegment(subsegment) - context.end_segment() - assert not mimic_segment.in_progress - - -def test_begin_close_twice(): - segment_one = Segment(service_name) - context.put_segment(segment_one) - context.end_segment() - entity_one = context.get_trace_entity() - context.put_segment(segment_one) - context.end_segment() - entity_two = context.get_trace_entity() - assert entity_one != entity_two - - -def test_cant_end_segment_twice(): - segment_one = Segment(service_name) - context.put_segment(segment_one) - context.end_segment() - with pytest.raises(AlreadyEndedException): - context.end_segment() From 4b51783c684edf42d0de4670d5a0f7357b396d01 Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Fri, 22 Feb 2019 15:16:23 -0800 Subject: [PATCH 5/6] Add Serverless Framework Support Revision 4 (#127) * Revised design; no more serverless context. * Subsegment generated when the middleware adds request * Unit tests to ensure that these subsegments are added on top of the facade segment --- aws_xray_sdk/core/models/entity.py | 17 ++++++++++ aws_xray_sdk/core/models/segment.py | 15 --------- aws_xray_sdk/ext/django/middleware.py | 25 +++++++++++---- aws_xray_sdk/ext/flask/middleware.py | 40 +++++++++++++++++------ tests/ext/django/test_middleware.py | 25 +++++++++++++-- tests/ext/flask/test_flask.py | 46 ++++++++++++++++++++++++++- 6 files changed, 133 insertions(+), 35 deletions(-) diff --git a/aws_xray_sdk/core/models/entity.py b/aws_xray_sdk/core/models/entity.py index 648881cc..a3593f11 100644 --- a/aws_xray_sdk/core/models/entity.py +++ b/aws_xray_sdk/core/models/entity.py @@ -18,6 +18,8 @@ _common_invalid_name_characters = '?;*()!$~^<>' _valid_annotation_key_characters = string.ascii_letters + string.digits + '_' +ORIGIN_TRACE_HEADER_ATTR_KEY = '_origin_trace_header' + class Entity(object): """ @@ -228,6 +230,20 @@ def add_exception(self, exception, stack, remote=False): self.cause['exceptions'] = exceptions self.cause['working_directory'] = os.getcwd() + def save_origin_trace_header(self, trace_header): + """ + Temporarily store additional data fields in trace header + to the segment for later propagation. The data will be + cleaned up upon serilaization. + """ + setattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, trace_header) + + def get_origin_trace_header(self): + """ + Retrieve saved trace header data. + """ + return getattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, None) + def serialize(self): """ Serialize to JSON document that can be accepted by the @@ -258,6 +274,7 @@ def _delete_empty_properties(self, properties): del properties['annotations'] if not self.metadata: del properties['metadata'] + properties.pop(ORIGIN_TRACE_HEADER_ATTR_KEY, None) del properties['sampled'] diff --git a/aws_xray_sdk/core/models/segment.py b/aws_xray_sdk/core/models/segment.py index 3c52fc3c..b07d952c 100644 --- a/aws_xray_sdk/core/models/segment.py +++ b/aws_xray_sdk/core/models/segment.py @@ -155,20 +155,6 @@ def set_rule_name(self, rule_name): self.aws['xray'] = {} self.aws['xray']['sampling_rule_name'] = rule_name - def save_origin_trace_header(self, trace_header): - """ - Temporarily store additional data fields in trace header - to the segment for later propagation. The data will be - cleaned up upon serilaization. - """ - setattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, trace_header) - - def get_origin_trace_header(self): - """ - Retrieve saved trace header data. - """ - return getattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, None) - def __getstate__(self): """ Used by jsonpikle to remove unwanted fields. @@ -179,5 +165,4 @@ def __getstate__(self): del properties['user'] del properties['ref_counter'] del properties['_subsegments_counter'] - properties.pop(ORIGIN_TRACE_HEADER_ATTR_KEY, None) return properties diff --git a/aws_xray_sdk/ext/django/middleware.py b/aws_xray_sdk/ext/django/middleware.py index 9f071de9..c2905d30 100644 --- a/aws_xray_sdk/ext/django/middleware.py +++ b/aws_xray_sdk/ext/django/middleware.py @@ -5,6 +5,7 @@ from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header +from aws_xray_sdk.core.lambda_launcher import check_in_lambda log = logging.getLogger(__name__) @@ -24,6 +25,10 @@ class XRayMiddleware(object): def __init__(self, get_response): self.get_response = get_response + self.in_lambda = False + + if check_in_lambda(): + self.in_lambda = True # hooks for django version >= 1.10 def __call__(self, request): @@ -46,12 +51,15 @@ def __call__(self, request): sampling_req=sampling_req, ) - segment = xray_recorder.begin_segment( - name=name, - traceid=xray_header.root, - parent_id=xray_header.parent, - sampling=sampling_decision, - ) + if self.in_lambda: + segment = xray_recorder.begin_subsegment(name) + else: + segment = xray_recorder.begin_segment( + name=name, + traceid=xray_header.root, + parent_id=xray_header.parent, + sampling=sampling_decision, + ) segment.save_origin_trace_header(xray_header) segment.put_http_meta(http.URL, request.build_absolute_uri()) @@ -75,7 +83,10 @@ def __call__(self, request): segment.put_http_meta(http.CONTENT_LENGTH, length) response[http.XRAY_HEADER] = prepare_response_header(xray_header, segment) - xray_recorder.end_segment() + if self.in_lambda: + xray_recorder.end_subsegment() + else: + xray_recorder.end_segment() return response diff --git a/aws_xray_sdk/ext/flask/middleware.py b/aws_xray_sdk/ext/flask/middleware.py index 9e0b4877..1a09913b 100644 --- a/aws_xray_sdk/ext/flask/middleware.py +++ b/aws_xray_sdk/ext/flask/middleware.py @@ -5,6 +5,7 @@ from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import calculate_sampling_decision, \ calculate_segment_name, construct_xray_header, prepare_response_header +from aws_xray_sdk.core.lambda_launcher import check_in_lambda class XRayMiddleware(object): @@ -17,6 +18,10 @@ def __init__(self, app, recorder): self.app.before_request(self._before_request) self.app.after_request(self._after_request) self.app.teardown_request(self._handle_exception) + self.in_lambda = False + + if check_in_lambda(): + self.in_lambda = True _patch_render(recorder) @@ -39,12 +44,15 @@ def _before_request(self): sampling_req=sampling_req, ) - segment = self._recorder.begin_segment( - name=name, - traceid=xray_header.root, - parent_id=xray_header.parent, - sampling=sampling_decision, - ) + if self.in_lambda: + segment = self._recorder.begin_subsegment(name) + else: + segment = self._recorder.begin_segment( + name=name, + traceid=xray_header.root, + parent_id=xray_header.parent, + sampling=sampling_decision, + ) segment.save_origin_trace_header(xray_header) segment.put_http_meta(http.URL, req.base_url) @@ -59,7 +67,10 @@ def _before_request(self): segment.put_http_meta(http.CLIENT_IP, req.remote_addr) def _after_request(self, response): - segment = self._recorder.current_segment() + if self.in_lambda: + segment = self._recorder.current_subsegment() + else: + segment = self._recorder.current_segment() segment.put_http_meta(http.STATUS, response.status_code) origin_header = segment.get_origin_trace_header() @@ -70,7 +81,10 @@ def _after_request(self, response): if cont_len: segment.put_http_meta(http.CONTENT_LENGTH, int(cont_len)) - self._recorder.end_segment() + if self.in_lambda: + self._recorder.end_subsegment() + else: + self._recorder.end_segment() return response def _handle_exception(self, exception): @@ -78,7 +92,10 @@ def _handle_exception(self, exception): return segment = None try: - segment = self._recorder.current_segment() + if self.in_lambda: + segment = self._recorder.current_subsegment() + else: + segment = self._recorder.current_segment() except Exception: pass if not segment: @@ -87,7 +104,10 @@ def _handle_exception(self, exception): segment.put_http_meta(http.STATUS, 500) stack = stacktrace.get_stacktrace(limit=self._recorder._max_trace_back) segment.add_exception(exception, stack) - self._recorder.end_segment() + if self.in_lambda: + self._recorder.end_subsegment() + else: + self._recorder.end_segment() def _patch_render(recorder): diff --git a/tests/ext/django/test_middleware.py b/tests/ext/django/test_middleware.py index a0128b7c..cb36ddf9 100644 --- a/tests/ext/django/test_middleware.py +++ b/tests/ext/django/test_middleware.py @@ -3,9 +3,11 @@ from django.core.urlresolvers import reverse from django.test import TestCase -from aws_xray_sdk.core import xray_recorder +from aws_xray_sdk.core import xray_recorder, lambda_launcher from aws_xray_sdk.core.context import Context -from aws_xray_sdk.core.models import http +from aws_xray_sdk.core.models import http, facade_segment +from tests.util import get_new_stubbed_recorder +import os class XRayTestCase(TestCase): @@ -111,3 +113,22 @@ def test_disabled_sdk(self): self.client.get(url) segment = xray_recorder.emitter.pop() assert not segment + + def test_lambda_serverless(self): + TRACE_ID = '1-5759e988-bd862e3fe1be46a994272793' + PARENT_ID = '53995c3f42cd8ad8' + HEADER_VAR = "Root=%s;Parent=%s;Sampled=1" % (TRACE_ID, PARENT_ID) + + os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR + lambda_context = lambda_launcher.LambdaContext() + + new_recorder = get_new_stubbed_recorder() + new_recorder.configure(service='test', sampling=False, context=lambda_context) + subsegment = new_recorder.begin_subsegment("subsegment") + assert type(subsegment.parent_segment) == facade_segment.FacadeSegment + new_recorder.end_subsegment() + + url = reverse('200ok') + self.client.get(url) + segment = new_recorder.emitter.pop() + assert not segment diff --git a/tests/ext/flask/test_flask.py b/tests/ext/flask/test_flask.py index 07c8d42c..b5bbf1f7 100644 --- a/tests/ext/flask/test_flask.py +++ b/tests/ext/flask/test_flask.py @@ -4,8 +4,10 @@ from aws_xray_sdk import global_sdk_config from aws_xray_sdk.ext.flask.middleware import XRayMiddleware from aws_xray_sdk.core.context import Context -from aws_xray_sdk.core.models import http +from aws_xray_sdk.core import lambda_launcher +from aws_xray_sdk.core.models import http, facade_segment from tests.util import get_new_stubbed_recorder +import os # define a flask app for testing purpose @@ -153,3 +155,45 @@ def test_disabled_sdk(): app.get(path) segment = recorder.emitter.pop() assert not segment + + +def test_lambda_serverless(): + TRACE_ID = '1-5759e988-bd862e3fe1be46a994272793' + PARENT_ID = '53995c3f42cd8ad8' + HEADER_VAR = "Root=%s;Parent=%s;Sampled=1" % (TRACE_ID, PARENT_ID) + + os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR + lambda_context = lambda_launcher.LambdaContext() + + new_recorder = get_new_stubbed_recorder() + new_recorder.configure(service='test', sampling=False, context=lambda_context) + new_app = Flask(__name__) + + @new_app.route('/subsegment') + def subsegment(): + # Test in between request and make sure Serverless creates a subsegment instead of a segment. + # Ensure that the parent segment is a facade segment. + assert new_recorder.current_subsegment() + assert type(new_recorder.current_segment()) == facade_segment.FacadeSegment + return 'ok' + + @new_app.route('/trace_header') + def trace_header(): + # Ensure trace header is preserved. + subsegment = new_recorder.current_subsegment() + header = subsegment.get_origin_trace_header() + assert header.data['k1'] == 'v1' + return 'ok' + + middleware = XRayMiddleware(new_app, new_recorder) + middleware.in_lambda = True + + app_client = new_app.test_client() + + path = '/subsegment' + app_client.get(path) + segment = recorder.emitter.pop() + assert not segment # Segment should be none because it's created and ended by the middleware + + path2 = '/trace_header' + app_client.get(path2, headers={http.XRAY_HEADER: 'k1=v1'}) From 293f76ff91c6ee33bed3cf234f32ff32c8a98e0a Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Mon, 25 Feb 2019 12:39:54 -0800 Subject: [PATCH 6/6] Add Serverless Framework Support Revision 5 (#127) * Comment change added for save_origin_trace_header --- aws_xray_sdk/core/models/entity.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aws_xray_sdk/core/models/entity.py b/aws_xray_sdk/core/models/entity.py index a3593f11..8b8bdf4f 100644 --- a/aws_xray_sdk/core/models/entity.py +++ b/aws_xray_sdk/core/models/entity.py @@ -233,8 +233,8 @@ def add_exception(self, exception, stack, remote=False): def save_origin_trace_header(self, trace_header): """ Temporarily store additional data fields in trace header - to the segment for later propagation. The data will be - cleaned up upon serilaization. + to the entity for later propagation. The data will be + cleaned up upon serialization. """ setattr(self, ORIGIN_TRACE_HEADER_ATTR_KEY, trace_header)