Repository: scoutapp/scout_apm_python

Issue:
ContextVar LookupError on some requests
There's a case where the usage of [`self._ctx.get()`](https://github.com/scoutapp/scout_apm_python/blob/master/src/scout_apm/core/context.py#L136) is throwing a `LookupError`.

I'm unable to reliably reproduce the issue.

```
LookupError: <ContextVar name='__scout_trackedrequest' at 0x7f241cd17ae0>
  File "django/core/handlers/exception.py", line 47, in inner
    response = get_response(request)
  File "scout_apm/django/middleware.py", line 85, in __call__
    return self.get_response(request)
  File "contextlib.py", line 124, in __exit__
    next(self.gen)
  File "scout_apm/core/tracked_request.py", line 132, in span
    self.stop_span()
  File "scout_apm/core/tracked_request.py", line 124, in stop_span
    self.finish()
  File "scout_apm/core/tracked_request.py", line 169, in finish
    context.clear_tracked_request(self)
  File "scout_apm/core/context.py", line 136, in clear_tracked_request
    if self._ctx and self._token is not None and self._ctx.get() is instance:
```


Discussion:

Resolving patch:
diff --git a/src/scout_apm/core/context.py b/src/scout_apm/core/context.py
--- a/src/scout_apm/core/context.py
+++ b/src/scout_apm/core/context.py
@@ -20,8 +20,10 @@
 
 try:
     from contextvars import ContextVar
+
+    scout_context_var = ContextVar("__scout_trackedrequest")
 except ImportError:
-    ContextVar = None
+    scout_context_var = None
 
 
 SCOUT_REQUEST_ATTR = "__scout_trackedrequest"
@@ -118,14 +120,13 @@ def __init__(self):
             self._local = SimplifiedAsgirefLocal()
         else:
             self._local = ThreadLocal()
-        self._ctx = ContextVar(SCOUT_REQUEST_ATTR) if ContextVar else None
-        self._token = None
+        self.use_context_var = scout_context_var is not None
 
     def get_tracked_request(self):
-        if self._ctx:
-            if not self._ctx.get(None):
-                self._token = self._ctx.set(TrackedRequest())
-            return self._ctx.get()
+        if scout_context_var:
+            if not scout_context_var.get(None):
+                scout_context_var.set(TrackedRequest())
+            return scout_context_var.get()
         if not hasattr(self._local, "tracked_request"):
             self._local.tracked_request = TrackedRequest()
         return self._local.tracked_request
@@ -133,9 +134,8 @@ def get_tracked_request(self):
     def clear_tracked_request(self, instance):
         if getattr(self._local, "tracked_request", None) is instance:
             del self._local.tracked_request
-        if self._ctx and self._token is not None and self._ctx.get() is instance:
-            self._ctx = ContextVar(SCOUT_REQUEST_ATTR)
-            self._token = None
+        if scout_context_var and scout_context_var.get(None) is instance:
+            scout_context_var.set(None)
 
 
 context = LocalContext()
diff --git a/src/scout_apm/core/tracked_request.py b/src/scout_apm/core/tracked_request.py
--- a/src/scout_apm/core/tracked_request.py
+++ b/src/scout_apm/core/tracked_request.py
@@ -34,6 +34,7 @@ class TrackedRequest(object):
         "_memory_start",
         "n_plus_one_tracker",
         "hit_max",
+        "sent",
     )
 
     # Stop adding new spans at this point, to avoid exhausting memory
@@ -56,6 +57,7 @@ def __init__(self):
         self._memory_start = get_rss_in_mb()
         self.n_plus_one_tracker = NPlusOneTracker()
         self.hit_max = False
+        self.sent = False
         logger.debug("Starting request: %s", self.request_id)
 
     def __repr__(self):
@@ -139,12 +141,17 @@ def current_span(self):
 
     # Request is done, release any info we have about it.
     def finish(self):
+        from scout_apm.core.context import context
+
         logger.debug("Stopping request: %s", self.request_id)
         if self.end_time is None:
             self.end_time = dt.datetime.utcnow()
