From f156e344dc48adf91be53e65fb24638876bf9612 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 18 Apr 2025 13:00:56 -0400 Subject: [PATCH 1/2] add `awslambdaric` integration --- ddtrace/contrib/_awslambdaric.py | 25 +++++++++++++++++++ .../contrib/internal/awslambdaric/__init__.py | 3 +++ 2 files changed, 28 insertions(+) create mode 100644 ddtrace/contrib/_awslambdaric.py create mode 100644 ddtrace/contrib/internal/awslambdaric/__init__.py diff --git a/ddtrace/contrib/_awslambdaric.py b/ddtrace/contrib/_awslambdaric.py new file mode 100644 index 00000000000..0956cbe7138 --- /dev/null +++ b/ddtrace/contrib/_awslambdaric.py @@ -0,0 +1,25 @@ +""" +The awslambdaric integration traces AWS Lambda Runtime Interface Client (RIC) operations. + +Enabling +~~~~~~~~ + +The awslambdaric integration is enabled automatically when using +``ddtrace-run`` or ``import ddtrace.auto``. + +Global Configuration +~~~~~~~~~~~~~~~~~~~ + +.. py:data:: DD_AWSLAMBDARIC_SERVICE + + The service name reported by default for awslambdaric spans. + + This option can also be set with the ``config.awslambdaric['service_name']`` setting. + + Default: ``'awslambdaric'`` + +Integration Configuration +~~~~~~~~~~~~~~~~~~~~~~~ + +No additional configuration options are available for this integration. +""" \ No newline at end of file diff --git a/ddtrace/contrib/internal/awslambdaric/__init__.py b/ddtrace/contrib/internal/awslambdaric/__init__.py new file mode 100644 index 00000000000..791ee83e8e8 --- /dev/null +++ b/ddtrace/contrib/internal/awslambdaric/__init__.py @@ -0,0 +1,3 @@ +""" +The awslambdaric integration traces AWS Lambda Runtime Interface Client (RIC) operations. +""" \ No newline at end of file From 252984ef3cbe7f0bdcbc123d254f4c4ab330ec71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?jordan=20gonz=C3=A1lez?= <30836115+duncanista@users.noreply.github.com> Date: Fri, 18 Apr 2025 13:01:24 -0400 Subject: [PATCH 2/2] add `patch.py` for integration proof of concept that AWS Lambda runtime can be patched --- .../contrib/internal/awslambdaric/patch.py | 109 ++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 ddtrace/contrib/internal/awslambdaric/patch.py diff --git a/ddtrace/contrib/internal/awslambdaric/patch.py b/ddtrace/contrib/internal/awslambdaric/patch.py new file mode 100644 index 00000000000..e4c69deeb2b --- /dev/null +++ b/ddtrace/contrib/internal/awslambdaric/patch.py @@ -0,0 +1,109 @@ +""" +Datadog instrumentation for AWS Lambda Runtime Interface Client (RIC). +""" + +from ddtrace import config +from ddtrace.contrib.internal.trace_utils import with_traced_module +from ddtrace.internal.logger import get_logger +from ddtrace.internal.utils import get_argument_value +from ddtrace.internal.wrapping import unwrap +from ddtrace.internal.wrapping import wrap +from ddtrace.trace import Pin + +log = get_logger(__name__) + +config._add( + "awslambdaric", + dict( + _default_service="awslambdaric", + ), +) + +class DatadogInstrumentation(object): + """Patches an AWS Lambda handler function for Datadog instrumentation.""" + + def __call__(self, func, args, kwargs): + self.func = func + self._before(args, kwargs) + try: + self.response = self.func(*args, **kwargs) + return self.response + finally: + self._after() + + def _set_context(self, args, kwargs): + """Sets the context attribute.""" + # The context is the second argument in a handler + # signature and it is always sent. + # + # note: AWS Lambda context is an object, the event is a dict. + # `get_remaining_time_in_millis` is guaranteed to be + # present in the context. + _context = get_argument_value(args, kwargs, 1, "context") + if hasattr(_context, "get_remaining_time_in_millis"): + self.context = _context + else: + # Handler was possibly manually wrapped, and the first + # argument is the `datadog-lambda` decorator object. + self.context = get_argument_value(args, kwargs, 2, "context") + + def _before(self, args, kwargs): + log.info("[awslambdaric][handler] _before was called") + self._set_context(args, kwargs) + + def _after(self): + log.info("[awslambdaric][handler] _after was called") + pass + +def get_version(): + # type: () -> str + try: + import awslambdaric + return getattr(awslambdaric, "__version__", "") + except ImportError: + return "" + +def patch(): + """ + Patch the AWS Lambda Runtime Interface Client to enable tracing. + """ + import awslambdaric + + if getattr(awslambdaric, "__datadog_patch", False): + return + + Pin().onto(awslambdaric) + + wrap(awslambdaric, "bootstrap.handle_event_request", _wrap_handle_event_request) + + config._awslambdaric_patched = True + + +def unpatch(): + """ + Remove AWS Lambda Runtime Interface Client tracing. + """ + import awslambdaric + + if getattr(awslambdaric, "__datadog_patch", False): + return + + awslambdaric.__datadog_patch = False + +@with_traced_module +def _wrap_handle_event_request(awslambdaric, pin, func, instance, args, kwargs): + """ + Wrap the handle_event_request method. + """ + _invoke_id = get_argument_value(args, kwargs, 2, "invoke_id") + _event_body = get_argument_value(args, kwargs, 3, "event_body") + log.info("[awslambdaric] handle_event_request was called with invoke_id: %s, event_body: %s", _invoke_id, _event_body) + + _customer_handler = get_argument_value(args, kwargs, 1, "request_handler") + _datadog_handler = DatadogInstrumentation() + + wrap(_customer_handler, _datadog_handler) + + func(*args, **kwargs) + +