From d0f380855717a2f4d35dcd01f6cc0b9b117cfa07 Mon Sep 17 00:00:00 2001
From: Andrey Marchenko <andrey.marchenko@datadoghq.com>
Date: Tue, 26 Sep 2023 13:51:51 +0200
Subject: [PATCH] make shutdown timeout configurable for
 Datadog::Tracing::Workers::TraceWriter

---
 lib/datadog/core/workers/polling.rb           |  4 ++--
 lib/datadog/tracing/workers/trace_writer.rb   |  4 +++-
 spec/datadog/core/workers/polling_spec.rb     | 20 +++++++++++++++++--
 .../tracing/workers/trace_writer_spec.rb      | 16 +++++++++++++++
 4 files changed, 39 insertions(+), 5 deletions(-)

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