+
         if self.is_real_request:
             self.tag("mem_delta", self._get_mem_delta())
-            if not self.is_ignored():
+            if not self.is_ignored() and not self.sent:
+                self.sent = True
+                logger.debug("Sending request: %s", self.request_id)
                 batch_command = BatchCommand.from_tracked_request(self)
                 CoreAgentSocketThread.send(batch_command)
             SamplersThread.ensure_started()
@@ -160,12 +167,10 @@ def finish(self):
                 ("tags", len(self.tags)),
                 ("hit_max", self.hit_max),
                 ("is_real_request", self.is_real_request),
+                ("sent", self.sent),
             ]
         )
         logger.debug("Request %s %s", self.request_id, details)
-
-        from scout_apm.core.context import context
-
         context.clear_tracked_request(self)
 
     def _get_mem_delta(self):


Test patch:
diff --git a/tests/compat.py b/tests/compat.py
--- a/tests/compat.py
+++ b/tests/compat.py
@@ -73,4 +73,17 @@ def gzip_decompress(data):
             return f.read()
 
 
-__all__ = ["gzip_decompress", "mock", "nullcontext", "suppress", "TemporaryDirectory"]
+try:
+    from contextvars import copy_context
+except ImportError:
+    copy_context = None
+
+
+__all__ = [
+    "gzip_decompress",
+    "mock",
+    "nullcontext",
+    "suppress",
+    "TemporaryDirectory",
+    "copy_context",
+]
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -15,6 +15,7 @@
 from scout_apm.core.agent import socket as scout_apm_core_socket
 from scout_apm.core.agent.manager import CoreAgentManager
 from scout_apm.core.config import SCOUT_PYTHON_VALUES, scout_config
+from scout_apm.core.context import scout_context_var
 from scout_apm.core.error_service import ErrorServiceThread
 from scout_apm.core.tracked_request import TrackedRequest
 from tests.compat import TemporaryDirectory
@@ -122,24 +123,23 @@ def isolate_global_state():
         if SCOUT_PYTHON_VALUES:
             raise ConfigLeak("Python config changes: %r" % SCOUT_PYTHON_VALUES)
 
+        request = None
         try:
             request = TrackedRequest._thread_lookup.instance
         except AttributeError:
             pass
