33import threading
44
55from aws_xray_sdk import global_sdk_config
6+ from .models .dummy_entities import DummySegment
67from .models .facade_segment import FacadeSegment
78from .models .trace_header import TraceHeader
89from .context import Context
@@ -44,7 +45,7 @@ class LambdaContext(Context):
4445 """
4546 Lambda service will generate a segment for each function invocation which
4647 cannot be mutated. The context doesn't keep any manually created segment
47- but instead every time ``get_trace_entity()`` gets called it refresh the facade
48+ but instead every time ``get_trace_entity()`` gets called it refresh the
4849 segment based on environment variables set by Lambda worker.
4950 """
5051 def __init__ (self ):
@@ -65,19 +66,31 @@ def end_segment(self, end_time=None):
6566
6667 def put_subsegment (self , subsegment ):
6768 """
68- Refresh the facade segment every time this function is invoked to prevent
69+ Refresh the segment every time this function is invoked to prevent
6970 a new subsegment from being attached to a leaked segment/subsegment.
7071 """
7172 current_entity = self .get_trace_entity ()
7273
73- if not self ._is_subsegment (current_entity ) and current_entity . initializing :
74- if global_sdk_config .sdk_enabled ():
74+ if not self ._is_subsegment (current_entity ) and ( getattr ( current_entity , ' initializing' , None ) or isinstance ( current_entity , DummySegment )) :
75+ if global_sdk_config .sdk_enabled () and not os . getenv ( LAMBDA_TRACE_HEADER_KEY ) :
7576 log .warning ("Subsegment %s discarded due to Lambda worker still initializing" % subsegment .name )
7677 return
7778
7879 current_entity .add_subsegment (subsegment )
7980 self ._local .entities .append (subsegment )
8081
82+ def set_trace_entity (self , trace_entity ):
83+ """
84+ For Lambda context, we additionally store the segment in the thread local.
85+ """
86+ if self ._is_subsegment (trace_entity ):
87+ segment = trace_entity .parent_segment
88+ else :
89+ segment = trace_entity
90+
91+ setattr (self ._local , 'segment' , segment )
92+ setattr (self ._local , 'entities' , [trace_entity ])
93+
8194 def get_trace_entity (self ):
8295 self ._refresh_context ()
8396 if getattr (self ._local , 'entities' , None ):
@@ -87,9 +100,9 @@ def get_trace_entity(self):
87100
88101 def _refresh_context (self ):
89102 """
90- Get current facade segment. To prevent resource leaking in Lambda worker,
103+ Get current segment. To prevent resource leaking in Lambda worker,
91104 every time there is segment present, we compare its trace id to current
92- environment variables. If it is different we create a new facade segment
105+ environment variables. If it is different we create a new segment
93106 and clean up subsegments stored.
94107 """
95108 header_str = os .getenv (LAMBDA_TRACE_HEADER_KEY )
@@ -124,8 +137,8 @@ def handle_context_missing(self):
124137
125138 def _initialize_context (self , trace_header ):
126139 """
127- Create a facade segment based on environment variables
128- set by AWS Lambda and initialize storage for subsegments.
140+ Create a segment based on environment variables set by
141+ AWS Lambda and initialize storage for subsegments.
129142 """
130143 sampled = None
131144 if not global_sdk_config .sdk_enabled ():
@@ -136,12 +149,17 @@ def _initialize_context(self, trace_header):
136149 elif trace_header .sampled == 1 :
137150 sampled = True
138151
139- segment = FacadeSegment (
140- name = 'facade' ,
141- traceid = trace_header .root ,
142- entityid = trace_header .parent ,
143- sampled = sampled ,
144- )
152+ segment = None
153+ if not trace_header .root or not trace_header .parent or trace_header .sampled is None :
154+ segment = DummySegment ()
155+ log .debug ("Creating NoOp/Dummy parent segment" )
156+ else :
157+ segment = FacadeSegment (
158+ name = 'facade' ,
159+ traceid = trace_header .root ,
160+ entityid = trace_header .parent ,
161+ sampled = sampled ,
162+ )
145163 segment .save_origin_trace_header (trace_header )
146164 setattr (self ._local , 'segment' , segment )
147165 setattr (self ._local , 'entities' , [])
0 commit comments