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.

Author: chanchiem

aws_xray_sdk/core/exceptions/exceptions.py

@@ -22,6 +22,10 @@ class FacadeSegmentMutationException(Exception):
     pass
 
 
+class MimicSegmentInvalidException(Exception):
+    pass
+
+
 class MissingPluginNames(Exception):
     pass
 

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

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):
         """

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

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):
 

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):

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")