Repository: scoutapp/scout_apm_python (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.

Pull request description:
Flush transactions at shutdown
Fixes #473.

Patch:
diff --git a/.travis.yml b/.travis.yml
index 517b984..2c733a2 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -40,9 +40,9 @@ jobs:
       env:
         - CIBW_SKIP="cp33-*"
         - TWINE_NON_INTERACTIVE=1
-        - TWINE_USERNAME=scout
+        - TWINE_USERNAME=__token__
         # TWINE_PASSWORD:
-        - secure: "ohvspT59/b71SdnIFf/HA5552tGaYWLYTUV2IQ/ZNxrTiJd/hcm5eBbfAkOUiC6l2cxdHt+cAyRh27us32z7F29qzlSzrWkTD18pvFKjSGtaa2X0aBTiollfiGgiXpdT5wiDoyjyGWti3q4CpwOZjrpB/rZx0DUxdRejf0cQfs/6W8xqei/PeSTAY+58eRxKUo9JOsEObDgpjJFSWTHVip6Rm0C8zG5fEal6OEMozP+LncSzxIPZiaA0x82nYSMSBQiJXqgX6x1iNDUVl9W1wwghPjCxhz1n6kUtoKOnLcz2VhtmzR+rFDXcdgIr3l9HsSaj8t/jRIPwyUuB1DDamKWuEC0dFrqeTraOKz/t7eZkSpKa+jRaNUC4RLPSSNJVC3hEH8WlNIAyCGyvYv0RAGEYzrWcl0QZQgn1s205uIVKF8v3bepAV4oPUOEjhAk3+kVhBDLrUHWnDtLzQxPJOcC/PpMIkXx5+ZCADJrkmppEAtv+gAbHIE5gcodvyGgUc1ynp/Ul2LQxUX7KUopZBkBtFomRJWOXli5rV9tc8It+m91Xua6+jRB8UePyHw2KLGoABulNNoGSvr+uRdH8Ph/4des33iikTEnEDBk740NAsMW1pExMMlXAdaYScpBjCu7z3MrSCs4oraCVJ1g9nqbc7ugxVqMPC97fu7XKpe8="
+        - secure: "M4qjfKwZzEEGFLnOOuxcWzL7QibJmbg28uHw0puN2UzwUHBMLmgIMHcBVExMF8GmzFy3Wd1CFPukBRSvgAnxI8q42SlGXczs6GzAEPuNs9f+DkUEq2QC4msPTTyEyPVxHPwSJAvbg9kBOR+ELtQnaZg+xOlS4RAsqCqOS6bNMFMLf7jop44r7w441ZCaOLs4doI9T1hqq3XWRJ9YyP3KEPy3kMs343r9oURE+z7lqzJ3OYXLCml/ZdsHjDQI3kQja3q3PIotJK3brUMXBF+p9ECrVQwtolU2TNh4rO93m2O9EvMacEunlXoW6JB9/tXi/9K0Gee5KXtOfGu4I11GbWlj7FcExxLy+axI3i1r22mF1eer7aalid+6xFP7/UTJivvLaYpxPyMGBF7muUnqflhyQA1UpWc0z7jmY+IjjvBwJ/HGGkK1Du6xP9tP0Ksw97n7Fqo1XHNIm2IevJl2Jqp5hgl/FjcpJRMhYerITcbiYnj9MwlwvRYL6bvRIca29PoC7CwBXGzc6xBQwwnsq2RWbUnENmvjjJFvYX+nteuPUEeDHhQApPngXG4Nqf+dszouNlyIcDl4r1tRnLjBernDealGhmbR6fpW+eBt17LA3WKUKqRnTyeQN7ckIW+nOCCOsizH59dooOTtFu+P26KnmokpP9M3YcA8c9rk1tA="
       script:
         - python setup.py sdist
         - pip install --upgrade cibuildwheel twine
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 38e604c..1799677 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,11 @@
   ([PR #460](https://github.com/scoutapp/scout_apm_python/pull/460)).
 - Limit size of recorded tracebacks to reduce memory usage
   ([PR #467](https://github.com/scoutapp/scout_apm_python/pull/467)).
+- Scout now flushes any queued commands to the core agent at shutdown, with a
+  configurable timeout `shutdown_timeout_seconds`, defaulting to 2 seconds.
+  This ensures that all traces are recorded, rather than lost, especially
+  useful for single-run background commands.
+  ([Issue # 473](https://github.com/scoutapp/scout_apm_python/issues/473)).
 
 ### Fixed
 
diff --git a/src/scout_apm/bottle.py b/src/scout_apm/bottle.py
index 44e73b8..e81a18f 100644
--- a/src/scout_apm/bottle.py
+++ b/src/scout_apm/bottle.py
@@ -22,10 +22,12 @@ class ScoutPlugin(object):
 
     def set_config_from_bottle(self, app):
         bottle_configs = {}
-        for key in scout_config.known_keys():
-            value = app.config.get("scout.{}".format(key))
-            if value is not None and value != "":
-                bottle_configs[key] = value
+        prefix = "scout."
+        prefix_len = len(prefix)
+        for key, value in app.config.items():
+            if key.startswith(prefix) and len(key) > prefix_len:
+                scout_key = key[prefix_len:]
+                bottle_configs[scout_key] = value
         scout_config.set(**bottle_configs)
 
     def setup(self, app):
diff --git a/src/scout_apm/core/__init__.py b/src/scout_apm/core/__init__.py
index 8f7c17e..656abe2 100644
--- 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 @@ logger = logging.getLogger(__name__)
 
 @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
index e9c82fe..1e3832c 100644
--- a/src/scout_apm/core/config.py
+++ b/src/scout_apm/core/config.py
@@ -73,6 +73,7 @@ class ScoutConfig(object):
             "name",
             "revision_sha",
             "scm_subdirectory",
+            "shutdown_timeout_seconds",
             "socket_path",
         ]
 
@@ -220,6 +221,7 @@ class Defaults(object):
             "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
index bf424b5..89e9515 100644
--- a/src/scout_apm/core/socket.py
+++ b/src/scout_apm/core/socket.py
@@ -44,6 +44,26 @@ class CoreAgentSocketThread(SingletonThread):
 
         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
index 2ecb225..3b5f1e7 100644
--- 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
index 21bbb94..afbe3c1 100644
--- 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/integration/test_bottle.py b/tests/integration/test_bottle.py
index 5189f7c..f87583e 100644
--- a/tests/integration/test_bottle.py
+++ b/tests/integration/test_bottle.py
@@ -10,6 +10,7 @@ from webtest import TestApp
 from scout_apm.api import Config
 from scout_apm.bottle import ScoutPlugin
 from scout_apm.compat import datetime_to_timestamp, kwargs_only
+from scout_apm.core.config import scout_config
 from tests.integration.util import (
     parametrize_filtered_params,
     parametrize_queue_time_header_name,
@@ -67,6 +68,20 @@ def app_with_scout(config=None, catchall=False):
         Config.reset_all()
 
 
+def test_config_copied():
+    with app_with_scout(config={"scout.foo": "bar"}):
+        value = scout_config.value("foo")
+
+    assert value == "bar"
+
+
+def test_empty_config_not_copied():
+    with app_with_scout(config={"scout.": "foo"}):
+        value = scout_config.value("")
+
+    assert value is None
+
+
 def test_home(tracked_requests):
     with app_with_scout() as app:
         response = TestApp(app).get("/")
diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py
index 3cdda80..f1fff08 100644
--- a/tests/unit/test_core.py
+++ b/tests/unit/test_core.py
@@ -5,8 +5,10 @@ import logging
 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