-        else:
-            if request is not None:
-                raise TrackedRequestLeak(
-                    "Unfinished request: "
-                    "active spans = %r, complete spans = %r, tags = %r"
-                    % (
-                        [(span.operation, span.tags) for span in request.active_spans],
-                        [
-                            (span.operation, span.tags)
-                            for span in request.complete_spans
-                        ],
-                        request.tags,
-                    )
+        if request is None and scout_context_var:
+            request = scout_context_var.get(None)
+        if request is not None:
+            raise TrackedRequestLeak(
+                "Unfinished request: "
+                "active spans = %r, complete spans = %r, tags = %r"
+                % (
+                    [(span.operation, span.tags) for span in request.active_spans],
+                    [(span.operation, span.tags) for span in request.complete_spans],
+                    request.tags,
                 )
+            )
 
 
 @pytest.fixture(autouse=True, scope="session")
diff --git a/tests/integration/instruments/test_asyncio_py36plus.py b/tests/integration/instruments/test_asyncio_py36plus.py
--- a/tests/integration/instruments/test_asyncio_py36plus.py
+++ b/tests/integration/instruments/test_asyncio_py36plus.py
@@ -12,10 +12,12 @@
 import asyncio
 import contextvars
 import functools
+import sys
 
 import pytest
 
 from scout_apm.api import BackgroundTransaction, instrument
+from scout_apm.core.tracked_request import TrackedRequest
 
 
 def create_task(coro):
@@ -145,6 +147,7 @@ async def test_task_not_awaited(tracked_requests, tracked_request):
         await create_task(coro("short"))
         long = create_task(coro("long", wait=0.5))
 
+    assert len(tracked_requests) == 1
     assert len(tracked_request.complete_spans) == 2
     span = tracked_request.complete_spans[0]
     assert span.operation == "Custom/coro"
@@ -153,13 +156,25 @@ async def test_task_not_awaited(tracked_requests, tracked_request):
 
     assert not long.done()
     await long
-    assert len(tracked_request.complete_spans) == 2
-    assert tracked_request.complete_spans[0].operation == "Custom/coro"
     assert tracked_request.complete_spans[0].tags["value"] == "short"
     assert tracked_request.complete_spans[1].operation == "Job/test"
-    assert len(tracked_requests) == 1
+
     assert tracked_requests[0].request_id == tracked_request.request_id
 
+    if sys.version_info[:2] < (3, 7):
+        # Python 3.6 utilizes the contextvars backport which functions
+        # slightly differently than the contextvars stdlib implementation
+        # in 3.7+
+        assert len(tracked_request.complete_spans) == 2
+        assert len(tracked_requests) == 1
+    else:
+        assert len(tracked_request.complete_spans) == 3
+        assert tracked_request.complete_spans[2].tags["value"] == "long"
+        # When we await long, it will call finish again which adds it
+        # to tracked_requests
+        assert len(tracked_requests) == 2
+        assert tracked_requests[0] is tracked_requests[1]
+
 
 @pytest.mark.asyncio
 async def test_task_cancelled(tracked_requests, tracked_request):
@@ -212,6 +227,8 @@ async def test_future_not_gathered(tracked_requests, tracked_request):
         gather_future = asyncio.gather(
             resolving_future(0.5),
             coro(wait=0.5),
+            coro(wait=0.5),
+            coro(wait=0.5),
         )
 
     assert len(tracked_request.complete_spans) == 1
@@ -219,9 +236,18 @@ async def test_future_not_gathered(tracked_requests, tracked_request):
 
     await gather_future
 
-    assert len(tracked_request.complete_spans) == 1
-    assert len(tracked_requests) == 1
-    assert tracked_requests[0].request_id == tracked_request.request_id
+    if sys.version_info[:2] < (3, 7):
+        # Python 3.6 utilizes the contextvars backport which functions
+        # slightly differently than the contextvars stdlib implementation
+        # in 3.7+
+        assert len(tracked_request.complete_spans) == 1
+        assert len(tracked_requests) == 1
+    else:
+        assert len(tracked_request.complete_spans) == 5
+        # When the last awaitable completes, it will call finish again which
+        # adds it to tracked_requests
+        assert len(tracked_requests) == 2
+        assert tracked_requests[0].request_id == tracked_request.request_id
 
 
 @pytest.mark.asyncio
@@ -257,6 +283,54 @@ def inner_func():
     assert len(tracked_requests[0].complete_spans) == 1
 
 
+@pytest.mark.asyncio
+async def test_run_in_executor_with_context_copied_and_active_span(
+    tracked_requests, tracked_request
+):
+    # Start an outer span to prevent the tracked_request from finishing.
+    tracked_request.start_span("outer")
+
+    def inner_func():
+        with BackgroundTransaction("test"):
+            pass
+
+    # Run inner_func
+    loop = asyncio.get_event_loop()
+    child = functools.partial(inner_func)
+    context = contextvars.copy_context()
+    func = context.run
+    args = (child,)
+    await loop.run_in_executor(None, func, *args)
+    # Verify finish wasn't called
+    assert len(tracked_requests) == 0
+    assert len(tracked_request.complete_spans) == 1
+
+    tracked_request.stop_span()
+
+    assert tracked_requests[0].request_id == tracked_request.request_id
+    assert len(tracked_requests[0].complete_spans) == 2
+
+
+@pytest.mark.asyncio
+async def test_run_in_executor_with_context_copied_escalates_ignored(tracked_request):
+    # Start an outer span to prevent the tracked_request from finishing.
+    tracked_request.start_span("outer")
+
+    def inner_func():
+        with BackgroundTransaction("test"):
+            TrackedRequest.instance().tag("ignore_transaction", True)
+
+    # Run inner_func
+    loop = asyncio.get_event_loop()
+    child = functools.partial(inner_func)
+    context = contextvars.copy_context()
+    func = context.run
+    args = (child,)
+    await loop.run_in_executor(None, func, *args)
+    tracked_request.stop_span()
+    assert tracked_request.is_ignored()
+
+
 @pytest.mark.asyncio
 async def test_run_in_executor_without_context_copied(
     tracked_requests, tracked_request
@@ -272,3 +346,50 @@ def inner_func():
     assert len(tracked_requests) == 1
     assert tracked_requests[0].request_id != tracked_request.request_id
     assert len(tracked_requests[0].complete_spans) == 1
+
+
+@pytest.mark.asyncio
+async def test_run_in_executor_without_context_copied_and_active_span(
+    tracked_requests, tracked_request
+):
+    # Start an outer span to prevent the tracked_request from finishing.
+    tracked_request.start_span("outer")
+
+    def inner_func():
+        with BackgroundTransaction("test"):
+            pass
+
+    # Run inner_func
+    loop = asyncio.get_event_loop()
+    func = functools.partial(inner_func)
+    await loop.run_in_executor(None, func)
+    # Verify finish wasn't called
+    assert len(tracked_requests) == 1
+    assert len(tracked_requests[0].complete_spans) == 1
+    assert len(tracked_request.complete_spans) == 0
+
+    tracked_request.stop_span()
+
+    assert tracked_requests[0].request_id != tracked_request.request_id
+    assert len(tracked_request.complete_spans) == 1
+
+
+@pytest.mark.asyncio
+async def test_run_in_executor_without_context_copied_escalates_ignored(
+    tracked_request,
+):
+    # Start an outer span to prevent the tracked_request from finishing.
+    tracked_request.start_span("outer")
+
+    def inner_func():
+        with BackgroundTransaction("test"):
+            TrackedRequest.instance().tag("ignore_transaction", True)
+
+    # Run inner_func
+    loop = asyncio.get_event_loop()
+    func = functools.partial(inner_func)
+    await loop.run_in_executor(None, func)
+
+    tracked_request.stop_span()
+
+    assert not tracked_request.is_ignored()
diff --git a/tests/tools.py b/tests/tools.py
--- a/tests/tools.py
+++ b/tests/tools.py
@@ -4,6 +4,11 @@
 import sys
 from contextlib import contextmanager
 
+try:
+    from contextvars import ContextVar
+except ImportError:
+    ContextVar = None
+
 import pytest
 
 from scout_apm.core import objtrace
@@ -24,6 +29,10 @@
     objtrace.is_extension, reason="Requires no objtrace C extension"
 )
 
+skip_if_missing_context_vars = pytest.mark.skipif(
+    ContextVar is None, reason="Requires ContextVars"
+)
+
 
 @contextmanager
 def delete_attributes(obj, *attributes):
diff --git a/tests/unit/core/test_tracked_request.py b/tests/unit/core/test_tracked_request.py
--- a/tests/unit/core/test_tracked_request.py
+++ b/tests/unit/core/test_tracked_request.py
@@ -7,8 +7,12 @@
 
 from scout_apm.core import objtrace
 from scout_apm.core.tracked_request import TrackedRequest
-from tests.compat import mock
-from tests.tools import skip_if_objtrace_is_extension, skip_if_objtrace_not_extension
+from tests.compat import copy_context, mock
+from tests.tools import (
+    skip_if_missing_context_vars,
+    skip_if_objtrace_is_extension,
+    skip_if_objtrace_not_extension,
+)
 
 
 def test_tracked_request_repr(tracked_request):
@@ -25,6 +29,17 @@ def test_tracked_request_instance_is_a_singleton():
         tracked_request_2.finish()
 
 
+@skip_if_missing_context_vars
+def test_tracked_request_copied_into_new_context(tracked_request):
+    ctx = copy_context()
+
+    def generate_new_tracked_request(existing_tracked_request):
+        new_tracked_request = TrackedRequest.instance()
+        assert existing_tracked_request is new_tracked_request
+
+    ctx.run(generate_new_tracked_request, tracked_request)
+
+
 def test_is_real_request_default_false(tracked_request):
     assert not tracked_request.is_real_request
 
@@ -186,10 +201,9 @@ def test_finish_does_captures_memory_on_real_requests(tracked_request):
 
 
 def test_finish_log_request_info(tracked_request, caplog):
+    tracked_request.is_real_request = True
     tracked_request.start_span(operation="Something")
     tracked_request.stop_span()
-    tracked_request.is_real_request = True
-    tracked_request.finish()
 
     assert (
         "scout_apm.core.tracked_request",
@@ -197,6 +211,12 @@ def test_finish_log_request_info(tracked_request, caplog):
         "Stopping request: {}".format(tracked_request.request_id),
     ) in caplog.record_tuples
 
+    assert (
+        "scout_apm.core.tracked_request",
+        logging.DEBUG,
+        "Sending request: {}".format(tracked_request.request_id),
+    ) in caplog.record_tuples
+
     assert (
         "scout_apm.core.tracked_request",
         logging.DEBUG,
@@ -211,11 +231,29 @@ def test_finish_log_request_info(tracked_request, caplog):
             + "complete_spans=1 "
             + "tags=1 "
             + "hit_max=False "
-            + "is_real_request=True"
+            + "is_real_request=True "
+            + "sent=True"
         ),
     ) in caplog.record_tuples
 
 
