From f0e8c5e0b36e8fac07af48f8ab8fe13629faabe8 Mon Sep 17 00:00:00 2001
From: Tal Usvyatsky <tal.usvyatsky@datadoghq.com>
Date: Wed, 24 Apr 2024 09:25:13 -0400
Subject: [PATCH] fix duration

---
 datadog_lambda/handler.py | 8 +++++---
 datadog_lambda/tracing.py | 6 ++++--
 datadog_lambda/wrapper.py | 6 +++---
 tests/test_handler.py     | 7 +++----
 tests/test_tracing.py     | 5 +++--
 5 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/datadog_lambda/handler.py b/datadog_lambda/handler.py
index c0cbeafd..cf2af5a0 100644
--- a/datadog_lambda/handler.py
+++ b/datadog_lambda/handler.py
@@ -7,7 +7,7 @@
 from importlib import import_module
 
 import os
-import time
+from time import time_ns
 
 from datadog_lambda.wrapper import datadog_lambda_wrapper, error_fallback_handler
 from datadog_lambda.module_name import modify_module_name
@@ -31,9 +31,11 @@ class HandlerError(Exception):
 modified_mod_name = modify_module_name(mod_name)
 
 try:
-    start_time_ns = time.time_ns()
+    handler_load_start_time_ns = time_ns()
     handler_module = import_module(modified_mod_name)
     handler_func = getattr(handler_module, handler_name)
     handler = datadog_lambda_wrapper(handler_func)
 except Exception as e:
-    handler = error_fallback_handler(e, modified_mod_name, start_time_ns)
+    handler = error_fallback_handler(
+        e, modified_mod_name, time_ns() - handler_load_start_time_ns
+    )
diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
index c04a99c6..9be25744 100644
--- a/datadog_lambda/tracing.py
+++ b/datadog_lambda/tracing.py
@@ -9,6 +9,7 @@
 import traceback
 import ujson as json
 from datetime import datetime, timezone
+from time import time_ns
 from typing import Optional, Dict
 
 from datadog_lambda.metric import submit_errors_metric
@@ -1324,7 +1325,7 @@ def is_async(span: Span) -> bool:
 
 
 def emit_telemetry_on_exception_outside_of_handler(
-    context, exception, resource_name, start_time_ns
+    context, exception, resource_name, handler_load_duration
 ):
     """
     Emit an enhanced error metric and create a span for exceptions occuring outside of the handler
@@ -1337,7 +1338,8 @@ def emit_telemetry_on_exception_outside_of_handler(
         resource=resource_name,
         span_type="serverless",
     )
-    span.start_ns = start_time_ns
+    span.start_ns = time_ns() - handler_load_duration
+
     tags = {
         "error.status": 500,
         "error.type": type(exception).__name__,
diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
index 54685918..da5263a1 100644
--- a/datadog_lambda/wrapper.py
+++ b/datadog_lambda/wrapper.py
@@ -389,17 +389,17 @@ class _ErrorFallbackHandler(object):
     Emits telemetry and re-raises the exception.
     """
 
-    def __init__(self, exception, modified_mod_name, start_time_ns):
+    def __init__(self, exception, modified_mod_name, handler_load_duration_ns):
         self.exception = exception
         self.modified_mod_name = modified_mod_name
-        self.start_time_ns = start_time_ns
+        self.handler_load_duration_ns = handler_load_duration_ns
 
     def __call__(self, event, context, **kwargs):
         emit_telemetry_on_exception_outside_of_handler(
             context,
             self.exception,
             self.modified_mod_name,
-            self.start_time_ns,
+            self.handler_load_duration_ns,
         )
         raise self.exception
 
diff --git a/tests/test_handler.py b/tests/test_handler.py
index d8c7b06a..7066dd06 100644
--- a/tests/test_handler.py
+++ b/tests/test_handler.py
@@ -39,7 +39,7 @@ def test_exception_importing_module(self, mock_time, mock_emit_telemetry):
             lambda_context = get_mock_context()
             datadog_lambda.handler.handler.__call__(None, lambda_context)
         mock_emit_telemetry.assert_called_once_with(
-            lambda_context, test_context.exception, "nonsense", 42
+            lambda_context, test_context.exception, "nonsense", 0
         )
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
@@ -55,16 +55,15 @@ def test_exception_getting_handler_func(
             lambda_context = get_mock_context()
             datadog_lambda.handler.handler.__call__(None, lambda_context)
         mock_emit_telemetry.assert_called_once_with(
-            lambda_context, test_context.exception, "nonsense", 42
+            lambda_context, test_context.exception, "nonsense", 0
         )
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
     @patch("importlib.import_module")
     @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
-    @patch("time.time_ns", return_value=42)
     @patch("datadog_lambda.wrapper.datadog_lambda_wrapper")
     def test_handler_success(
-        self, mock_lambda_wrapper, mock_time, mock_emit_telemetry, mock_import
+        self, mock_lambda_wrapper, mock_emit_telemetry, mock_import
     ):
         def nonsense():
             pass
diff --git a/tests/test_tracing.py b/tests/test_tracing.py
index f81e4941..8f40ed72 100644
--- a/tests/test_tracing.py
+++ b/tests/test_tracing.py
@@ -2005,7 +2005,8 @@ def test_deterministic_m5_hash__always_leading_with_zero(self):
 
 class TestExceptionOutsideHandler(unittest.TestCase):
     @patch("datadog_lambda.tracing.submit_errors_metric")
-    def test_exception_outside_handler(self, mock_submit_errors_metric):
+    @patch("datadog_lambda.tracing.time_ns", return_value=100)
+    def test_exception_outside_handler(self, mock_time, mock_submit_errors_metric):
         fake_error = ValueError("Some error message")
         resource_name = "my_handler"
         span_type = "aws.lambda"
@@ -2045,4 +2046,4 @@ def test_exception_outside_handler(self, mock_submit_errors_metric):
         )
         mock_span.finish.assert_called_once()
         assert mock_span.error == 1
-        assert mock_span.start_ns == 42
+        assert mock_span.start_ns == 58