make shutdown timeout configurable for Datadog::Tracing::Workers::Tra…

…ceWriter

lib/datadog/core/workers/polling.rb

@@ -6,7 +6,7 @@ module Core
     module Workers
       # Adds polling (async looping) behavior to workers
       module Polling
-        SHUTDOWN_TIMEOUT = 1
+        DEFAULT_SHUTDOWN_TIMEOUT = 1
 
         def self.included(base)
           base.include(Workers::IntervalLoop)
@@ -21,7 +21,7 @@ def perform(*args)
           end
         end
 
-        def stop(force_stop = false, timeout = SHUTDOWN_TIMEOUT)
+        def stop(force_stop = false, timeout = DEFAULT_SHUTDOWN_TIMEOUT)
           if running?
             # Attempt graceful stop and wait
             stop_loop

lib/datadog/tracing/workers/trace_writer.rb

@@ -104,6 +104,8 @@ def initialize(options = {})
           # Workers::Queue settings
           @buffer_size = options.fetch(:buffer_size, DEFAULT_BUFFER_MAX_SIZE)
           self.buffer = TraceBuffer.new(@buffer_size)
+
+          @shutdown_timeout = options.fetch(:shutdown_timeout, Core::Workers::Polling::DEFAULT_SHUTDOWN_TIMEOUT)
         end
 
         # NOTE: #perform is wrapped by other modules:
@@ -119,7 +121,7 @@ def perform(traces)
           nil
         end
 
-        def stop(*args)
+        def stop(force_stop = false, timeout = @shutdown_timeout)
           buffer.close if running?
           super
         end

spec/datadog/core/workers/polling_spec.rb

@@ -47,15 +47,15 @@
       shared_context 'graceful stop' do
         before do
           allow(worker).to receive(:join)
-            .with(described_class::SHUTDOWN_TIMEOUT)
+            .with(described_class::DEFAULT_SHUTDOWN_TIMEOUT)
             .and_return(true)
         end
       end
 
       context 'when the worker has not been started' do
         before do
           allow(worker).to receive(:join)
-            .with(described_class::SHUTDOWN_TIMEOUT)
+            .with(described_class::DEFAULT_SHUTDOWN_TIMEOUT)
             .and_return(true)
         end
 
@@ -113,6 +113,22 @@
           end
         end
       end
+
+      context 'given shutdown timeout' do
+        subject(:stop) { worker.stop(false, 1000) }
+        include_context 'graceful stop'
+
+        before do
+          expect(worker).to receive(:join)
+            .with(1000)
+            .and_return(true)
+
+          worker.perform
+          try_wait_until { worker.running? && worker.run_loop? }
+        end
+
+        it { is_expected.to be true }
+      end
     end
 
     describe '#enabled?' do

spec/datadog/tracing/workers/trace_writer_spec.rb

@@ -397,6 +397,22 @@
         end
       end
     end
+
+    context 'given shutdown_timeout' do
+      let(:options) { { shutdown_timeout: 1000 } }
+      include_context 'shuts down the worker'
+
+      context 'and the worker has been started' do
+        before do
+          expect(writer).to receive(:join).with(1000).and_return(true)
+
+          writer.perform
+          try_wait_until { writer.running? && writer.run_loop? }
+        end
+
+        it { is_expected.to be true }
+      end
+    end
   end
 
   describe '#work_pending?' do