Skip to content

Commit 294b9a3

Browse files
authored
Merge branch 'main' into yiming.luo/durable-first-invocation-tag
2 parents e2108f4 + 38cabf2 commit 294b9a3

File tree

10 files changed

+164
-27
lines changed

10 files changed

+164
-27
lines changed

.github/workflows/system_tests.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ on:
77
pull_request:
88
branches:
99
- "**"
10+
schedule:
11+
- cron: '00 03 * * *'
1012

1113
jobs:
1214
build:

Dockerfile

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,12 @@ RUN pip install --no-cache-dir . -t ./python/lib/$runtime/site-packages
3333
RUN rm -rf ./python/lib/$runtime/site-packages/botocore*
3434
RUN rm -rf ./python/lib/$runtime/site-packages/setuptools
3535
RUN rm -rf ./python/lib/$runtime/site-packages/jsonschema/tests
36-
RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast/_ast/iastpatch*.so
37-
RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast/_taint_tracking/_vendor
38-
RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast/_taint_tracking/*.so
39-
RUN rm -f ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast/_stacktrace*.so
36+
RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast
37+
RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/internal/test_visibility
38+
# Dogshell
39+
RUN rm -rf ./python/lib/$runtime/site-packages/datadog/dogshell
40+
RUN rm -rf ./python/lib/$runtime/site-packages/bin/dog*
41+
4042
# remove *.dist-info directories except any entry_points.txt files and METADATA files required for Appsec Software Composition Analysis
4143
RUN find ./python/lib/$runtime/site-packages/*.dist-info \
4244
-type f \
@@ -50,7 +52,8 @@ RUN rm -rf \
5052
./python/lib/$runtime/site-packages/urllib3* \
5153
./python/lib/$runtime/site-packages/certifi* \
5254
./python/lib/$runtime/site-packages/idna* \
53-
./python/lib/$runtime/site-packages/charset_normalizer*
55+
./python/lib/$runtime/site-packages/charset_normalizer* \
56+
./python/lib/$runtime/site-packages/*__mypyc*.so # from charset_normalizer
5457

5558
# Precompile all .pyc files and remove .py files. This speeds up load time.
5659
# Compile with optimization level 2 (-OO) and PYTHONNODEBUGRANGES=1 to redtce
@@ -73,6 +76,7 @@ RUN find ./python/lib/$runtime/site-packages/ddtrace -name \*.cc -delete
7376
RUN find ./python/lib/$runtime/site-packages/ddtrace -name \*.h -delete
7477
RUN find ./python/lib/$runtime/site-packages/ddtrace -name \*.hpp -delete
7578
RUN find ./python/lib/$runtime/site-packages/ddtrace -name \*.pyx -delete
79+
RUN find ./python/lib/$runtime/site-packages/ddtrace -name \*.pyi -delete
7680

7781
# Strip debug symbols and symbols that are not needed for relocation
7882
# processing using strip --strip-unneeded for all .so files. This is to

datadog_lambda/config.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,12 @@ def _reset(self):
141141
"enabled" if config.fips_mode_enabled else "not enabled",
142142
)
143143

144+
# disable css to prevent double counting in lambda
145+
os.environ["DD_TRACE_STATS_COMPUTATION_ENABLED"] = "false"
146+
147+
# unset css aliases to ensure it is disabled
148+
if "DD_TRACE_COMPUTE_STATS" in os.environ:
149+
del os.environ["DD_TRACE_COMPUTE_STATS"]
144150

145151
if (
146152
"DD_INSTRUMENTATION_TELEMETRY_ENABLED" not in os.environ

datadog_lambda/durable.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,3 +53,17 @@ def extract_durable_function_tags(event):
5353
is_first_invocation
5454
).lower(),
5555
}
56+
57+
58+
VALID_DURABLE_STATUSES = {"SUCCEEDED", "FAILED", "STOPPED", "TIMED_OUT"}
59+
60+
61+
def extract_durable_execution_status(response, event):
62+
if not isinstance(event, dict) or "DurableExecutionArn" not in event:
63+
return None
64+
if not isinstance(response, dict):
65+
return None
66+
status = response.get("Status")
67+
if status not in VALID_DURABLE_STATUSES:
68+
return None
69+
return status

datadog_lambda/wrapper.py

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,10 @@
4242
tracer,
4343
propagator,
4444
)
45-
from datadog_lambda.durable import extract_durable_function_tags
45+
from datadog_lambda.durable import (
46+
extract_durable_function_tags,
47+
extract_durable_execution_status,
48+
)
4649
from datadog_lambda.trigger import (
4750
extract_trigger_tags,
4851
extract_http_status_code_tag,
@@ -153,7 +156,7 @@ def __init__(self, func):
153156
if config.trace_extractor:
154157
extractor_parts = config.trace_extractor.rsplit(".", 1)
155158
if len(extractor_parts) == 2:
156-
(mod_name, extractor_name) = extractor_parts
159+
mod_name, extractor_name = extractor_parts
157160
modified_extractor_name = modify_module_name(mod_name)
158161
extractor_module = import_module(modified_extractor_name)
159162
self.trace_extractor = getattr(extractor_module, extractor_name)
@@ -305,6 +308,16 @@ def _after(self, event, context):
305308

306309
status_code = extract_http_status_code_tag(self.trigger_tags, self.response)
307310

311+
# Skip creating cold start spans in managed instances mode
312+
# In managed instances, the tracer library handles cold start independently
313+
should_trace_cold_start = (
314+
config.cold_start_tracing
315+
and is_new_sandbox()
316+
and not is_managed_instances_mode()
317+
)
318+
if should_trace_cold_start:
319+
trace_ctx = tracer.current_trace_context()
320+
308321
if self.span:
309322
if config.appsec_enabled and not self.blocking_response:
310323
asm_start_response(
@@ -330,6 +343,13 @@ def _after(self, event, context):
330343
if status_code:
331344
self.span.set_tag("http.status_code", status_code)
332345

346+
durable_status = extract_durable_execution_status(self.response, event)
347+
if durable_status:
348+
self.span.set_tag(
349+
"aws_lambda.durable_function.execution_status",
350+
durable_status,
351+
)
352+
333353
self.span.finish()
334354

335355
if status_code:
@@ -342,15 +362,6 @@ def _after(self, event, context):
342362
create_dd_dummy_metadata_subsegment(
343363
self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
344364
)
345-
# Skip creating cold start spans in managed instances mode
346-
# In managed instances, the tracer library handles cold start independently
347-
should_trace_cold_start = (
348-
config.cold_start_tracing
349-
and is_new_sandbox()
350-
and not is_managed_instances_mode()
351-
)
352-
if should_trace_cold_start:
353-
trace_ctx = tracer.current_trace_context()
354365

355366
if self.inferred_span:
356367
if status_code:

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ datadog = ">=0.51.0,<1.0.0"
3131
wrapt = "^1.11.2"
3232
ddtrace = [
3333
{version = ">=3.19.1,<4", python = ">=3.8,<3.10"},
34-
{version = ">=4.1.1,<5", python = ">=3.10"}
34+
{version = ">=4.1.1,<5,!=4.6.*", python = ">=3.10"}
3535
]
3636
ujson = ">=5.9.0"
3737
botocore = { version = "^1.34.0", optional = true }

scripts/build_layers.sh

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@
2020
# DD_TRACE_COMMIT Specific dd-trace-py commit SHA to build from GitHub.
2121
# DD_TRACE_COMMIT_BRANCH dd-trace-py branch name to build from GitHub.
2222
# DD_TRACE_WHEEL Path to a pre-built ddtrace .whl file.
23+
# UPSTREAM_PIPELINE_ID GitLab pipeline ID from dd-trace-py. Downloads the
24+
# matching pre-built wheel from S3 (via
25+
# index-manylinux2014.html) for each python/arch.
2326
#
2427
# Examples:
2528
# # Build a single layer for Python 3.12 on arm64
@@ -91,14 +94,6 @@ replace_ddtrace_dep() {
9194
perl -i -0777 -pe "s|ddtrace = \[[^\]]*\]|$1|gs" pyproject.toml
9295
}
9396

94-
# Replace ddtrace source if necessary
95-
if [ -n "$DD_TRACE_COMMIT" ]; then
96-
replace_ddtrace_dep "ddtrace = { git = \"https://github.com/DataDog/dd-trace-py.git\", rev = \"$DD_TRACE_COMMIT\" }"
97-
elif [ -n "$DD_TRACE_COMMIT_BRANCH" ]; then
98-
replace_ddtrace_dep "ddtrace = { git = \"https://github.com/DataDog/dd-trace-py.git\", branch = \"$DD_TRACE_COMMIT_BRANCH\" }"
99-
elif [ -n "$DD_TRACE_WHEEL" ]; then
100-
replace_ddtrace_dep "ddtrace = { file = \"$DD_TRACE_WHEEL\" }"
101-
fi
10297
function make_path_absolute {
10398
echo "$(cd "$(dirname "$1")"; pwd)/$(basename "$1")"
10499
}
@@ -109,6 +104,36 @@ function docker_build_zip {
109104
destination=$(make_path_absolute $2)
110105
arch=$3
111106

107+
# Restore pyproject.toml to a clean state for each build iteration
108+
cp pyproject.toml.bak pyproject.toml
109+
110+
# Replace ddtrace source if necessary
111+
if [ -n "$DD_TRACE_COMMIT" ]; then
112+
replace_ddtrace_dep "ddtrace = { git = \"https://github.com/DataDog/dd-trace-py.git\", rev = \"$DD_TRACE_COMMIT\" }"
113+
elif [ -n "$DD_TRACE_COMMIT_BRANCH" ]; then
114+
replace_ddtrace_dep "ddtrace = { git = \"https://github.com/DataDog/dd-trace-py.git\", branch = \"$DD_TRACE_COMMIT_BRANCH\" }"
115+
elif [ -n "$DD_TRACE_WHEEL" ]; then
116+
replace_ddtrace_dep "ddtrace = { file = \"$DD_TRACE_WHEEL\" }"
117+
elif [ -n "$UPSTREAM_PIPELINE_ID" ]; then
118+
S3_BASE="https://dd-trace-py-builds.s3.amazonaws.com/${UPSTREAM_PIPELINE_ID}"
119+
if [ "${arch}" = "amd64" ]; then
120+
PLATFORM="manylinux2014_x86_64"
121+
else
122+
PLATFORM="manylinux2014_aarch64"
123+
fi
124+
PY_TAG="cp$(echo "$1" | tr -d '.')"
125+
WHEEL_FILE=$(curl -sSfL "${S3_BASE}/index-manylinux2014.html" \
126+
| grep -o "ddtrace-[^\"]*${PY_TAG}[^\"]*${PLATFORM}[^\"]*\.whl" \
127+
| head -n 1)
128+
if [ -z "${WHEEL_FILE}" ]; then
129+
echo "No S3 wheel found for ${PY_TAG} ${PLATFORM}, using default pyproject.toml version"
130+
else
131+
curl -sSfL "${S3_BASE}/${WHEEL_FILE}" -o "${WHEEL_FILE}"
132+
echo "Using S3 wheel: ${WHEEL_FILE}"
133+
replace_ddtrace_dep "ddtrace = { file = \"${WHEEL_FILE}\" }"
134+
fi
135+
fi
136+
112137
# Install datadogpy in a docker container to avoid the mess from switching
113138
# between different python runtimes.
114139
temp_dir=$(mktemp -d)

scripts/check_layer_size.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
# Compares layer size to threshold, and fails if below that threshold
99

1010
set -e
11-
MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 9 \* 1024) # 9216 KB
11+
MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 9 \* 1024 + 15) # 9231 KB
1212
MAX_LAYER_UNCOMPRESSED_SIZE_KB=$(expr 25 \* 1024) # 25600 KB
1313

1414

tests/test_cold_start.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -310,13 +310,18 @@ def handler(event, context):
310310
handler.cold_start_tracing = True
311311
handler({}, lambda_context)
312312

313-
function_span = import_span = None
313+
function_span = import_span = load_span = None
314314
for span in spans:
315315
if span.resource == "tabnanny":
316316
import_span = span
317317
elif span.name == "aws.lambda":
318318
function_span = span
319+
elif span.name == "aws.lambda.load":
320+
load_span = span
319321

320322
assert function_span is not None
321323
assert import_span is not None
322324
assert import_span.parent_id == function_span.span_id
325+
assert import_span.trace_id == function_span.trace_id
326+
assert load_span is not None
327+
assert load_span.trace_id == function_span.trace_id

tests/test_durable.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from datadog_lambda.durable import (
88
_parse_durable_execution_arn,
99
extract_durable_function_tags,
10+
extract_durable_execution_status,
1011
)
1112

1213

@@ -113,3 +114,72 @@ def test_returns_empty_dict_when_durable_execution_arn_cannot_be_parsed(self):
113114
def test_returns_empty_dict_when_event_is_empty(self):
114115
result = extract_durable_function_tags({})
115116
self.assertEqual(result, {})
117+
118+
119+
class TestExtractDurableExecutionStatus(unittest.TestCase):
120+
def test_returns_succeeded(self):
121+
event = {
122+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
123+
}
124+
response = {"Status": "SUCCEEDED", "Result": "some-result"}
125+
self.assertEqual(extract_durable_execution_status(response, event), "SUCCEEDED")
126+
127+
def test_returns_failed(self):
128+
event = {
129+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
130+
}
131+
response = {"Status": "FAILED", "Error": "some-error"}
132+
self.assertEqual(extract_durable_execution_status(response, event), "FAILED")
133+
134+
def test_returns_stopped(self):
135+
event = {
136+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
137+
}
138+
response = {"Status": "STOPPED"}
139+
self.assertEqual(extract_durable_execution_status(response, event), "STOPPED")
140+
141+
def test_returns_timed_out(self):
142+
event = {
143+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
144+
}
145+
response = {"Status": "TIMED_OUT"}
146+
self.assertEqual(extract_durable_execution_status(response, event), "TIMED_OUT")
147+
148+
def test_returns_none_for_non_durable_event(self):
149+
event = {"key": "value"}
150+
response = {"Status": "SUCCEEDED"}
151+
self.assertIsNone(extract_durable_execution_status(response, event))
152+
153+
def test_returns_none_for_non_dict_response(self):
154+
event = {
155+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
156+
}
157+
self.assertIsNone(extract_durable_execution_status("string", event))
158+
159+
def test_returns_none_for_missing_status(self):
160+
event = {
161+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
162+
}
163+
response = {"Result": "some-result"}
164+
self.assertIsNone(extract_durable_execution_status(response, event))
165+
166+
def test_returns_none_for_invalid_status(self):
167+
event = {
168+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
169+
}
170+
response = {"Status": "INVALID"}
171+
self.assertIsNone(extract_durable_execution_status(response, event))
172+
173+
def test_returns_none_for_non_dict_event(self):
174+
response = {"Status": "SUCCEEDED"}
175+
self.assertIsNone(extract_durable_execution_status(response, "not-a-dict"))
176+
177+
def test_returns_none_for_none_event(self):
178+
response = {"Status": "SUCCEEDED"}
179+
self.assertIsNone(extract_durable_execution_status(response, None))
180+
181+
def test_returns_none_for_none_response(self):
182+
event = {
183+
"DurableExecutionArn": "arn:aws:lambda:us-east-1:123:function:f:1/durable-execution/n/id"
184+
}
185+
self.assertIsNone(extract_durable_execution_status(None, event))

0 commit comments

Comments
 (0)