Skip to content

Commit

Permalink
Add Serverless Framework Support (aws#2)
Browse files Browse the repository at this point in the history
    * 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.
  • Loading branch information
chanchiem committed Feb 20, 2019
1 parent 237f8c7 commit c3f333f
Show file tree
Hide file tree
Showing 8 changed files with 476 additions and 1 deletion.
4 changes: 4 additions & 0 deletions aws_xray_sdk/core/exceptions/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ class FacadeSegmentMutationException(Exception):
pass


class MimicSegmentInvalidException(Exception):
pass


class MissingPluginNames(Exception):
pass

Expand Down
35 changes: 35 additions & 0 deletions aws_xray_sdk/core/models/mimic_segment.py
Original file line number Diff line number Diff line change
@@ -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
3 changes: 2 additions & 1 deletion aws_xray_sdk/core/recorder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
131 changes: 131 additions & 0 deletions aws_xray_sdk/core/serverless_context.py
Original file line number Diff line number Diff line change
@@ -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
7 changes: 7 additions & 0 deletions aws_xray_sdk/ext/django/middleware.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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):

Expand Down
7 changes: 7 additions & 0 deletions aws_xray_sdk/ext/flask/middleware.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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):
Expand Down
93 changes: 93 additions & 0 deletions tests/test_mimic_segment.py
Original file line number Diff line number Diff line change
@@ -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")
Loading

0 comments on commit c3f333f

Please sign in to comment.