Track in flight requests

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()
 

bugsnag/request_tracker.py

@@ -0,0 +1,49 @@
+from uuid import uuid4
+from threading import Lock
+from typing import Callable
+
+
+class RequestTracker:
+    def __init__(self):
+        self._mutex = Lock()
+        self._requests = set()  # type: set[str]
+
+    def new_request(self) -> Callable[[], None]:
+        """
+        Track a new request, returning a callback that marks the request as
+        complete.
+
+        >>> request_tracker = RequestTracker()
+        >>> mark_request_complete = request_tracker.new_request()
+        >>> # ...make the request...
+        >>> mark_request_complete()
+        """
+        request_id = uuid4().hex
+
+        with self._mutex:
+            self._requests.add(request_id)
+
+        def mark_request_complete():
+            with self._mutex:
+                # we use 'discard' instead of 'remove' to allow this callback
+                # to be called multiple times without raising an error
+                self._requests.discard(request_id)
+
+        return mark_request_complete
+
+    def has_in_flight_requests(self) -> bool:
+        """
+        See if there are any requests that have not been marked as completed.
+
+        >>> request_tracker = RequestTracker()
+        >>> request_tracker.has_in_flight_requests()
+        False
+        >>> mark_request_complete = request_tracker.new_request()
+        >>> request_tracker.has_in_flight_requests()
+        True
+        >>> mark_request_complete()
+        >>> request_tracker.has_in_flight_requests()
+        False
+        """
+        with self._mutex:
+            return bool(self._requests)

bugsnag/sessiontracker.py

@@ -16,6 +16,7 @@
 from bugsnag.notifier import _NOTIFIER_INFORMATION
 from bugsnag.utils import FilterDict, SanitizingJSONEncoder
 from bugsnag.event import Event
+from bugsnag.request_tracker import RequestTracker
 
 
 __all__ = []  # type: List[str]
@@ -36,6 +37,7 @@ def __init__(self, configuration):
         self.mutex = Lock()
         self.auto_sessions = False
         self.delivery_thread = None
+        self._request_tracker = RequestTracker()
 
     def start_session(self):
         if not self.auto_sessions:
@@ -138,15 +140,25 @@ def __deliver(self, sessions: List[Dict], asynchronous=True):
             deliver = self.config.delivery.deliver_sessions
 
             if 'options' in deliver.__code__.co_varnames:
-                deliver(
-                    self.config,
-                    encoded_payload,
-                    options={'asynchronous': asynchronous}
-                )
+                try:
+                    post_delivery_callback = self._request_tracker.new_request()
+
+                    deliver(
+                        self.config,
+                        encoded_payload,
+                        options={
+                            'asynchronous': asynchronous,
+                            'post_delivery_callback': post_delivery_callback,
+                        }
+                    )
+                except Exception:
+                    # ensure the request is not still marked as pending
+                    post_delivery_callback()
+                    raise
+
             else:
                 deliver(self.config, encoded_payload)
 
-
         except Exception as e:
             self.config.logger.exception('Sending sessions failed %s', e)
 

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'),

tests/test_request_tracker.py

@@ -0,0 +1,60 @@
+from bugsnag.request_tracker import RequestTracker
+
+
+def test_a_request_can_be_tracked():
+    tracker = RequestTracker()
+    assert not tracker.has_in_flight_requests()
+
+    tracker.new_request()
+    assert tracker.has_in_flight_requests()
+
+
+def test_a_request_can_be_marked_as_complete():
+    tracker = RequestTracker()
+    assert not tracker.has_in_flight_requests()
+
+    complete_request = tracker.new_request()
+    assert tracker.has_in_flight_requests()
+
+    complete_request()
+    assert not tracker.has_in_flight_requests()
+
+
+def test_requests_can_be_marked_as_complete():
+    tracker = RequestTracker()
+
+    complete_request_1 = tracker.new_request()
+    complete_request_2 = tracker.new_request()
+    complete_request_3 = tracker.new_request()
+
+    assert tracker.has_in_flight_requests()
+
+    complete_request_1()
+    complete_request_2()
+
+    assert tracker.has_in_flight_requests()
+
+    complete_request_3()
+    assert not tracker.has_in_flight_requests()
+
+
+def test_callbacks_can_be_called_multiple_times():
+    tracker = RequestTracker()
+    assert not tracker.has_in_flight_requests()
+
+    complete_request_1 = tracker.new_request()
+    complete_request_2 = tracker.new_request()
+
+    assert tracker.has_in_flight_requests()
+
+    complete_request_1()
+    complete_request_1()
+    complete_request_1()
+
+    assert tracker.has_in_flight_requests()
+
+    complete_request_2()
+    assert not tracker.has_in_flight_requests()
+
+    complete_request_2()
+    assert not tracker.has_in_flight_requests()

tests/test_sessiontracker.py

@@ -5,7 +5,7 @@
 from bugsnag.configuration import Configuration
 from bugsnag.notifier import _NOTIFIER_INFORMATION
 from bugsnag.sessiontracker import SessionTracker
-from tests.utils import IntegrationTest
+from tests.utils import BrokenDelivery, IntegrationTest, QueueingDelivery
 from unittest.mock import Mock
 
 
@@ -192,3 +192,37 @@ def test_session_tracker_starts_delivery_when_auto_capture_is_off(self):
         self.server.wait_for_session()
 
         assert self.server.sent_session_count == 1
+
+    def test_session_tracker_tracks_in_flight_requests(self):
+        delivery = QueueingDelivery()
+        client = Client(
+            api_key='a05afff2bd2ffaf0ab0f52715bbdcffd',
+            auto_capture_sessions=False,
+            session_endpoint=self.server.sessions_url,
+            asynchronous=False,
+            delivery=delivery,
+        )
+
+        client.session_tracker.start_session()
+        client.session_tracker.send_sessions()
+
+        request_tracker = client.session_tracker._request_tracker
+        assert request_tracker.has_in_flight_requests()
+
+        delivery.flush_request_queue()
+        assert not request_tracker.has_in_flight_requests()
+
+    def test_session_tracker_tracks_in_flight_requests_failure(self):
+        client = Client(
+            api_key='a05afff2bd2ffaf0ab0f52715bbdcffd',
+            auto_capture_sessions=False,
+            session_endpoint=self.server.sessions_url,
+            asynchronous=False,
+            delivery=BrokenDelivery(),
+        )
+
+        client.session_tracker.start_session()
+        client.session_tracker.send_sessions()
+
+        request_tracker = client.session_tracker._request_tracker
+        assert not request_tracker.has_in_flight_requests()