From 0220273adfae1ad85efafc655e2f782ca9aa8535 Mon Sep 17 00:00:00 2001
From: "Christian Mauduit (DataDog)" <christian.mauduit@datadoghq.com>
Date: Wed, 15 Nov 2017 09:28:57 -0500
Subject: [PATCH] [tracer] more tests for partial flush

---
 lib/ddtrace/context.rb       |  5 ++++
 lib/ddtrace/context_flush.rb |  3 ++
 test/context_flush_test.rb   | 53 ++++++++++++++++++++++++++++++++++--
 test/context_test.rb         |  8 +++---
 4 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/lib/ddtrace/context.rb b/lib/ddtrace/context.rb
index d82d6998c6b..08e455bc647 100644
--- a/lib/ddtrace/context.rb
+++ b/lib/ddtrace/context.rb
@@ -203,7 +203,12 @@ def delete_span_if
           finished = span.finished?
           delete_span = yield span
           if delete_span
+            # We need to detach the span from the context, else, some code
+            # finishing it afterwards would mess up with the number of
+            # finished_spans and possibly cause other side effects.
             span.context = nil
+            # Acknowledge there's one span less to finish, if needed.
+            # It's very important to keep this balanced.
             @finished_spans -= 1 if finished
           end
           delete_span
diff --git a/lib/ddtrace/context_flush.rb b/lib/ddtrace/context_flush.rb
index 3bde47b5f52..9a0de4fef0a 100644
--- a/lib/ddtrace/context_flush.rb
+++ b/lib/ddtrace/context_flush.rb
@@ -107,9 +107,12 @@ def partial_flush(context)
     def each_partial_trace(context)
       start_time = context.start_time
       length = context.length
+      # Stop and do not flush anything if there are not enough spans.
       return if length < @min_spans_before_partial_flush
+      # If there are enough spans, but not too many, check for start time.
       return if length < @max_spans_before_partial_flush &&
                 start_time && start_time > Time.now.utc - @partial_flush_timeout
+      # Here, either the trace is old or we have too many spans, flush it.
       traces = partial_flush(context)
       return unless traces
       traces.each do |trace|
diff --git a/test/context_flush_test.rb b/test/context_flush_test.rb
index b877f2be80a..d414e5e5af6 100644
--- a/test/context_flush_test.rb
+++ b/test/context_flush_test.rb
@@ -4,9 +4,8 @@
 
 module Datadog
   class ContextFlush
-    attr_accessor :max_spans_per_trace_soft
-    attr_accessor :max_spans_per_trace_hard
-    attr_accessor :min_spans_for_flush_timeout
+    attr_accessor :max_spans_before_partial_flush
+    attr_accessor :min_spans_before_partial_flush
     attr_accessor :partial_flush_timeout
 
     public :partial_roots
@@ -41,6 +40,10 @@ def test_partial_roots_typical
         assert_equal([root_id, child3_id].to_set, marked_ids)
         partial_roots_spans = context_flush.partial_roots_spans(context)
         assert_includes(partial_roots_spans, child1_id)
+
+        context_flush.each_partial_trace(context) do |_t|
+          flunk('there should be no trace here, not enough spans to trigger a flush')
+        end
       end
     end
   end
@@ -55,5 +58,49 @@ def test_partial_roots_empty
     assert_nil(marked_ids)
     partial_roots_spans = context_flush.partial_roots_spans(context)
     assert_nil(partial_roots_spans)
+
+    context_flush.each_partial_trace(context) do |_t|
+      flunk('there should be no trace here, there are no spans at all')
+    end
+  end
+
+  def test_partial_flush
+    tracer = get_test_tracer
+    # Trigger early flush.
+    context_flush = Datadog::ContextFlush.new(min_spans_before_partial_flush: 1,
+                                              max_spans_before_partial_flush: 1)
+    context = tracer.call_context
+
+    action12 = Minitest::Mock.new
+    action12.expect(:call_with_names, nil, [%w[child1 child2]])
+    action3456 = Minitest::Mock.new
+    action3456.expect(:call_with_names, nil, [['child3']])
+    action3456.expect(:call_with_names, nil, [%w[child4 child5 child6]])
+
+    tracer.trace('root') do
+      tracer.trace('child1') do
+        tracer.trace('child2') do
+        end
+      end
+      tracer.trace('child3') do
+        context_flush.each_partial_trace(context) do |t|
+          action12.call_with_names(t.map(&:name))
+        end
+      end
+      tracer.trace('child4') do
+        tracer.trace('child5') do
+        end
+        tracer.trace('child6') do
+        end
+      end
+      context_flush.each_partial_trace(context) do |t|
+        action3456.call_with_names(t.map(&:name))
+      end
+    end
+
+    action12.verify
+    action3456.verify
+
+    assert_equal(0, context.length, 'everything should be written by now')
   end
 end
diff --git a/test/context_test.rb b/test/context_test.rb
index 68fa21beb64..d4e2faffaee 100644
--- a/test/context_test.rb
+++ b/test/context_test.rb
@@ -304,9 +304,9 @@ def test_each_span
     ctx.add_span(span)
 
     action = MiniTest::Mock.new
-    action.expect(:nop, nil, ['test.op'])
+    action.expect(:call_with_name, nil, ['test.op'])
     ctx.each_span do |s|
-      action.nop(s.name)
+      action.call_with_name(s.name)
     end
     action.verify
   end
@@ -316,14 +316,14 @@ def test_delete_span_if
     ctx = tracer.call_context
 
     action = MiniTest::Mock.new
-    action.expect(:nop, nil, ['test.op2'])
+    action.expect(:call_with_name, nil, ['test.op2'])
     tracer.trace('test.op1') do
       tracer.trace('test.op2') do
         assert_equal(2, ctx.length)
         ctx.delete_span_if { |span| span.name == 'test.op1' }
         assert_equal(1, ctx.length)
         ctx.each_span do |s|
-          action.nop(s.name)
+          action.call_with_name(s.name)
         end
         assert_equal(false, ctx.finished?, 'context is not finished as op2 is not finished')
         tracer.trace('test.op3') do