Track in-flight event requests

Author: imjoehaines

bugsnag/client.py

@@ -19,6 +19,7 @@
 from bugsnag.sessiontracker import SessionTracker
 from bugsnag.utils import to_rfc3339
 from bugsnag.context import ContextLocalState
+from bugsnag.request_tracker import RequestTracker
 
 __all__ = ('Client',)
 
@@ -36,6 +37,7 @@ def __init__(self, configuration: Optional[Configuration] = None,
         self.session_tracker = SessionTracker(self.configuration)
         self.configuration.configure(**kwargs)
         self._context = ContextLocalState(self)
+        self._request_tracker = RequestTracker()
 
         if install_sys_hook:
             self.install_sys_hook()
@@ -174,11 +176,6 @@ def run_middleware():
             initial_reason = event.severity_reason.copy()
 
             def send_payload():
-                if asynchronous is None:
-                    options = {}
-                else:
-                    options = {'asynchronous': asynchronous}
-
                 if event.api_key is None:
                     self.configuration.logger.warning(
                         "No API key configured, couldn't notify"
@@ -192,16 +189,30 @@ def send_payload():
                     }
                 else:
                     event.severity_reason = initial_reason
+
                 payload = event._payload()
+
+                post_delivery_callback = self._request_tracker.new_request()
+                options = {'post_delivery_callback': post_delivery_callback}
+
+                if asynchronous is not None:
+                    options['asynchronous'] = asynchronous
+
                 try:
-                    self.configuration.delivery.deliver(self.configuration,
-                                                        payload, options)
+                    self.configuration.delivery.deliver(
+                        self.configuration,
+                        payload,
+                        options
+                    )
                 except Exception as e:
                     self.configuration.logger.exception(
                         'Notifying Bugsnag failed %s',
                         e
                     )
 
+                    # ensure this request is not still marked as in-flight
+                    post_delivery_callback()
+
                 # Trigger session delivery
                 self.session_tracker.send_sessions()
 

tests/test_client.py

@@ -17,7 +17,12 @@
 )
 
 import bugsnag.legacy as legacy
-from tests.utils import IntegrationTest, ScaryException
+from tests.utils import (
+    BrokenDelivery,
+    IntegrationTest,
+    QueueingDelivery,
+    ScaryException,
+)
 
 timestamp_regex = r'^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}(?:[+-]\d{2}:\d{2}|Z)'  # noqa: E501
 
@@ -1688,6 +1693,62 @@ def on_error(event):
             FeatureFlag('e')
         ]
 
+    def test_in_flight_event_request_tracking_in_notify(self):
+        delivery = QueueingDelivery()
+        configuration = Configuration()
+        configuration.configure(delivery=delivery, api_key='abc')
+
+        client = Client(configuration)
+        assert not client._request_tracker.has_in_flight_requests()
+
+        client.notify(Exception('Oh no'))
+        assert client._request_tracker.has_in_flight_requests()
+
+        delivery.flush_request_queue()
+        assert not client._request_tracker.has_in_flight_requests()
+
+    def test_in_flight_event_request_tracking_in_notify_failure(self):
+        configuration = Configuration()
+        configuration.configure(delivery=BrokenDelivery(), api_key='abc')
+
+        client = Client(configuration)
+        assert not client._request_tracker.has_in_flight_requests()
+
+        client.notify(Exception('Oh no'))
+        assert not client._request_tracker.has_in_flight_requests()
+
+    def test_in_flight_event_request_tracking_in_notify_exc_info(self):
+        delivery = QueueingDelivery()
+        configuration = Configuration()
+        configuration.configure(delivery=delivery, api_key='abc')
+
+        client = Client(configuration)
+        assert not client._request_tracker.has_in_flight_requests()
+
+        try:
+            raise Exception(':)')
+        except Exception:
+            client.notify_exc_info(*sys.exc_info())
+
+        assert client._request_tracker.has_in_flight_requests()
+
+        delivery.flush_request_queue()
+        assert not client._request_tracker.has_in_flight_requests()
+
+    def test_in_flight_event_request_tracking_in_notify_exc_info_failure(self):
+        configuration = Configuration()
+        configuration.configure(delivery=BrokenDelivery(), api_key='abc')
+
+        client = Client(configuration)
+        assert not client._request_tracker.has_in_flight_requests()
+
+        try:
+            raise Exception(':)')
+        except Exception:
+            client.notify_exc_info(*sys.exc_info())
+
+        assert not client._request_tracker.has_in_flight_requests()
+
 
 @pytest.mark.parametrize("metadata,type", [
     (1234, 'int'),