refactor: AWS Lambda instrumentation.

Signed-off-by: Paulo Vital <[email protected]>

Author: pvital

src/instana/__init__.py

@@ -13,6 +13,7 @@
 import importlib
 import os
 import sys
+from typing import Tuple
 
 from instana.collector.helpers.runtime import (
     is_autowrapt_instrumented,
@@ -55,7 +56,7 @@
 ]
 
 
-def load(_):
+def load(_: object) -> None:
     """
     Method used to activate the Instana sensor via AUTOWRAPT_BOOTSTRAP
     environment variable.
@@ -66,15 +67,15 @@ def load(_):
     return None
 
 
-def apply_gevent_monkey_patch():
+def apply_gevent_monkey_patch() -> None:
     from gevent import monkey
 
     if os.environ.get("INSTANA_GEVENT_MONKEY_OPTIONS"):
 
-        def short_key(k):
+        def short_key(k: str) -> str:
             return k[3:] if k.startswith("no-") else k
 
-        def key_to_bool(k):
+        def key_to_bool(k: str) -> bool:
             return not k.startswith("no-")
 
         import inspect
@@ -99,14 +100,16 @@ def key_to_bool(k):
         monkey.patch_all()
 
 
-def get_lambda_handler_or_default():
+def get_aws_lambda_handler() -> Tuple[str, str]:
     """
-    For instrumenting AWS Lambda, users specify their original lambda handler in the LAMBDA_HANDLER environment
-    variable.  This function searches for and parses that environment variable or returns the defaults.
-
-    The default handler value for AWS Lambda is 'lambda_function.lambda_handler' which
-    equates to the function "lambda_handler in a file named "lambda_function.py" or in Python
-    terms "from lambda_function import lambda_handler"
+    For instrumenting AWS Lambda, users specify their original lambda handler 
+    in the LAMBDA_HANDLER environment variable.  This function searches for and 
+    parses that environment variable or returns the defaults.
+
+    The default handler value for AWS Lambda is 'lambda_function.lambda_handler'
+    which equates to the function "lambda_handler in a file named 
+    lambda_function.py" or in Python terms 
+    "from lambda_function import lambda_handler"
     """
     handler_module = "lambda_function"
     handler_function = "lambda_handler"
@@ -118,20 +121,20 @@ def get_lambda_handler_or_default():
             parts = handler.split(".")
             handler_function = parts.pop().strip()
             handler_module = ".".join(parts).strip()
-    except Exception:
-        pass
+    except Exception as exc:
+        print(f"get_aws_lambda_handler error: {exc}")
 
     return handler_module, handler_function
 
 
-def lambda_handler(event, context):
+def lambda_handler(event: str, context: str) -> None:
     """
     Entry point for AWS Lambda monitoring.
 
     This function will trigger the initialization of Instana monitoring and then call
     the original user specified lambda handler function.
     """
-    module_name, function_name = get_lambda_handler_or_default()
+    module_name, function_name = get_aws_lambda_handler()
 
     try:
         # Import the module specified in module_name
@@ -151,7 +154,7 @@ def lambda_handler(event, context):
             )
 
 
-def boot_agent():
+def boot_agent() -> None:
     """Initialize the Instana agent and conditionally load auto-instrumentation."""
 
     import instana.singletons  # noqa: F401
@@ -189,8 +192,7 @@ def boot_agent():
             client,  # noqa: F401
             server,  # noqa: F401
         )
-
-        # from instana.instrumentation.aws import lambda_inst  # noqa: F401
+        from instana.instrumentation.aws import lambda_inst  # noqa: F401
         from instana.instrumentation import celery  # noqa: F401
         from instana.instrumentation.django import middleware  # noqa: F401
         from instana.instrumentation.google.cloud import (
@@ -229,7 +231,7 @@ def boot_agent():
             apply_gevent_monkey_patch()
         # AutoProfile
         if "INSTANA_AUTOPROFILE" in os.environ:
-            from .singletons import get_profiler
+            from instana.singletons import get_profiler
 
             profiler = get_profiler()
             if profiler:

src/instana/agent/aws_lambda.py

@@ -2,21 +2,23 @@
 # (c) Copyright Instana Inc. 2020
 
 """
