1+ import time
2+ import os
3+ from typing import List , Hashable
4+ import logging
5+
6+ logger = logging .getLogger (__name__ )
7+
18_cold_start = True
29_lambda_container_initialized = False
310
@@ -21,3 +28,196 @@ def is_cold_start():
2128def get_cold_start_tag ():
2229 """Returns the cold start tag to be used in metrics"""
2330 return "cold_start:{}" .format (str (is_cold_start ()).lower ())
31+
32+
33+ class ImportNode (object ):
34+ def __init__ (self , module_name , full_file_path , start_time_ns , end_time_ns = None ):
35+ self .module_name = module_name
36+ self .full_file_path = full_file_path
37+ self .start_time_ns = start_time_ns
38+ self .end_time_ns = end_time_ns
39+ self .children = []
40+
41+
42+ root_nodes : List [ImportNode ] = []
43+ import_stack : List [ImportNode ] = []
44+ already_wrapped_loaders = set ()
45+
46+
47+ def reset_node_stacks ():
48+ global root_nodes
49+ root_nodes = []
50+ global import_stack
51+ import_stack = []
52+
53+
54+ def push_node (module_name , file_path ):
55+ node = ImportNode (module_name , file_path , time .time_ns ())
56+ global import_stack
57+ if import_stack :
58+ import_stack [- 1 ].children .append (node )
59+ import_stack .append (node )
60+
61+
62+ def pop_node (module_name ):
63+ global import_stack
64+ if not import_stack :
65+ return
66+ node = import_stack .pop ()
67+ if node .module_name != module_name :
68+ return
69+ end_time_ns = time .time_ns ()
70+ node .end_time_ns = end_time_ns
71+ if not import_stack : # import_stack empty, a root node has been found
72+ global root_nodes
73+ root_nodes .append (node )
74+
75+
76+ def wrap_exec_module (original_exec_module ):
77+ def wrapped_method (module ):
78+ should_pop = False
79+ try :
80+ spec = module .__spec__
81+ push_node (spec .name , spec .origin )
82+ should_pop = True
83+ except Exception :
84+ pass
85+ try :
86+ return original_exec_module (module )
87+ finally :
88+ if should_pop :
89+ pop_node (spec .name )
90+
91+ return wrapped_method
92+
93+
94+ def wrap_find_spec (original_find_spec ):
95+ def wrapped_find_spec (* args , ** kwargs ):
96+ spec = original_find_spec (* args , ** kwargs )
97+ if spec is None :
98+ return None
99+ loader = getattr (spec , "loader" , None )
100+ if (
101+ loader is not None
102+ and isinstance (loader , Hashable )
103+ and loader not in already_wrapped_loaders
104+ ):
105+ if hasattr (loader , "exec_module" ):
106+ try :
107+ loader .exec_module = wrap_exec_module (loader .exec_module )
108+ already_wrapped_loaders .add (loader )
109+ except Exception as e :
110+ logger .debug ("Failed to wrap the loader. %s" , e )
111+ return spec
112+
113+ return wrapped_find_spec
114+
115+
116+ def initialize_cold_start_tracing ():
117+ if (
118+ is_cold_start ()
119+ and os .environ .get ("DD_TRACE_ENABLED" , "true" ).lower () == "true"
120+ and os .environ .get ("DD_COLD_START_TRACING" , "true" ).lower () == "true"
121+ ):
122+ from sys import version_info , meta_path
123+
124+ if version_info >= (3 , 7 ): # current implementation only support version > 3.7
125+ for importer in meta_path :
126+ try :
127+ importer .find_spec = wrap_find_spec (importer .find_spec )
128+ except Exception :
129+ pass
130+
131+
132+ class ColdStartTracer (object ):
133+ def __init__ (
134+ self ,
135+ tracer ,
136+ function_name ,
137+ cold_start_span_finish_time_ns ,
138+ trace_ctx ,
139+ min_duration_ms : int ,
140+ ignored_libs : List [str ] = [],
141+ ):
142+ self ._tracer = tracer
143+ self .function_name = function_name
144+ self .cold_start_span_finish_time_ns = cold_start_span_finish_time_ns
145+ self .min_duration_ms = min_duration_ms
146+ self .trace_ctx = trace_ctx
147+ self .ignored_libs = ignored_libs
148+ self .need_to_reactivate_context = True
149+
150+ def trace (self , root_nodes : List [ImportNode ] = root_nodes ):
151+ if not root_nodes :
152+ return
153+ cold_start_span_start_time_ns = root_nodes [0 ].start_time_ns
154+ cold_start_span = self .create_cold_start_span (cold_start_span_start_time_ns )
155+ while root_nodes :
156+ root_node = root_nodes .pop ()
157+ self .trace_tree (root_node , cold_start_span )
158+ self .finish_span (cold_start_span , self .cold_start_span_finish_time_ns )
159+
160+ def trace_tree (self , import_node : ImportNode , parent_span ):
161+ if (
162+ import_node .end_time_ns - import_node .start_time_ns
163+ < self .min_duration_ms * 1e6
164+ or import_node .module_name in self .ignored_libs
165+ ):
166+ return
167+
168+ span = self .start_span (
169+ "aws.lambda.import" , import_node .module_name , import_node .start_time_ns
170+ )
171+ tags = {
172+ "resource_names" : import_node .module_name ,
173+ "resource.name" : import_node .module_name ,
174+ "filename" : import_node .full_file_path ,
175+ "operation_name" : self .get_operation_name (import_node .full_file_path ),
176+ }
177+ span .set_tags (tags )
178+ if parent_span :
179+ span .parent_id = parent_span .span_id
180+ for child_node in import_node .children :
181+ self .trace_tree (child_node , span )
182+ self .finish_span (span , import_node .end_time_ns )
183+
184+ def create_cold_start_span (self , start_time_ns ):
185+ span = self .start_span ("aws.lambda.load" , self .function_name , start_time_ns )
186+ tags = {
187+ "resource_names" : self .function_name ,
188+ "resource.name" : self .function_name ,
189+ "operation_name" : "aws.lambda.load" ,
190+ }
191+ span .set_tags (tags )
192+ return span
193+
194+ def start_span (self , span_type , resource , start_time_ns ):
195+ if self .need_to_reactivate_context :
196+ self ._tracer .context_provider .activate (
197+ self .trace_ctx
198+ ) # reactivate required after each finish() call
199+ self .need_to_reactivate_context = False
200+ span_kwargs = {
201+ "service" : "aws.lambda" ,
202+ "resource" : resource ,
203+ "span_type" : span_type ,
204+ }
205+ span = self ._tracer .trace (span_type , ** span_kwargs )
206+ span .start_ns = start_time_ns
207+ return span
208+
209+ def finish_span (self , span , finish_time_ns ):
210+ span .finish (finish_time_ns / 1e9 )
211+ self .need_to_reactivate_context = True
212+
213+ def get_operation_name (self , filename : str ):
214+ if filename is None :
215+ return "aws.lambda.import_core_module"
216+ if not isinstance (filename , str ):
217+ return "aws.lambda.import"
218+ if filename .startswith ("/opt/" ):
219+ return "aws.lambda.import_layer"
220+ elif filename .startswith ("/var/lang/" ):
221+ return "aws.lambda.import_runtime"
222+ else :
223+ return "aws.lambda.import"
0 commit comments