Ensure workers can send data popped off the queue at shutdown

Previously, any data still in the queue would be handled by the main
thread. But, any data that a worker had already popped off the queue,
but not yet sent on the wire, would be lost.

This change pushes down the logic for stopping into the writer and
worker classes. The worker tells the threads to exit and waits for them
to do so.

Author: onlynone

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 For the full commit log, [see here](https://github.com/influxdata/influxdb-ruby/commits/master).
 
+## unreleased
+
+- Ensure workers can send data popped off the queue at shutdown (#239,
+  @onlynone)
+
 ## v0.8.0, released 2020-02-05
 
 - Allow dropping of specific series from specific DBs (#233, @cantino)

README.md

@@ -295,7 +295,10 @@ async_options = {
   # checking if there are new jobs in the queue
   sleep_interval:      5,
   # whether client will block if queue is full
-  block_on_full_queue: false
+  block_on_full_queue: false,
+  # Max time (in seconds) the main thread will wait for worker threads to stop
+  # on shutdown. Defaults to 2x sleep_interval.
+  shutdown_timeout: 10
 }
 
 influxdb = InfluxDB::Client.new database, async: async_options

lib/influxdb/client.rb

@@ -59,17 +59,19 @@ def initialize(database = nil, **opts)
     end
 
     def stop!
-      if config.async?
-        # If retry was infinite (-1), set it to zero to give the main thread one
-        # last chance to flush the queue
-        config.retry = 0 if config.retry < 0
-        writer.worker.stop!
+      if @writer == self
+        @stopped = true
+      else
+        @writer.stop!
       end
-      @stopped = true
     end
 
     def stopped?
-      @stopped
+      if @writer == self
+        @stopped
+      else
+        @writer.stopped?
+      end
     end
 
     def now

lib/influxdb/writer/async.rb

@@ -9,6 +9,16 @@ class Async
       def initialize(client, config)
         @client = client
         @config = config
+        @stopped = false
+      end
+
+      def stopped?
+        @stopped
+      end
+
+      def stop!
+        worker.stop!
+        @stopped = true
       end
 
       def write(data, precision = nil, retention_policy = nil, database = nil)
@@ -29,15 +39,16 @@ def worker
         end
       end
 
-      class Worker
+      class Worker # rubocop:disable Metrics/ClassLength
         attr_reader :client,
                     :queue,
                     :threads,
                     :max_post_points,
                     :max_queue_size,
                     :num_worker_threads,
                     :sleep_interval,
-                    :block_on_full_queue
+                    :block_on_full_queue,
+                    :shutdown_timeout
 
         include InfluxDB::Logging
 
@@ -47,7 +58,7 @@ class Worker
         SLEEP_INTERVAL      = 5
         BLOCK_ON_FULL_QUEUE = false
 
-        def initialize(client, config)
+        def initialize(client, config) # rubocop:disable Metrics/MethodLength
           @client = client
           config = config.is_a?(Hash) ? config : {}
 
@@ -56,10 +67,11 @@ def initialize(client, config)
           @num_worker_threads  = config.fetch(:num_worker_threads,  NUM_WORKER_THREADS)
           @sleep_interval      = config.fetch(:sleep_interval,      SLEEP_INTERVAL)
           @block_on_full_queue = config.fetch(:block_on_full_queue, BLOCK_ON_FULL_QUEUE)
+          @shutdown_timeout    = config.fetch(:shutdown_timeout,    2 * @sleep_interval)
 
           queue_class = @block_on_full_queue ? SizedQueue : InfluxDB::MaxQueue
           @queue = queue_class.new max_queue_size
-
+          @should_stop = false
           spawn_threads!
         end
 
@@ -68,11 +80,11 @@ def push(payload, precision = nil, retention_policy = nil, database = nil)
         end
 
         def current_threads
-          Thread.list.select { |t| t[:influxdb] == object_id }
+          @threads
         end
 
         def current_thread_count
-          Thread.list.count { |t| t[:influxdb] == object_id }
+          @threads.count
         end
 
         # rubocop:disable Metrics/CyclomaticComplexity
@@ -87,7 +99,7 @@ def spawn_threads!
             @threads << Thread.new do
               Thread.current[:influxdb] = object_id
 
-              until client.stopped?
+              until @should_stop
                 check_background_queue(thread_num)
                 sleep rand(sleep_interval)
               end
@@ -97,7 +109,7 @@ def spawn_threads!
           end
         end
 
-        def check_background_queue(thread_num = 0)
+        def check_background_queue(thread_num = -1)
           log(:debug) do
             "Checking background queue on thread #{thread_num} (#{current_thread_count} active)"
           end
@@ -123,10 +135,10 @@ def check_background_queue(thread_num = 0)
             return if data.values.flatten.empty?
 
             begin
-              log(:debug) { "Found data in the queue! (#{sizes(data)})" }
+              log(:debug) { "Found data in the queue! (#{sizes(data)}) on thread #{thread_num}" }
               write(data)
             rescue StandardError => e
-              log :error, "Cannot write data: #{e.inspect}"
+              log :error, "Cannot write data: #{e.inspect} on thread #{thread_num}"
             end
 
             break if queue.length > max_post_points
@@ -138,7 +150,22 @@ def check_background_queue(thread_num = 0)
         # rubocop:enable Metrics/AbcSize
 
         def stop!
-          log(:debug) { "Thread exiting, flushing queue." }
+          log(:debug) { "Worker is being stopped, flushing queue." }
+
+          # If retry was infinite (-1), set it to zero to give the threads one
+          # last chance to write their data
+          client.config.retry = 0 if client.config.retry < 0
+
+          # Signal the background threads that they should exit.
+          @should_stop = true
+
+          # Wait for the threads to exit and then kill them
+          @threads.each do |t|
+            r = t.join(shutdown_timeout)
+            t.kill if r.nil?
+          end
+
+          # Flush any remaining items in the queue on the main thread
           check_background_queue until queue.empty?
         end
 

lib/influxdb/writer/udp.rb

@@ -11,6 +11,9 @@ def initialize(client, host: "localhost".freeze, port: 4444)
         @port = port
       end
 
+      # No-op for UDP writers
+      def stop!; end
+
       def write(payload, _precision = nil, _retention_policy = nil, _database = nil)
         with_socket { |sock| sock.send(payload, 0) }
       end

spec/influxdb/cases/async_client_spec.rb

@@ -2,7 +2,7 @@
 require "timeout"
 
 describe InfluxDB::Client do
-  let(:async_options) { true }
+  let(:async_options) { { sleep_interval: 0.1 } }
   let(:client) { described_class.new(async: async_options) }
   let(:subject) { client }
   let(:stub_url) { "http://localhost:8086/write?db=&p=root&precision=s&u=root" }
@@ -18,14 +18,9 @@
         subject.write_point('a', values: { i: i })
       end
 
-      # The timout code is fragile, and heavily dependent on system load
-      # (and scheduler decisions). On the CI, the system is less
-      # responsive and needs a bit more time.
-      timeout_stretch = ENV["TRAVIS"] == "true" ? 10 : 3
+      sleep 1 until worker.threads.none? { |t| t[:influxdb].nil? }
 
-      Timeout.timeout(timeout_stretch * worker.sleep_interval) do
-        subject.stop!
-      end
+      subject.stop!
 
       worker.threads.each do |t|
         expect(t.stop?).to be true
@@ -41,6 +36,7 @@
       let(:precision)        { 'test_precision' }
       let(:retention_policy) { 'test_period' }
       let(:database)         { 'test_database' }
+      let(:async_options)    { { num_worker_threads: 1, sleep_interval: 0.1 } }
 
       it "writes aggregate payload to the client" do
         queue = Queue.new
@@ -51,9 +47,11 @@
         subject.write_point(series, { values: { t: 60 } }, precision, retention_policy, database)
         subject.write_point(series, { values: { t: 61 } }, precision, retention_policy, database)
 
-        Timeout.timeout(worker.sleep_interval) do
-          expect(queue.pop).to eq ["#{series} t=60i\n#{series} t=61i", precision, retention_policy, database]
-        end
+        sleep 1 until worker.threads.none? { |t| t[:influxdb].nil? }
+
+        subject.stop!
+
+        expect(queue.pop).to eq ["#{series} t=60i\n#{series} t=61i", precision, retention_policy, database]
       end
 
       context 'when different precision, retention_policy and database are given' do
@@ -72,36 +70,62 @@
           subject.write_point(series, { values: { t: 62 } }, precision,  retention_policy2, database)
           subject.write_point(series, { values: { t: 63 } }, precision,  retention_policy,  database2)
 
-          Timeout.timeout(worker.sleep_interval) do
-            expect(queue.pop).to eq ["#{series} t=60i", precision,  retention_policy,  database]
-            expect(queue.pop).to eq ["#{series} t=61i", precision2, retention_policy,  database]
-            expect(queue.pop).to eq ["#{series} t=62i", precision,  retention_policy2, database]
-            expect(queue.pop).to eq ["#{series} t=63i", precision,  retention_policy,  database2]
-          end
+          sleep 1 until worker.threads.none? { |t| t[:influxdb].nil? }
+
+          subject.stop!
+
+          expect(queue.pop).to eq ["#{series} t=60i", precision,  retention_policy,  database]
+          expect(queue.pop).to eq ["#{series} t=61i", precision2, retention_policy,  database]
+          expect(queue.pop).to eq ["#{series} t=62i", precision,  retention_policy2, database]
+          expect(queue.pop).to eq ["#{series} t=63i", precision,  retention_policy,  database2]
         end
       end
     end
   end
 
   describe "async options" do
-    let(:async_options) do
-      {
-        max_post_points:     10,
-        max_queue_size:      100,
-        num_worker_threads:  1,
-        sleep_interval:      0.5,
-        block_on_full_queue: false
-      }
-    end
-
     subject { worker }
     before { worker.stop! }
 
-    specify { expect(subject.max_post_points).to be 10 }
-    specify { expect(subject.max_queue_size).to be 100 }
-    specify { expect(subject.num_worker_threads).to be 1 }
-    specify { expect(subject.sleep_interval).to be_within(0.0001).of(0.5) }
-    specify { expect(subject.block_on_full_queue).to be false }
-    specify { expect(subject.queue).to be_kind_of(InfluxDB::MaxQueue) }
+    context 'when all options are given' do
+      let(:async_options) do
+        {
+          max_post_points:     10,
+          max_queue_size:      100,
+          num_worker_threads:  1,
+          sleep_interval:      0.5,
+          block_on_full_queue: false,
+          shutdown_timeout:    0.6,
+        }
+      end
+
+      it "uses the specified values" do
+        expect(subject.max_post_points).to be 10
+        expect(subject.max_queue_size).to be 100
+        expect(subject.num_worker_threads).to be 1
+        expect(subject.sleep_interval).to be_within(0.0001).of(0.5)
+        expect(subject.block_on_full_queue).to be false
+        expect(subject.queue).to be_kind_of(InfluxDB::MaxQueue)
+        expect(subject.shutdown_timeout).to be_within(0.0001).of(0.6)
+      end
+    end
+
+    context 'when only sleep_interval is given' do
+      let(:async_options) { { sleep_interval: 0.2 } }
+
+      it "uses a value for shutdown_timeout that is 2x sleep_interval" do
+        expect(subject.sleep_interval).to be_within(0.0001).of(0.2)
+        expect(subject.shutdown_timeout).to be_within(0.0001).of(0.4)
+      end
+    end
+
+    context 'when only shutdown_timeout is given' do
+      let(:async_options) { { shutdown_timeout: 0.3 } }
+
+      it "uses that value" do
+        expect(subject.sleep_interval).to be_within(0.0001).of(5)
+        expect(subject.shutdown_timeout).to be_within(0.0001).of(0.3)
+      end
+    end
   end
 end

spec/influxdb/config_spec.rb

@@ -1,9 +1,10 @@
 require 'spec_helper'
 
 describe InfluxDB::Config do
-  let(:conf) do
-    InfluxDB::Client.new(*args).config
-  end
+  after { client.stop! }
+
+  let(:client) { InfluxDB::Client.new(*args) }
+  let(:conf) { client.config }
 
   let(:args) { {} }