Track in-flight session requests

Author: imjoehaines

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