From 95bd2c8f76cd11fc74553abb282bffab8f88f075 Mon Sep 17 00:00:00 2001 From: Jeffery Saeteurn Date: Thu, 7 Feb 2019 15:39:33 -0800 Subject: [PATCH] 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()