-The Instana agent (for AWS Lambda functions) that manages
+The Instana Agent for AWS Lambda functions that manages
 monitoring state and reporting that data.
 """
-import time
-from ..log import logger
-from ..util import to_json
-from .base import BaseAgent
-from ..version import VERSION
-from ..collector.aws_lambda import AWSLambdaCollector
-from ..options import AWSLambdaOptions
+
+from typing import Any, Dict
+from instana.agent.base import BaseAgent
+from instana.collector.aws_lambda import AWSLambdaCollector
+from instana.log import logger
+from instana.options import AWSLambdaOptions
+from instana.util import to_json
+from instana.version import VERSION
 
 
 class AWSLambdaAgent(BaseAgent):
-    """ In-process agent for AWS Lambda """
-    def __init__(self):
+    """In-process Agent for AWS Lambda"""
+
+    def __init__(self) -> None:
         super(AWSLambdaAgent, self).__init__()
 
         self.collector = None
@@ -27,29 +29,33 @@ def __init__(self):
         # Update log level from what Options detected
         self.update_log_level()
 
-        logger.info("Stan is on the AWS Lambda scene.  Starting Instana instrumentation version: %s", VERSION)
+        logger.info(
+            f"Stan is on the AWS Lambda scene.  Starting Instana instrumentation version: {VERSION}",
+        )
 
         if self._validate_options():
             self._can_send = True
             self.collector = AWSLambdaCollector(self)
             self.collector.start()
         else:
-            logger.warning("Required INSTANA_AGENT_KEY and/or INSTANA_ENDPOINT_URL environment variables not set.  "
-                           "We will not be able monitor this function.")
+            logger.warning(
+                "Required INSTANA_AGENT_KEY and/or INSTANA_ENDPOINT_URL environment variables not set.  "
+                "We will not be able monitor this function."
+            )
 
-    def can_send(self):
+    def can_send(self) -> bool:
         """
         Are we in a state where we can send data?
         @return: Boolean
         """
         return self._can_send
 
-    def get_from_structure(self):
+    def get_from_structure(self) -> Dict[str, Any]:
         """
         Retrieves the From data that is reported alongside monitoring data.
         @return: dict()
         """
-        return {'hl': True, 'cp': 'aws', 'e': self.collector.get_fq_arn()}
+        return {"hl": True, "cp": "aws", "e": self.collector.get_fq_arn()}
 
     def report_data_payload(self, payload):
         """
@@ -64,30 +70,38 @@ def report_data_payload(self, payload):
                 self.report_headers["X-Instana-Host"] = self.collector.get_fq_arn()
                 self.report_headers["X-Instana-Key"] = self.options.agent_key
 
-            response = self.client.post(self.__data_bundle_url(),
-                                        data=to_json(payload),
-                                        headers=self.report_headers,
-                                        timeout=self.options.timeout,
-                                        verify=self.options.ssl_verify,
-                                        proxies=self.options.endpoint_proxy)
+            response = self.client.post(
+                self.__data_bundle_url(),
+                data=to_json(payload),
+                headers=self.report_headers,
+                timeout=self.options.timeout,
+                verify=self.options.ssl_verify,
+                proxies=self.options.endpoint_proxy,
+            )
 
             if 200 <= response.status_code < 300:
-                logger.debug("report_data_payload: Instana responded with status code %s", response.status_code)
+                logger.debug(
+                    "report_data_payload: Instana responded with status code %s",
+                    response.status_code,
+                )
             else:
-                logger.info("report_data_payload: Instana responded with status code %s", response.status_code)
+                logger.info(
+                    "report_data_payload: Instana responded with status code %s",
+                    response.status_code,
+                )
         except Exception as exc:
             logger.debug("report_data_payload: connection error (%s)", type(exc))
 
         return response
 
-    def _validate_options(self):
+    def _validate_options(self) -> bool:
         """
         Validate that the options used by this Agent are valid.  e.g. can we report data?
         """
-        return self.options.endpoint_url is not None and self.options.agent_key is not None
+        return self.options.endpoint_url and self.options.agent_key
 
-    def __data_bundle_url(self):
+    def __data_bundle_url(self) -> str:
         """
         URL for posting metrics to the host agent.  Only valid when announced.
         """
-        return "%s/bundle" % self.options.endpoint_url
+        return f"{self.options.endpoint_url}/bundle"

src/instana/instrumentation/aws/lambda_inst.py

@@ -4,66 +4,89 @@
 """
 Instrumentation for AWS Lambda functions
 """
+
 import sys
+import traceback
+from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple
+
 import wrapt
-import opentracing.ext.tags as ext
+from opentelemetry.semconv.trace import SpanAttributes
 
