feat: register atexit handler (getsentry#15)

untitaker · web-flow · commit 7f25a3988983 · 2018-07-31T13:32:15.000+02:00
* feat: register atexit handler

* chore: add tests, stylefixes

* fix: tests under py3
diff --git a/sentry_sdk/client.py b/sentry_sdk/client.py
@@ -1,6 +1,7 @@
 import os
 import uuid
 import random
+import atexit
 
 from .utils import Dsn, SkipEvent, ContextVar
 from .transport import Transport
@@ -52,6 +53,8 @@ def __init__(self, dsn=None, *args, **kwargs):
         for integration in integrations:
             integration(self)
 
+        atexit.register(self.close)
+
     @property
     def dsn(self):
         """The DSN that created this event."""
diff --git a/sentry_sdk/transport.py b/sentry_sdk/transport.py
@@ -1,10 +1,14 @@
+from __future__ import print_function
+
 import json
-import time
 import zlib
 import urllib3
 import logging
 import threading
 import certifi
+import sys
+import traceback
+
 from datetime import datetime, timedelta
 
 from ._compat import queue
@@ -64,22 +68,23 @@ def thread():
         disabled_until = None
         while 1:
             item = transport._queue.get()
+            if item is _SHUTDOWN:
+                transport._queue.task_done()
+                break
+
             if disabled_until is not None:
-                if disabled_until < datetime.utcnow():
+                if datetime.utcnow() < disabled_until:
+                    transport._queue.task_done()
                     continue
                 disabled_until = None
 
-            if item is None:
-                if transport._done:
-                    break
-                continue
-            elif item is _SHUTDOWN:
-                break
-
             try:
                 disabled_until = send_event(transport._pool, item, auth)
+                transport._queue.task_done()
             except Exception:
-                logger.exception("Could not send sentry event")
+                print("Could not send sentry event", file=sys.stderr)
+                print(traceback.format_exc(), file=sys.stderr)
+                transport._queue.task_done()
                 continue
 
     t = threading.Thread(target=thread)
@@ -91,7 +96,6 @@ class Transport(object):
     def __init__(self, dsn, http_proxy=None, https_proxy=None):
         self.dsn = dsn
         self._queue = None
-        self._done = False
         self._pool = _make_pool(dsn, http_proxy=http_proxy, https_proxy=https_proxy)
 
     def start(self):
@@ -108,7 +112,6 @@ def capture_event(self, event):
             pass
 
     def close(self):
-        self._done = True
         if self._queue is not None:
             try:
                 self._queue.put_nowait(_SHUTDOWN)
@@ -117,12 +120,12 @@ def close(self):
             self._queue = None
 
     def drain_events(self, timeout):
-        deadline = time.time() + timeout
         q = self._queue
-        while not self._done and q.qsize() > 0:
-            if time.time() >= deadline:
-                return False
-            time.sleep(0.1)
+        if q is None:
+            return True
+
+        with q.all_tasks_done:
+            q.all_tasks_done.wait(timeout)
         return True
 
     def __del__(self):
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -1,7 +1,11 @@
+import time
 import pytest
+import sys
+import subprocess
+from textwrap import dedent
 from sentry_sdk import Client
-from sentry_sdk.utils import Event, Dsn
 from sentry_sdk.transport import Transport
+from sentry_sdk.utils import Event, Dsn
 
 
 def test_transport_option(monkeypatch):
@@ -42,3 +46,34 @@ def capture_event(self, event):
     c.capture_event(e(Exception))
     c.capture_event(e(ValueError))
     pytest.raises(EventCaptured, lambda: c.capture_event(e(BaseException)))
+
+
+@pytest.mark.parametrize("num_messages", [10, 20])
+def test_atexit(tmpdir, monkeypatch, num_messages):
+    app = tmpdir.join("app.py")
+    app.write(
+        dedent(
+            """
+    import time
+    from sentry_sdk import init, transport, capture_message
+
+    def send_event(pool, event, auth):
+        time.sleep(0.1)
+        print(event["message"])
+
+    transport.send_event = send_event
+    init("http://foobar@localhost/123", drain_timeout={num_messages})
+
+    for _ in range({num_messages}):
+        capture_message("HI")
+    """.format(
+                num_messages=num_messages
+            )
+        )
+    )
+
+    start = time.time()
+    output = subprocess.check_output([sys.executable, str(app)])
+    end = time.time()
+    assert int(end - start) == num_messages / 10
+    assert output.count(b"HI") == num_messages