+def test_finish_clears_context():
+    tracked_request_1 = TrackedRequest.instance()
+    tracked_request_1.is_real_request = True
+    tracked_request_2 = TrackedRequest.instance()
+    try:
+        assert tracked_request_2 is tracked_request_1
+    finally:
+        tracked_request_1.finish()
+        tracked_request_2.finish()
+
+    tracked_request_3 = TrackedRequest.instance()
+    try:
+        assert tracked_request_3 is not tracked_request_2
+    finally:
+        tracked_request_3.finish()
+
+
 def test_is_ignored_default_false(tracked_request):
     assert not tracked_request.is_ignored()
 
@@ -223,3 +261,54 @@ def test_is_ignored_default_false(tracked_request):
 def test_is_ignored_mark_true(tracked_request):
     tracked_request.tag("ignore_transaction", True)
     assert tracked_request.is_ignored()
+
+
+def test_request_only_sent_once(tracked_request, caplog):
+    tracked_request.is_real_request = True
+    tracked_request.start_span(operation="Something")
+    tracked_request.stop_span()
+
+    assert tracked_request.sent
+
+    sent_log = (
+        "scout_apm.core.tracked_request",
+        logging.DEBUG,
+        "Sending request: {}".format(tracked_request.request_id),
+    )
+    info_log = (
+        "scout_apm.core.tracked_request",
+        logging.DEBUG,
+        (
+            "Request {} ".format(tracked_request.request_id)
+            + "start_time={} ".format(tracked_request.start_time)
+            + "end_time={} ".format(tracked_request.end_time)
+            + "duration={} ".format(
+                (tracked_request.end_time - tracked_request.start_time).total_seconds()
+            )
+            + "active_spans=0 "
+            + "complete_spans=1 "
+            + "tags=1 "
+            + "hit_max=False "
+            + "is_real_request=True "
+            + "sent=True"
+        ),
+    )
+    assert (
+        len([log_tuple for log_tuple in caplog.record_tuples if log_tuple == sent_log])
+        == 1
+    )
+    assert (
+        len([log_tuple for log_tuple in caplog.record_tuples if log_tuple == info_log])
+        == 1
+    )
+
+    # Call finish again.
+    tracked_request.finish()
+    assert (
+        len([log_tuple for log_tuple in caplog.record_tuples if log_tuple == sent_log])
+        == 1
+    )
+    assert (
+        len([log_tuple for log_tuple in caplog.record_tuples if log_tuple == info_log])
+        == 2
+    )