Allows trace buffer to be closed. Makes sure it is closed before worker is shut down. Adds tests for these changes

Author: Rachel Lo

lib/ddtrace/buffer.rb

@@ -10,11 +10,13 @@ def initialize(max_size)
 
       @mutex = Mutex.new()
       @traces = []
+      @closed = false
     end
 
     # Add a new ``trace`` in the local queue. This method doesn't block the execution
     # even if the buffer is full. In that case, a random trace is discarded.
     def push(trace)
+      return if @closed
       @mutex.synchronize do
         len = @traces.length
         if len < @max_size || @max_size <= 0
@@ -48,5 +50,17 @@ def pop
         return traces
       end
     end
+
+    def close
+      @mutex.synchronize do
+        @closed = true
+      end
+    end
+
+    def closed?
+      @mutex.synchronise do
+        return @closed
+      end
+    end
   end
 end

lib/ddtrace/tracer.rb

@@ -18,8 +18,6 @@ module Datadog
   # of these function calls and sub-requests would be encapsulated within a single trace.
   # rubocop:disable Metrics/ClassLength
   class Tracer
-    DEFAULT_TIMEOUT = 5
-
     attr_reader :writer, :sampler, :services, :tags, :provider
     attr_accessor :enabled
     attr_writer :default_service
@@ -59,8 +57,8 @@ def self.debug_logging
       log.level == Logger::DEBUG
     end
 
-    def self.shutdown!
-      return if !@enabled? || @writer.worker.nil?
+    def shutdown!
+      return if !@enabled || @writer.worker.nil?
       @writer.worker.shutdown!
     end
 

lib/ddtrace/workers.rb

@@ -9,6 +9,8 @@ module Workers
     # will perform a task at regular intervals. The thread can be stopped
     # with the +stop()+ method and can start with the +start()+ method.
     class AsyncTransport
+      DEFAULT_TIMEOUT = 5
+
       attr_reader :trace_buffer, :service_buffer, :shutting_down
 
       def initialize(span_interval, service_interval, transport, buff_size, trace_task, service_task)
@@ -91,8 +93,10 @@ def stop
       end
 
       def shutdown!
-        return if @shutting_down
+        return if @shutting_down || (@trace_buffer.empty? && @service_buffer.empty?)
         @shutting_down = true
+        @trace_buffer.close
+        @service_buffer.close
         sleep(0.1)
         timeout_time = Time.now + DEFAULT_TIMEOUT
         while (!@trace_buffer.empty? || !@service_buffer.empty?) && Time.now <= timeout_time

test/buffer_test.rb

@@ -69,4 +69,20 @@ def test_trace_buffer_pop
     assert_includes(output_traces, input_traces[0])
     assert_includes(output_traces, input_traces[1])
   end
+
+  def test_closed_trace_buffer
+    # the trace buffer should not accept anymore traces when closed
+    buffer = Datadog::TraceBuffer.new(4)
+    buffer.push(1)
+    buffer.push(2)
+    buffer.push(3)
+    buffer.push(4)
+    buffer.close
+    buffer.push(5)
+    buffer.push(6)
+    out = buffer.pop
+    assert_equal(out.length, 4)
+    assert(!out.include?(5))
+    assert(!out.include?(6))
+  end
 end

test/integration_test.rb

@@ -66,6 +66,49 @@ def agent_receives_span_step3(tracer, previous_success)
     assert_equal(0, stats[:transport][:internal_error])
   end
 
+  def agent_receives_short_span(tracer)
+    tracer.set_service_info('my.service', 'rails', 'web')
+    span = tracer.start_span('my.short.op')
+    span.service = 'my.service'
+    span.finish()
+
+    tracer.shutdown!
+
+    stats = tracer.writer.stats
+    assert(span.finished?, 'span did not finish')
+    assert_equal(1, stats[:traces_flushed], 'wrong number of traces flushed')
+    assert_equal(1, stats[:services_flushed], 'wrong number of services flushed')
+    assert_equal(0, stats[:transport][:client_error])
+    assert_equal(0, stats[:transport][:server_error])
+    assert_equal(0, stats[:transport][:internal_error])
+  end
+
+  def shutdown_exec_only_once(tracer)
+    tracer.set_service_info('my.service', 'rails', 'web')
+    span = tracer.start_span('my.short.op')
+    span.service = 'my.service'
+    span.finish()
+
+    first_point = Time.now.utc
+    tracer.shutdown!
+    second_point  = Time.now.utc
+    tracer.shutdown!
+    third_point = Time.now.utc
+
+    first_shutdown = second_point - first_point
+    second_shutdown = third_point - second_point
+
+    stats = tracer.writer.stats
+    assert(first_shutdown >= 0.1, 'should have executed shutdown')
+    assert_equal(true, second_shutdown < 0.1, 'should not have executed second shutdown')
+    assert_equal(true, span.finished?, 'span did not finish')
+    assert_equal(1, stats[:traces_flushed], 'wrong number of traces flushed')
+    assert_equal(1, stats[:services_flushed], 'wrong number of services flushed')
+    assert_equal(0, stats[:transport][:client_error])
+    assert_equal(0, stats[:transport][:server_error])
+    assert_equal(0, stats[:transport][:internal_error])
+  end
+
   def test_agent_receives_span
     # test that the agent really receives the spans
     # this test can be long since it waits internal buffers flush
@@ -78,4 +121,22 @@ def test_agent_receives_span
     success = agent_receives_span_step2(tracer)
     agent_receives_span_step3(tracer, success)
   end
+
+  def test_short_span
+    skip unless ENV['TEST_DATADOG_INTEGRATION'] # requires a running agent
+
+    tracer = Datadog::Tracer.new
+    tracer.configure(enabled: true, hostname: '127.0.0.1', port: '8126')
+
+    agent_receives_short_span(tracer)
+  end
+
+  def test_shutdown_exec_once
+    skip unless ENV['TEST_DATADOG_INTEGRATION'] # requires a running agent
+
+    tracer = Datadog::Tracer.new
+    tracer.configure(enabled: true, hostname: '127.0.0.1', port: '8126')
+
+    shutdown_exec_only_once(tracer)
+  end
 end