forked from DataDog/datadog-lambda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathwrapper.py
More file actions
371 lines (331 loc) · 14.3 KB
/
wrapper.py
File metadata and controls
371 lines (331 loc) · 14.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
# Unless explicitly stated otherwise all files in this repository are licensed
# under the Apache License Version 2.0.
# This product includes software developed at Datadog (https://www.datadoghq.com/).
# Copyright 2019 Datadog, Inc.
import base64
import os
import logging
import traceback
from importlib import import_module
import json
from time import time_ns
from datadog_lambda.extension import should_use_extension, flush_extension
from datadog_lambda.cold_start import (
set_cold_start,
is_cold_start,
is_proactive_init,
is_new_sandbox,
ColdStartTracer,
)
from datadog_lambda.constants import (
TraceContextSource,
XraySubsegment,
Headers,
TraceHeader,
)
from datadog_lambda.metric import (
flush_stats,
submit_invocations_metric,
submit_errors_metric,
)
from datadog_lambda.module_name import modify_module_name
from datadog_lambda.patch import patch_all
from datadog_lambda.tracing import (
extract_dd_trace_context,
create_dd_dummy_metadata_subsegment,
inject_correlation_ids,
dd_tracing_enabled,
mark_trace_as_error_for_5xx_responses,
set_correlation_ids,
set_dd_trace_py_root,
create_function_execution_span,
create_inferred_span,
InferredSpanInfo,
is_authorizer_response,
tracer,
)
from datadog_lambda.trigger import (
extract_trigger_tags,
extract_http_status_code_tag,
)
from datadog_lambda.tag_object import tag_object
profiling_env_var = os.environ.get("DD_PROFILING_ENABLED", "false").lower() == "true"
if profiling_env_var:
from ddtrace.profiling import profiler
logger = logging.getLogger(__name__)
dd_capture_lambda_payload_enabled = (
os.environ.get("DD_CAPTURE_LAMBDA_PAYLOAD", "false").lower() == "true"
)
DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG"
DD_LOGS_INJECTION = "DD_LOGS_INJECTION"
DD_MERGE_XRAY_TRACES = "DD_MERGE_XRAY_TRACES"
AWS_LAMBDA_FUNCTION_NAME = "AWS_LAMBDA_FUNCTION_NAME"
DD_TRACE_EXTRACTOR = "DD_TRACE_EXTRACTOR"
DD_TRACE_MANAGED_SERVICES = "DD_TRACE_MANAGED_SERVICES"
DD_ENCODE_AUTHORIZER_CONTEXT = "DD_ENCODE_AUTHORIZER_CONTEXT"
DD_DECODE_AUTHORIZER_CONTEXT = "DD_DECODE_AUTHORIZER_CONTEXT"
DD_COLD_START_TRACING = "DD_COLD_START_TRACING"
DD_MIN_COLD_START_DURATION = "DD_MIN_COLD_START_DURATION"
DD_COLD_START_TRACE_SKIP_LIB = "DD_COLD_START_TRACE_SKIP_LIB"
DD_REQUESTS_SERVICE_NAME = "DD_REQUESTS_SERVICE_NAME"
DD_SERVICE = "DD_SERVICE"
DD_ENV = "DD_ENV"
env_env_var = os.environ.get(DD_ENV, None)
init_timestamp_ns = time_ns()
"""
Usage:
import requests
from datadog_lambda.wrapper import datadog_lambda_wrapper
from datadog_lambda.metric import lambda_metric
@datadog_lambda_wrapper
def my_lambda_handle(event, context):
lambda_metric("my_metric", 10)
requests.get("https://www.datadoghq.com")
"""
class _NoopDecorator(object):
def __init__(self, func):
self.func = func
def __call__(self, *args, **kwargs):
return self.func(*args, **kwargs)
class _LambdaDecorator(object):
"""
Decorator to automatically initialize Datadog API client, flush metrics,
and extracts/injects trace context.
"""
_force_wrap = False
def __new__(cls, func):
"""
If the decorator is accidentally applied to the same function multiple times,
wrap only once.
If _force_wrap, always return a real decorator, useful for unit tests.
"""
try:
if cls._force_wrap or not isinstance(func, _LambdaDecorator):
wrapped = super(_LambdaDecorator, cls).__new__(cls)
logger.debug("datadog_lambda_wrapper wrapped")
return wrapped
else:
logger.debug("datadog_lambda_wrapper already wrapped")
return _NoopDecorator(func)
except Exception as e:
logger.error(format_err_with_traceback(e))
return func
def __init__(self, func):
"""Executes when the wrapped function gets wrapped"""
try:
self.func = func
self.flush_to_log = os.environ.get(DD_FLUSH_TO_LOG, "").lower() == "true"
self.logs_injection = (
os.environ.get(DD_LOGS_INJECTION, "true").lower() == "true"
)
self.merge_xray_traces = (
os.environ.get(DD_MERGE_XRAY_TRACES, "false").lower() == "true"
)
self.function_name = os.environ.get(AWS_LAMBDA_FUNCTION_NAME, "function")
self.service = os.environ.get(DD_SERVICE, None)
self.extractor_env = os.environ.get(DD_TRACE_EXTRACTOR, None)
self.trace_extractor = None
self.span = None
self.inferred_span = None
depends_on_dd_tracing_enabled = (
lambda original_boolean: dd_tracing_enabled and original_boolean
)
self.make_inferred_span = depends_on_dd_tracing_enabled(
os.environ.get(DD_TRACE_MANAGED_SERVICES, "true").lower() == "true"
)
self.encode_authorizer_context = depends_on_dd_tracing_enabled(
os.environ.get(DD_ENCODE_AUTHORIZER_CONTEXT, "true").lower() == "true"
)
self.decode_authorizer_context = depends_on_dd_tracing_enabled(
os.environ.get(DD_DECODE_AUTHORIZER_CONTEXT, "true").lower() == "true"
)
self.cold_start_tracing = depends_on_dd_tracing_enabled(
os.environ.get(DD_COLD_START_TRACING, "true").lower() == "true"
)
self.min_cold_start_trace_duration = 3
if DD_MIN_COLD_START_DURATION in os.environ:
try:
self.min_cold_start_trace_duration = int(
os.environ[DD_MIN_COLD_START_DURATION]
)
except Exception:
logger.debug(f"Malformatted env {DD_MIN_COLD_START_DURATION}")
self.cold_start_trace_skip_lib = [
"ddtrace.internal.compat",
"ddtrace.filters",
]
if DD_COLD_START_TRACE_SKIP_LIB in os.environ:
try:
self.cold_start_trace_skip_lib = os.environ[
DD_COLD_START_TRACE_SKIP_LIB
].split(",")
except Exception:
logger.debug(f"Malformatted for env {DD_COLD_START_TRACE_SKIP_LIB}")
self.response = None
if profiling_env_var:
self.prof = profiler.Profiler(env=env_env_var, service=self.service)
if self.extractor_env:
extractor_parts = self.extractor_env.rsplit(".", 1)
if len(extractor_parts) == 2:
(mod_name, extractor_name) = extractor_parts
modified_extractor_name = modify_module_name(mod_name)
extractor_module = import_module(modified_extractor_name)
self.trace_extractor = getattr(extractor_module, extractor_name)
# Inject trace correlation ids to logs
if self.logs_injection:
inject_correlation_ids()
# This prevents a breaking change in ddtrace v0.49 regarding the service name
# in requests-related spans
os.environ[DD_REQUESTS_SERVICE_NAME] = os.environ.get(
DD_SERVICE, "aws.lambda"
)
# Patch third-party libraries for tracing
patch_all()
logger.debug("datadog_lambda_wrapper initialized")
except Exception as e:
logger.error(format_err_with_traceback(e))
def __call__(self, event, context, **kwargs):
"""Executes when the wrapped function gets called"""
self._before(event, context)
try:
self.response = self.func(event, context, **kwargs)
return self.response
except Exception:
submit_errors_metric(context)
if self.span:
self.span.set_traceback()
raise
finally:
self._after(event, context)
def _inject_authorizer_span_headers(self, request_id):
reference_span = self.inferred_span if self.inferred_span else self.span
assert reference_span.finished
# the finish_time_ns should be set as the end of the inferred span if it exist
# or the end of the current span
finish_time_ns = (
reference_span.start_ns + reference_span.duration_ns
if reference_span is not None
and hasattr(reference_span, "start_ns")
and hasattr(reference_span, "duration_ns")
else time_ns()
)
injected_headers = {}
source_span = self.inferred_span if self.inferred_span else self.span
span_context = source_span.context
injected_headers[TraceHeader.TRACE_ID] = str(span_context.trace_id)
injected_headers[TraceHeader.PARENT_ID] = str(span_context.span_id)
sampling_priority = span_context.sampling_priority
if sampling_priority is not None:
injected_headers[TraceHeader.SAMPLING_PRIORITY] = str(
span_context.sampling_priority
)
injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns
if request_id is not None:
injected_headers[Headers.Authorizing_Request_Id] = request_id
datadog_data = base64.b64encode(json.dumps(injected_headers).encode()).decode()
self.response.setdefault("context", {})
self.response["context"]["_datadog"] = datadog_data
def _before(self, event, context):
try:
self.response = None
set_cold_start(init_timestamp_ns)
submit_invocations_metric(context)
self.trigger_tags = extract_trigger_tags(event, context)
# Extract Datadog trace context and source from incoming requests
dd_context, trace_context_source, event_source = extract_dd_trace_context(
event,
context,
extractor=self.trace_extractor,
decode_authorizer_context=self.decode_authorizer_context,
)
self.event_source = event_source
# Create a Datadog X-Ray subsegment with the trace context
if dd_context and trace_context_source == TraceContextSource.EVENT:
create_dd_dummy_metadata_subsegment(
dd_context, XraySubsegment.TRACE_KEY
)
if dd_tracing_enabled:
set_dd_trace_py_root(trace_context_source, self.merge_xray_traces)
if self.make_inferred_span:
self.inferred_span = create_inferred_span(
event, context, event_source, self.decode_authorizer_context
)
self.span = create_function_execution_span(
context,
self.function_name,
is_cold_start(),
is_proactive_init(),
trace_context_source,
self.merge_xray_traces,
self.trigger_tags,
parent_span=self.inferred_span,
)
else:
set_correlation_ids()
if profiling_env_var and is_new_sandbox():
self.prof.start(stop_on_exit=False, profile_children=True)
logger.debug("datadog_lambda_wrapper _before() done")
except Exception as e:
logger.error(format_err_with_traceback(e))
def _after(self, event, context):
try:
status_code = extract_http_status_code_tag(self.trigger_tags, self.response)
if status_code:
self.trigger_tags["http.status_code"] = status_code
mark_trace_as_error_for_5xx_responses(context, status_code, self.span)
# Create a new dummy Datadog subsegment for function trigger tags so we
# can attach them to X-Ray spans when hybrid tracing is used
if self.trigger_tags:
create_dd_dummy_metadata_subsegment(
self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
)
should_trace_cold_start = (
dd_tracing_enabled and self.cold_start_tracing and is_new_sandbox()
)
if should_trace_cold_start:
trace_ctx = tracer.current_trace_context()
if self.span:
if dd_capture_lambda_payload_enabled:
tag_object(self.span, "function.request", event)
tag_object(self.span, "function.response", self.response)
if status_code:
self.span.set_tag("http.status_code", status_code)
self.span.finish()
if self.inferred_span:
if status_code:
self.inferred_span.set_tag("http.status_code", status_code)
if self.service:
self.inferred_span.set_tag("peer.service", self.service)
if InferredSpanInfo.is_async(self.inferred_span) and self.span:
self.inferred_span.finish(finish_time=self.span.start)
else:
self.inferred_span.finish()
if should_trace_cold_start:
try:
following_span = self.span or self.inferred_span
ColdStartTracer(
tracer,
self.function_name,
following_span.start_ns,
trace_ctx,
self.min_cold_start_trace_duration,
self.cold_start_trace_skip_lib,
).trace()
except Exception as e:
logger.debug("Failed to create cold start spans. %s", e)
if not self.flush_to_log or should_use_extension:
flush_stats()
if should_use_extension:
flush_extension()
if self.encode_authorizer_context and is_authorizer_response(self.response):
self._inject_authorizer_span_headers(
event.get("requestContext", {}).get("requestId")
)
logger.debug("datadog_lambda_wrapper _after() done")
except Exception as e:
logger.error(format_err_with_traceback(e))
def format_err_with_traceback(e):
return "Error {}. Traceback: {}".format(
e, traceback.format_exc().replace("\n", "\r")
)
datadog_lambda_wrapper = _LambdaDecorator