[core] initial support for partial flushes (#668)

* [core] partial flush prototype

* [core] partial flush when >= min spans required

* [core] add tests for partial flushes

* [core] use span._finished = True instead of span.finish() for tests

* [core] add test for partial flush with remaining opened spans

* Update tests/test_context.py

* [core] change context config to 'tracer'

* [core] move partial flush settings to class level

* [core] remove usage of config

* [tests] remove unused import

Author: brettlangdon

ddtrace/context.py

@@ -2,7 +2,7 @@
 import threading
 
 from .constants import SAMPLING_PRIORITY_KEY
-
+from .utils.formats import asbool, get_env
 
 log = logging.getLogger(__name__)
 
@@ -22,6 +22,9 @@ class Context(object):
 
     This data structure is thread-safe.
     """
+    _partial_flush_enabled = asbool(get_env('tracer', 'partial_flush_enabled', 'false'))
+    _partial_flush_min_spans = int(get_env('tracer', 'partial_flush_min_spans', 500))
+
     def __init__(self, trace_id=None, span_id=None, sampled=True, sampling_priority=None):
         """
         Initialize a new thread-safe ``Context``.
@@ -190,6 +193,31 @@ def get(self):
                 self._sampling_priority = None
                 self._sampled = True
                 return trace, sampled
+
+            elif self._partial_flush_enabled and self._finished_spans >= self._partial_flush_min_spans:
+                # partial flush when enabled and we have more than the minimal required spans
+                trace = self._trace
+                sampled = self._sampled
+                sampling_priority = self._sampling_priority
+                # attach the sampling priority to the context root span
+                if sampled and sampling_priority is not None and trace:
+                    trace[0].set_metric(SAMPLING_PRIORITY_KEY, sampling_priority)
+
+                # Any open spans will remain as `self._trace`
+                # Any finished spans will get returned to be flushed
+                opened_spans = []
+                closed_spans = []
+                for span in trace:
+                    if span._finished:
+                        closed_spans.append(span)
+                    else:
+                        opened_spans.append(span)
+
+                # Update trace spans and stats
+                self._trace = opened_spans
+                self._finished_spans = 0
+
+                return closed_spans, sampled
             else:
                 return None, None
 

tests/test_context.py

@@ -1,3 +1,4 @@
+import contextlib
 import mock
 import threading
 
@@ -15,6 +16,20 @@ class TestTracingContext(TestCase):
     Tests related to the ``Context`` class that hosts the trace for the
     current execution flow.
     """
+    @contextlib.contextmanager
+    def override_partial_flush(self, ctx, enabled, min_spans):
+        original_enabled = ctx._partial_flush_enabled
+        original_min_spans = ctx._partial_flush_min_spans
+
+        ctx._partial_flush_enabled = enabled
+        ctx._partial_flush_min_spans = min_spans
+
+        try:
+            yield
+        finally:
+            ctx._partial_flush_enabled = original_enabled
+            ctx._partial_flush_min_spans = original_min_spans
+
     def test_add_span(self):
         # it should add multiple spans
         ctx = Context()
@@ -101,6 +116,163 @@ def test_get_trace_empty(self):
         ok_(trace is None)
         ok_(sampled is None)
 
+    def test_partial_flush(self):
+        """
+        When calling `Context.get`
+            When partial flushing is enabled
+            When we have just enough finished spans to flush
+                We return the finished spans
+        """
+        tracer = get_dummy_tracer()
+        ctx = Context()
+
+        # Create a root span with 5 children, all of the children are finished, the root is not
+        root = Span(tracer=tracer, name='root')
+        ctx.add_span(root)
+        for i in range(5):
+            child = Span(tracer=tracer, name='child_{}'.format(i), trace_id=root.trace_id, parent_id=root.span_id)
+            child._parent = root
+            child._finished = True
+            ctx.add_span(child)
+            ctx.close_span(child)
+
+        with self.override_partial_flush(ctx, enabled=True, min_spans=5):
+            trace, sampled = ctx.get()
+
+        self.assertIsNotNone(trace)
+        self.assertIsNotNone(sampled)
+
+        self.assertEqual(len(trace), 5)
+        self.assertEqual(
+            set(['child_0', 'child_1', 'child_2', 'child_3', 'child_4']),
+            set([span.name for span in trace])
+        )
+
+        # Ensure we clear/reset internal stats as expected
+        self.assertEqual(ctx._finished_spans, 0)
+        self.assertEqual(ctx._trace, [root])
+        with self.override_partial_flush(ctx, enabled=True, min_spans=5):
+            trace, sampled = ctx.get()
+            self.assertIsNone(trace)
+            self.assertIsNone(sampled)
+
+    def test_partial_flush_too_many(self):
+        """
+        When calling `Context.get`
+            When partial flushing is enabled
+            When we have more than the minimum number of spans needed to flush
+                We return the finished spans
+        """
+        tracer = get_dummy_tracer()
+        ctx = Context()
+
+        # Create a root span with 5 children, all of the children are finished, the root is not
+        root = Span(tracer=tracer, name='root')
+        ctx.add_span(root)
+        for i in range(5):
+            child = Span(tracer=tracer, name='child_{}'.format(i), trace_id=root.trace_id, parent_id=root.span_id)
+            child._parent = root
+            child._finished = True
+            ctx.add_span(child)
+            ctx.close_span(child)
+
+        with self.override_partial_flush(ctx, enabled=True, min_spans=1):
+            trace, sampled = ctx.get()
+
+        self.assertIsNotNone(trace)
+        self.assertIsNotNone(sampled)
+
+        self.assertEqual(len(trace), 5)
+        self.assertEqual(
+            set(['child_0', 'child_1', 'child_2', 'child_3', 'child_4']),
+            set([span.name for span in trace])
+        )
+
+        # Ensure we clear/reset internal stats as expected
+        self.assertEqual(ctx._finished_spans, 0)
+        self.assertEqual(ctx._trace, [root])
+        with self.override_partial_flush(ctx, enabled=True, min_spans=5):
+            trace, sampled = ctx.get()
+            self.assertIsNone(trace)
+            self.assertIsNone(sampled)
+
+    def test_partial_flush_too_few(self):
+        """
+        When calling `Context.get`
+            When partial flushing is enabled
+            When we do not have enough finished spans to flush
+                We return no spans
+        """
+        tracer = get_dummy_tracer()
+        ctx = Context()
+
+        # Create a root span with 5 children, all of the children are finished, the root is not
+        root = Span(tracer=tracer, name='root')
+        ctx.add_span(root)
+        for i in range(5):
+            child = Span(tracer=tracer, name='child_{}'.format(i), trace_id=root.trace_id, parent_id=root.span_id)
+            child._parent = root
+            child._finished = True
+            ctx.add_span(child)
+            ctx.close_span(child)
+
+        # Test with having 1 too few spans for partial flush
+        with self.override_partial_flush(ctx, enabled=True, min_spans=6):
+            trace, sampled = ctx.get()
+
+        self.assertIsNone(trace)
+        self.assertIsNone(sampled)
+
+        self.assertEqual(len(ctx._trace), 6)
+        self.assertEqual(ctx._finished_spans, 5)
+        self.assertEqual(
+            set(['root', 'child_0', 'child_1', 'child_2', 'child_3', 'child_4']),
+            set([span.name for span in ctx._trace])
+        )
+
+    def test_partial_flush_remaining(self):
+        """
+        When calling `Context.get`
+            When partial flushing is enabled
+            When we have some unfinished spans
+                We keep the unfinished spans around
+        """
+        tracer = get_dummy_tracer()
+        ctx = Context()
+
+        # Create a root span with 5 children, all of the children are finished, the root is not
+        root = Span(tracer=tracer, name='root')
+        ctx.add_span(root)
+        for i in range(10):
+            child = Span(tracer=tracer, name='child_{}'.format(i), trace_id=root.trace_id, parent_id=root.span_id)
+            child._parent = root
+            ctx.add_span(child)
+
+            # CLose the first 5 only
+            if i < 5:
+                child._finished = True
+                ctx.close_span(child)
+
+        with self.override_partial_flush(ctx, enabled=True, min_spans=5):
+            trace, sampled = ctx.get()
+
+        # Assert partially flushed spans
+        self.assertTrue(len(trace), 5)
+        self.assertIsNotNone(sampled)
+        self.assertEqual(
+            set(['child_0', 'child_1', 'child_2', 'child_3', 'child_4']),
+            set([span.name for span in trace])
+        )
+
+        # Assert remaining unclosed spans
+        self.assertEqual(len(ctx._trace), 6)
+        self.assertEqual(ctx._finished_spans, 0)
+        self.assertEqual(
+            set(['root', 'child_5', 'child_6', 'child_7', 'child_8', 'child_9']),
+            set([span.name for span in ctx._trace]),
+        )
+
+
     def test_finished(self):
         # a Context is finished if all spans inside are finished
         ctx = Context()