Add warning when lambda function may timeout

bugsnag/client.py

@@ -378,12 +378,14 @@ def aws_lambda_handler(
         self,
         real_handler: Optional[Callable] = None,
         flush_timeout_ms: int = 2000,
+        lambda_timeout_notify_ms: int = 1000,
     ) -> Callable:
         # handle being called with just 'flush_timeout_ms'
         if real_handler is None:
             return functools.partial(
                 self.aws_lambda_handler,
                 flush_timeout_ms=flush_timeout_ms,
+                lambda_timeout_notify_ms=lambda_timeout_notify_ms,
             )
 
         # attributes from the aws context that we want to capture as metadata
@@ -403,13 +405,65 @@ def aws_lambda_handler(
 
         @functools.wraps(real_handler)
         def wrapped_handler(aws_event, aws_context):
-            try:
-                aws_context_metadata = {
-                    attribute:
-                        getattr(aws_context, attribute, None)
-                        for attribute in aws_context_attributes
-                }
+            timer = None
+            aws_context_metadata = {
+                attribute:
+                    getattr(aws_context, attribute, None)
+                    for attribute in aws_context_attributes
+            }
+
+            if lambda_timeout_notify_ms > 0:
+                # reporting possible timeouts is done using a separate thread,
+                # but we don't want to lose the information from the main
+                # thread so we store references here to use later
+                # TODO: we shouldn't have 3 places where per-request data is
+                #       stored - it should all be in 'self._context'
+                main_request_config = RequestConfiguration.get_instance()
+                main_request_breadcrumbs = \
+                    self.configuration._breadcrumbs._breadcrumbs
+                main_request_feature_flags = \
+                    self._context.feature_flag_delegate._storage
+
+                def report_timeout_to_bugsnag():
+                    # copy over the main thread's data to this thread
+                    RequestConfiguration.set_instance(main_request_config)
+
+                    self.configuration._breadcrumbs._breadcrumbs.extend(
+                        main_request_breadcrumbs
+                    )
+
+                    self._context.feature_flag_delegate.merge(
+                        main_request_feature_flags.values()
+                    )
+
+                    # generate an empty traceback object so the lambda timeout
+                    # doesn't have a misleading traceback
+                    try:
+                        raise Exception()
+                    except Exception as exception:
+                        empty_traceback = exception.__traceback__
+
+                    lambda_timeout_approaching = LambdaTimeoutApproaching(
+                        aws_context.get_remaining_time_in_millis(),
+                        empty_traceback
+                    )
+
+                    # set the source_func so the user's lambda handler is the
+                    # only item in the traceback
+                    self.notify(
+                        lambda_timeout_approaching,
+                        source_func=real_handler
+                    )
+
+                remaining_ms = aws_context.get_remaining_time_in_millis()
+                timer = threading.Timer(
+                    (remaining_ms - lambda_timeout_notify_ms) / 1000,
+                    report_timeout_to_bugsnag
+                )
+
+                timer.start()
 
+            try:
                 self.add_metadata_tab('AWS Lambda Event', aws_event)
                 self.add_metadata_tab(
                     'AWS Lambda Context',
@@ -426,6 +480,10 @@ def wrapped_handler(aws_event, aws_context):
 
                 raise
             finally:
+                # a timer can only be cancelled if it hasn't fired yet
+                if timer and timer.is_alive():
+                    timer.cancel()
+
                 try:
                     self.flush(flush_timeout_ms)
                 except Exception as exception:
@@ -466,3 +524,9 @@ def __exit__(self, *exc_info):
                 self.client.notify_exc_info(*exc_info, **self.options)
 
         return False
+
+
+class LambdaTimeoutApproaching(Exception):
+    def __init__(self, remaining_ms: int, tb):
+        super().__init__('Lambda will timeout in %dms' % remaining_ms)
+        self.__traceback__ = tb

bugsnag/configuration.py

@@ -603,6 +603,10 @@ def get_instance(cls):
 
         return instance
 
+    @classmethod
+    def set_instance(cls, instance: 'RequestConfiguration') -> None:
+        _request_info.set(instance)  # type: ignore
+
     @classmethod
     def clear(cls):
         """

bugsnag/utils.py

@@ -261,7 +261,7 @@ def is_json_content_type(value: str) -> bool:
     return type == 'application' and (subtype == 'json' or suffix == 'json')
 
 
-_ignore_modules = ('__main__', 'builtins')
+_ignore_modules = ('__main__', 'builtins', 'bugsnag.client')
 
 
 def partly_qualified_class_name(obj):

tests/test_client.py

@@ -1916,6 +1916,69 @@ def my_handler(event, context):
         assert self.sent_report_count == 0
         assert self.sent_session_count == 0
 
+    def test_aws_lambda_handler_decorator_warns_of_potential_timeout(self):
+        aws_lambda_context = LambdaContext(remaining_time_in_millis=2)
+
+        @self.client.aws_lambda_handler(lambda_timeout_notify_ms=1)
+        def my_handler(event, context):
+            assert event == {'z': 9}
+            assert context == aws_lambda_context
+
+            self.client.leave_breadcrumb('hello 1')
+            self.client.leave_breadcrumb('hello 2')
+            self.client.leave_breadcrumb('hello 3')
+
+            self.client.add_feature_flag('a')
+            self.client.add_feature_flag('b', '1')
+            self.client.add_feature_flag('c')
+
+            time.sleep(0.1)
+
+        my_handler({'z': 9}, aws_lambda_context)
+
+        assert self.sent_report_count == 1
+        assert self.sent_session_count == 1
+
+        payload = self.server.events_received[0]['json_body']
+        event = payload['events'][0]
+
+        assert event['metaData']['AWS Lambda Event'] == {'z': 9}
+        assert event['metaData']['AWS Lambda Context'] == {
+            'function_name': 'function_name',
+            'function_version': 'function_version',
+            'invoked_function_arn': 'invoked_function_arn',
+            'memory_limit_in_mb': 'memory_limit_in_mb',
+            'aws_request_id': 'aws_request_id',
+            'log_group_name': 'log_group_name',
+            'log_stream_name': 'log_stream_name',
+            'identity': 'identity',
+            'client_context': 'client_context',
+        }
+
+        assert len(event['breadcrumbs']) == 3
+        assert event['breadcrumbs'][0]['name'] == 'hello 1'
+        assert event['breadcrumbs'][1]['name'] == 'hello 2'
+        assert event['breadcrumbs'][2]['name'] == 'hello 3'
+
+        assert event['featureFlags'] == [
+            {'featureFlag': 'a'},
+            {'featureFlag': 'b', 'variant': '1'},
+            {'featureFlag': 'c'},
+        ]
+
+        exception = event['exceptions'][0]
+
+        assert exception['message'] == 'Lambda will timeout in 2ms'
+        assert exception['errorClass'] == 'LambdaTimeoutApproaching'
+
+        # the stacktrace should have a single frame pointing to the user's
+        # lambda handler
+        stacktrace = exception['stacktrace']
+
+        assert len(stacktrace) == 1
+        assert stacktrace[0]['file'] == 'test_client.py'
+        assert stacktrace[0]['method'] == 'my_handler'
+
 
 @pytest.mark.parametrize("metadata,type", [
     (1234, 'int'),
@@ -1960,6 +2023,7 @@ def __init__(
         log_stream_name='log_stream_name',
         identity='identity',
         client_context='client_context',
+        remaining_time_in_millis=10000
     ):
         self.function_name = function_name
         self.function_version = function_version
@@ -1971,3 +2035,7 @@ def __init__(
         self.identity = identity
         self.client_context = client_context
         self.another_attribute = 'another_attribute'
+        self.remaining_time_in_millis = remaining_time_in_millis
+
+    def get_remaining_time_in_millis(self) -> int:
+        return self.remaining_time_in_millis