-from ...log import logger
-from ...singletons import env_is_aws_lambda
-from ... import get_lambda_handler_or_default
-from ...singletons import get_agent, get_tracer
-from .triggers import enrich_lambda_span, get_context
-import traceback
+from instana import get_aws_lambda_handler
+from instana.instrumentation.aws.triggers import enrich_lambda_span, get_context
+from instana.log import logger
+from instana.singletons import env_is_aws_lambda, get_agent, get_tracer
+
+if TYPE_CHECKING:
+    from instana.agent.aws_lambda import AWSLambdaAgent
 
 
-def lambda_handler_with_instana(wrapped, instance, args, kwargs):
+def lambda_handler_with_instana(
+    wrapped: Callable[..., object],
+    instance: object,
+    args: Tuple[object, ...],
+    kwargs: Dict[str, Any],
+) -> object:
     event = args[0]
-    agent = get_agent()
+    agent: "AWSLambdaAgent" = get_agent()
     tracer = get_tracer()
 
     agent.collector.collect_snapshot(*args)
     incoming_ctx = get_context(tracer, event)
 
     result = None
-    with tracer.start_active_span("aws.lambda.entry", child_of=incoming_ctx) as scope:
-        enrich_lambda_span(agent, scope.span, *args)
+    with tracer.start_as_current_span(
+        "aws.lambda.entry", span_context=incoming_ctx
+    ) as span:
+        enrich_lambda_span(agent, span, *args)
         try:
             result = wrapped(*args, **kwargs)
 
             if isinstance(result, dict):
-                server_timing_value = "intid;desc=%s" % scope.span.context.trace_id
-                if 'headers' in result:
-                    result['headers']['Server-Timing'] = server_timing_value
-                elif 'multiValueHeaders' in result:
-                    result['multiValueHeaders']['Server-Timing'] = [server_timing_value]
-                if 'statusCode' in result and result.get('statusCode'):
-                    status_code = int(result['statusCode'])
-                    scope.span.set_tag(ext.HTTP_STATUS_CODE, status_code)
+                server_timing_value = f"intid;desc={span.context.trace_id}"
+                if "headers" in result:
+                    result["headers"]["Server-Timing"] = server_timing_value
+                elif "multiValueHeaders" in result:
+                    result["multiValueHeaders"]["Server-Timing"] = [server_timing_value]
+                if "statusCode" in result and result.get("statusCode"):
+                    status_code = int(result["statusCode"])
+                    span.set_attribute(SpanAttributes.HTTP_STATUS_CODE, status_code)
                     if 500 <= status_code:
-                        scope.span.log_exception(f'HTTP status {status_code}')
+                        span.record_exception(f"HTTP status {status_code}")
         except Exception as exc:
-            if scope.span:
+            logger.debug(f"AWS Lambda lambda_handler_with_instana error: {exc}")
+            if span:
                 exc = traceback.format_exc()
-                scope.span.log_exception(exc)
+                span.record_exception(exc)
             raise
         finally:
-            scope.close()
             agent.collector.shutdown()
 
-    agent.collector.shutdown()
+    if agent.collector.started:
+        agent.collector.shutdown()
+    
     return result
 
 
-if env_is_aws_lambda is True:
-    handler_module, handler_function = get_lambda_handler_or_default()
+if env_is_aws_lambda:
+    handler_module, handler_function = get_aws_lambda_handler()
 
-    if handler_module is not None and handler_function is not None:
+    if handler_module and handler_function:
         try:
-            logger.debug("Instrumenting AWS Lambda handler (%s.%s)" % (handler_module, handler_function))
-            sys.path.insert(0, '/var/runtime')
-            sys.path.insert(0, '/var/task')
-            wrapt.wrap_function_wrapper(handler_module, handler_function, lambda_handler_with_instana)
+            logger.debug(
+                f"Instrumenting AWS Lambda handler ({handler_module}.{handler_function})"
+            )
+            sys.path.insert(0, "/var/runtime")
+            sys.path.insert(0, "/var/task")
+            wrapt.wrap_function_wrapper(
+                handler_module, handler_function, lambda_handler_with_instana
+            )
         except (ModuleNotFoundError, ImportError) as exc:
-            logger.warning("Instana: Couldn't instrument AWS Lambda handler. Not monitoring.")
+            logger.debug(f"AWS Lambda error: {exc}")
+            logger.warning(
+                "Instana: Couldn't instrument AWS Lambda handler. Not monitoring."
+            )
     else:
-        logger.warning("Instana: Couldn't determine AWS Lambda Handler.  Not monitoring.")
+        logger.warning(
+            "Instana: Couldn't determine AWS Lambda Handler.  Not monitoring."
+        )