fix(aws_lambda): warn when impending timeout but all traces have been finished [backport #5224] (#5265)

Backport of #5224 to 1.9

  ## Description

The root cause might be that customer is finishing all its spans, even
the one provided by `datadog-lambda`. Then hitting an impending timeout,
which makes the contrib `aws_lambda` to try to tag an error in the root
span. After this, the contrib might try to close all unfinished spans
too. Which in turn, if there's no root span, there will not be current
span, therefore, adding a conditional for both root span and current
span to act only if they are not `None`.

Discovered in issue #5220 

## Testing strategy
Added a unit test which finishes `aws_lambda` root span and then times
out. Expecting not to crash, and to not tag the error.

## Risk
User will not be getting an Impending Timeout error tagged in their root
span, but they will be receiving a warning message in their logs.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Author is aware of the performance implications of this PR as
reported in the benchmarks PR comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer is aware of, and discussed the performance implications
of this PR as reported in the benchmarks PR comment.

Author: duncanista

ddtrace/contrib/aws_lambda/patch.py

@@ -20,14 +20,17 @@ def _crash_flush(_, __):
     Tags the current root span with an Impending Timeout error.
     Finishes spans with ancestors from the current span.
     """
-
     root_span = tracer.current_root_span()
-    root_span.error = 1
-    root_span.set_tag_str(ERROR_MSG, "Datadog detected an Impending Timeout")
-    root_span.set_tag_str(ERROR_TYPE, "Impending Timeout")
+    if root_span is not None:
+        root_span.error = 1
+        root_span.set_tag_str(ERROR_MSG, "Datadog detected an Impending Timeout")
+        root_span.set_tag_str(ERROR_TYPE, "Impending Timeout")
+    else:
+        log.warning("An impending timeout was reached, but no root span was found. No error will be tagged.")
 
     current_span = tracer.current_span()
-    current_span.finish_with_ancestors()
+    if current_span is not None:
+        current_span.finish_with_ancestors()
 
 
 def _handle_signal(sig, f):

releasenotes/notes/fix-aws-lambda-crashing-when-timeout-has-no-root-span-4a0fdacf707709a9.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    aws_lambda: Resolves an exception not being handled, which occurs when no root span is found before a lambda times out.

tests/contrib/aws_lambda/handlers.py

@@ -28,6 +28,15 @@ def timeout():
     return {"success": False}
 
 
+def finishing_spans_early_handler(event, context):
+    root_span = tracer.current_root_span()
+    root_span.finish()
+
+    time.sleep(0.5)
+
+    return {"success": False}
+
+
 def datadog(_handler):
     @tracer.wrap("aws.lambda")
     def wrapped(*args):

tests/contrib/aws_lambda/test_aws_lambda.py

@@ -5,6 +5,7 @@
 from ddtrace.contrib.aws_lambda import patch
 from ddtrace.contrib.aws_lambda import unpatch
 from tests.contrib.aws_lambda.handlers import datadog
+from tests.contrib.aws_lambda.handlers import finishing_spans_early_handler
 from tests.contrib.aws_lambda.handlers import handler
 from tests.contrib.aws_lambda.handlers import manually_wrapped_handler
 from tests.contrib.aws_lambda.handlers import timeout_handler
@@ -55,6 +56,24 @@ def test_timeout_traces(context):
         assert e.type is Exception
 
 
+@pytest.mark.snapshot()
+def test_continue_on_early_trace_ending(context):
+    """
+    These scenario expects no timeout error being tagged on the root span
+    when closing all spans in the customers code and reaching a timeout.
+    """
+    os.environ.update(
+        {
+            "AWS_LAMBDA_FUNCTION_NAME": "finishing_spans_early_handler",
+            "DD_LAMBDA_HANDLER": "tests.contrib.aws_lambda.handlers.finishing_spans_early_handler",
+        }
+    )
+
+    patch()
+
+    datadog(finishing_spans_early_handler)({}, context())
+
+
 @pytest.mark.snapshot
 async def test_file_patching(context):
     os.environ.update(

tests/snapshots/tests.contrib.aws_lambda.test_aws_lambda.test_continue_on_early_trace_ending.json

@@ -0,0 +1,25 @@
+[[
+  {
+    "name": "aws.lambda",
+    "service": "",
+    "resource": "aws.lambda",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.p.dm": "-0",
+      "language": "python",
+      "runtime-id": "2a38d844990949379810ed84bb6e71d7"
+    },
+    "metrics": {
+      "_dd.agent_psr": 1.0,
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 55112
+    },
+    "duration": 6433000,
+    "start": 1677616801872185000
+  }]]