Repository: scoutapp/scout_apm_python

Issue:
Add a "flush background thread" function
To support one-off cron tasks that exit python when finished, we need to allow the background thread to finish its work. Currently, if the main thread exits, any in-flight data the background thread has is lost.

We should add a `flush_sync` or similar function that blocks until the background thread has sent whatever payloads it has over to the core agent.


Discussion:
I guess we should by default run this [atexit](https://docs.python.org/3/library/atexit.html?highlight=atexit#module-atexit) - or at least some best-attempt for N seconds to flush out the queue.
It's ok if it's an opt-in call. But yes it needs a timeout.

Resolving patch:
diff --git a/src/scout_apm/core/__init__.py b/src/scout_apm/core/__init__.py
--- a/src/scout_apm/core/__init__.py
+++ b/src/scout_apm/core/__init__.py
@@ -1,8 +1,10 @@
 # coding=utf-8
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+import atexit
 import logging
 import os
+import sys
 
 from scout_apm import instruments
 from scout_apm.compat import kwargs_only
@@ -17,6 +19,7 @@
 
 @kwargs_only
 def install(config=None):
+    global shutdown_registered
     if config is not None:
         scout_config.set(**config)
     scout_config.log()
@@ -45,4 +48,32 @@ def install(config=None):
         # Stop the thread to avoid running threads pre-fork
         CoreAgentSocketThread.ensure_stopped()
 
+    if scout_config.value("shutdown_timeout_seconds") > 0.0 and not shutdown_registered:
+        atexit.register(shutdown)
+        shutdown_registered = True
+
     return True
+
+
+shutdown_registered = False
+
+
+def shutdown():
+    timeout_seconds = scout_config.value("shutdown_timeout_seconds")
+
+    def callback(queue_size):
+        print(  # noqa: T001
+            (
+                "Scout draining {queue_size} event{s} for up to"
+                + " {timeout_seconds} seconds"
+            ).format(
+                queue_size=queue_size,
+                s=("" if queue_size == 1 else "s"),
+                timeout_seconds=timeout_seconds,
+            ),
+            file=sys.stderr,
+        )
+
+    CoreAgentSocketThread.wait_until_drained(
+        timeout_seconds=timeout_seconds, callback=callback
+    )
diff --git a/src/scout_apm/core/config.py b/src/scout_apm/core/config.py
--- a/src/scout_apm/core/config.py
+++ b/src/scout_apm/core/config.py
@@ -73,6 +73,7 @@ def known_keys(self):
             "name",
             "revision_sha",
             "scm_subdirectory",
+            "shutdown_timeout_seconds",
             "socket_path",
         ]
 
@@ -220,6 +221,7 @@ def __init__(self):
             "name": "",
             "revision_sha": self._git_revision_sha(),
             "scm_subdirectory": "",
+            "shutdown_timeout_seconds": 2,
             "uri_reporting": "filtered_params",
         }
 
diff --git a/src/scout_apm/core/socket.py b/src/scout_apm/core/socket.py
--- a/src/scout_apm/core/socket.py
+++ b/src/scout_apm/core/socket.py
@@ -44,6 +44,26 @@ def send(cls, command):
 
         cls.ensure_started()
 
+    @classmethod
+    def wait_until_drained(cls, timeout_seconds=2.0, callback=None):
+        interval_seconds = min(timeout_seconds, 0.05)
+        start = time.time()
+        while True:
+            queue_size = cls._command_queue.qsize()
+            queue_empty = queue_size == 0
+            elapsed = time.time() - start
+            if queue_empty or elapsed >= timeout_seconds:
+                break
+
+            if callback is not None:
+                callback(queue_size)
+                callback = None
+
+            cls.ensure_started()
+
+            time.sleep(interval_seconds)
+        return queue_empty
+
     def run(self):
         self.socket_path = scout_config.value("socket_path")
         self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -159,9 +159,13 @@ def short_timeouts():
 
 
 @pytest.fixture(autouse=True)
-def auto_stop_core_agent_socket():
+def stop_and_empty_core_agent_socket():
     yield
     CoreAgentSocketThread.ensure_stopped()
+    command_queue = CoreAgentSocketThread._command_queue
+    while not command_queue.empty():
+        command_queue.get()
+        command_queue.task_done()
 
 
 @pytest.fixture
diff --git a/tests/integration/core/test_socket.py b/tests/integration/core/test_socket.py
--- a/tests/integration/core/test_socket.py
+++ b/tests/integration/core/test_socket.py
@@ -55,3 +55,30 @@ def test_send_serialization_error(socket):
 def test_send_network_error(sendall, socket):
     sendall.side_effect = OSError
     CoreAgentSocketThread.send(Command())
+
+
+def test_wait_until_drained_empty(socket):
+    CoreAgentSocketThread.ensure_stopped()
+
+    empty = CoreAgentSocketThread.wait_until_drained()
+    assert empty
+
+
+def test_wait_until_drained_one_item(socket):
+    CoreAgentSocketThread._command_queue.put(Command(), False)
+
+    empty = CoreAgentSocketThread.wait_until_drained(timeout_seconds=0.1)
+    assert empty
+
+
+def test_wait_until_drained_one_slow(socket):
+    class SlowCommand(object):
+        def message(self):
+            time.sleep(0.05)
+            return {}
+
+    for _ in range(10):
+        CoreAgentSocketThread.send(SlowCommand())
+
+    empty = CoreAgentSocketThread.wait_until_drained(timeout_seconds=0.05)
+    assert not empty
diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py
--- a/tests/unit/test_core.py
+++ b/tests/unit/test_core.py
@@ -5,8 +5,10 @@
 import os
 
 from scout_apm.api import Config
-from scout_apm.core import install
+from scout_apm.core import install, shutdown
+from scout_apm.core.config import scout_config
 from scout_apm.core.core_agent_manager import CoreAgentManager
+from scout_apm.core.metadata import report_app_metadata
 from tests.compat import mock
 
 
@@ -52,3 +54,16 @@ def test_install_success(caplog):
         logging.DEBUG,
         "APM Launching on PID: %s" % os.getpid(),
     ) in caplog.record_tuples
+
+
+def test_shutdown(capsys):
+    report_app_metadata()  # queued but thread not running
+    try:
+        scout_config.set(shutdown_timeout_seconds=0.1)
+
+        shutdown()
+    finally:
+        Config.reset_all()
+
+    captured = capsys.readouterr()
+    assert "Scout draining" in captured.err