diff --git a/lib/datadog/core/workers/polling.rb b/lib/datadog/core/workers/polling.rb index 4433ead51ec..fd516f16f79 100644 --- a/lib/datadog/core/workers/polling.rb +++ b/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 diff --git a/lib/datadog/tracing/workers/trace_writer.rb b/lib/datadog/tracing/workers/trace_writer.rb index fc25ab14885..4ece7038e00 100644 --- a/lib/datadog/tracing/workers/trace_writer.rb +++ b/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 diff --git a/spec/datadog/core/workers/polling_spec.rb b/spec/datadog/core/workers/polling_spec.rb index 072bb07bcc7..e4059d95eeb 100644 --- a/spec/datadog/core/workers/polling_spec.rb +++ b/spec/datadog/core/workers/polling_spec.rb @@ -47,7 +47,7 @@ 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 @@ -55,7 +55,7 @@ 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 diff --git a/spec/datadog/tracing/workers/trace_writer_spec.rb b/spec/datadog/tracing/workers/trace_writer_spec.rb index 1f41db4a297..60b33b4b474 100644 --- a/spec/datadog/tracing/workers/trace_writer_spec.rb +++ b/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