Use thread pools

Author: Juuso Mäyränen

lib/logstash/inputs/redis.rb

@@ -3,7 +3,8 @@
 require "logstash/inputs/base"
 require "logstash/inputs/threadable"
 require 'redis'
-require 'concurrent/atomics'
+require 'concurrent'
+require 'concurrent/executors'
 
 # This input will read events from a Redis instance; it supports both Redis channels and lists.
 # The list command (BLPOP) used by Logstash is supported in Redis v1.3.1+, and
@@ -19,7 +20,6 @@
 #
 module LogStash module Inputs class Redis < LogStash::Inputs::Threadable
   BATCH_EMPTY_SLEEP = 0.25
-  MAX_ITEMS_PER_WORKER = 1000
 
   config_name "redis"
 
@@ -51,17 +51,29 @@ module LogStash module Inputs class Redis < LogStash::Inputs::Threadable
   config :key, :validate => :string, :required => true
 
   # Specify either list or channel.  If `data_type` is `list`, then we will BLPOP the
-  # key.  If `data_type` is `channel`, then we will SUBSCRIBE to the key.
-  # If `data_type` is `pattern_channel`, then we will PSUBSCRIBE to the key.
-  config :data_type, :validate => [ "list", "list_pattern", "channel", "pattern_channel" ], :required => true
+  # key. If `data_type` is `pattern_list`, then we will spawn a number of worker
+  # threads that will LPOP from keys matching that pattern. If `data_type` is
+  # `channel`, then we will SUBSCRIBE to the key. If `data_type` is `pattern_channel`,
+  # then we will PSUBSCRIBE to the key.
+  config :data_type, :validate => [ "list", "pattern_list", "channel", "pattern_channel" ], :required => true
 
   # The number of events to return from Redis using EVAL.
   config :batch_count, :validate => :number, :default => 125
 
   # Redefined Redis commands to be passed to the Redis client.
   config :command_map, :validate => :hash, :default => {}
 
-  config :worker_thread_count, :validate => :number, :default => 20
+  # Maximum number of worker threads to spawn when using `data_type` `pattern_list`.
+  config :pattern_list_threads, :validate => :number, :default => 20
+
+  # Maximum number of items for a single worker thread to process when `data_type` is `pattern_list`.
+  # After the list is empty or this number of items have been processed, the thread will exit and a
+  # new one will be started if there are non-empty lists matching the pattern without a consumer.
+  config :pattern_list_max_items, :validate => :number, :default => 1000
+
+  # Time to sleep in main loop after checking if more threads can/need to be spawned.
+  # Applies to `data_type` is `pattern_list`
+  config :pattern_list_threadpool_sleep, :validate => :number, :default => 0.2
 
   public
   # public API
@@ -81,6 +93,15 @@ def new_redis_instance
     @redis_builder.call
   end
 
+  def init_threadpool
+    @threadpool ||= Concurrent::ThreadPoolExecutor.new(
+        min_threads: @pattern_list_threads,
+        max_threads: @pattern_list_threads,
+        max_queue: 2 * @pattern_list_threads
+    )
+    @current_workers ||= Concurrent::Set.new
+  end
+
   def register
     @redis_url = @path.nil? ? "redis://#{@password}@#{@host}:#{@port}/#{@db}" : "#{@password}@#{@path}/#{@db}"
 
@@ -90,7 +111,7 @@ def register
     if @data_type == 'list' || @data_type == 'dummy'
       @run_method = method(:list_runner)
       @stop_method = method(:list_stop)
-    elsif @data_type == 'list_pattern'
+    elsif @data_type == 'pattern_list'
       @run_method = method(:list_pattern_runner)
       @stop_method = method(:list_pattern_stop)
     elsif @data_type == 'channel'
@@ -200,13 +221,18 @@ def queue_event(msg, output_queue, channel=nil)
   end
 
   # private
-  def list_stop
+  def reset_redis
     return if @redis.nil? || !@redis.connected?
 
     @redis.quit rescue nil
     @redis = nil
   end
 
+  # private
+  def list_stop
+    reset_redis
+  end
+
   # private
   def list_runner(output_queue)
     while !stop?
@@ -224,84 +250,86 @@ def list_runner(output_queue)
     end
   end
 
+  #private
+  def reset_threadpool
+    return if @threadpool.nil?
+    @threadpool.shutdown
+    @threadpool.wait_for_termination
+    @threadpool = nil
+  end
+
   # private
   def list_pattern_stop
-    return if @redis.nil? || !@redis.connected?
+    reset_redis
+    reset_threadpool
+  end
 
-    @redis.quit rescue nil
-    @redis = nil
-    @worker_threads.each do |key, thread|
-      if thread.alive?
-        thread.join(5)
-        @logger.warn("Thread #{key} joined")
+  # private
+  def worker_consume(output_queue, key)
+    @logger.debug("Launched worker for #{key}")
+    redis = new_redis_instance
+    begin
+      (0...@pattern_list_max_items).each do
+        if stop?
+          @logger.debug("Breaking from thread #{key} as it was requested to stop")
+          break
+        end
+        value = redis.lpop(key)
+        break if value.nil?
+        queue_event(value, output_queue)
       end
+    ensure
+      redis.quit rescue nil
     end
+    @logger.debug("Exit worker for #{key}")
+  end
+
+  # private
+  def threadpool_capacity?
+    @threadpool.remaining_capacity > 0
   end
 
+  # private
   def launch_worker(output_queue, key)
-     Thread.new do
-      redis = new_redis_instance
+    @current_workers.add(key)
+    @threadpool.post do
       begin
-        (0..MAX_ITEMS_PER_WORKER).each do
-          if stop?
-            @logger.warn("Breaking from thread #{key} as it was requested to stop")
-            break
-          end
-          value = redis.lpop(key)
-          break if value.nil?
-          queue_event(value, output_queue)
-          @items_processed.increment
-        end
+        worker_consume(output_queue, key)
       ensure
-        redis.quit rescue nil
+        @current_workers.delete(key)
       end
     end
   end
 
-  def clear_finished_workers
-    finished_threads = []
-    @worker_threads.each do |key, thread|
-      next if thread.alive?
-      finished_threads << key
-    end
-    finished_threads.each { |key| @worker_threads.delete(key) }
-  end
-
+  # private
   def ensure_workers(output_queue)
-    free_slots = @worker_thread_count - @worker_threads.length
-    return if free_slots == 0
-    keys = @redis.keys(@key)
-    keys.shuffle
-    keys.each do |key|
-      next if @worker_threads.has_key?(key)
-      @worker_threads[key] = launch_worker(output_queue, key)
-      free_slots -= 1
-      break if free_slots == 0
+    return unless threadpool_capacity?
+    @redis.keys(@key).shuffle.each do |key|
+      next if @current_workers.include?(key)
+      launch_worker(output_queue, key)
+      break unless threadpool_capacity?
     end
   end
 
-  def init_list_pattern_runner(output_queue)
-    @worker_threads ||= {}
+  # private
+  def init_list_pattern_runner
+    init_threadpool
     @redis ||= connect
-    @items_processed ||= Concurrent::AtomicFixnum.new
-    clear_finished_workers
-    ensure_workers(output_queue)
   end
 
   # private
   def list_pattern_runner(output_queue)
-    items = 0
     while !stop?
       begin
-        init_list_pattern_runner(output_queue)
-        sleep(0.1)
+        init_list_pattern_runner if @redis.nil?
+        ensure_workers(output_queue)
+        sleep(@pattern_list_threadpool_sleep)
       rescue ::Redis::BaseError => e
         @logger.warn("Redis connection problem", :exception => e)
         @redis = nil
         sleep 1
       end
     end
-    @logger.warn("Total items processed: #{@items_processed}")
   end
 
   def list_batch_listener(redis, output_queue)

spec/inputs/redis_spec.rb

@@ -63,6 +63,26 @@ def process(conf, event_count)
     populate(key, event_count)
     process(conf, event_count)
   end
+
+  it "should read events from a list pattern" do
+    key_base = SecureRandom.hex
+    conf = <<-CONFIG
+      input {
+        redis {
+          type => "blah"
+          key => "#{key}.*"
+          data_type => "list_pattern"
+        }
+      }
+    CONFIG
+    total_event_count = 0
+    (0..10).each do |idx|
+      event_count = 100 + rand(50)
+      total_event_count += event_count
+      populate("#{key_base}.#{idx}", event_count)
+    end
+    process(conf, total_event_count)
+  end
 end
 
 # unit tests ---------------------
@@ -264,6 +284,64 @@ def process(conf, event_count)
     end
   end
 
+  context 'runtime for pattern_list data_type' do
+    let(:data_type) { 'pattern_list' }
+    let(:key) { 'foo.*' }
+    before do
+      subject.register
+      subject.init_threadpool
+    end
+
+    context 'close when redis is unset' do
+      let(:quit_calls) { [:quit, :unsubscribe, :punsubscribe, :connection, :disconnect!] }
+
+      it 'does not attempt to quit' do
+        allow(redis).to receive(:nil?).and_return(true)
+        quit_calls.each do |call|
+          expect(redis).not_to receive(call)
+        end
+        expect {subject.do_stop}.not_to raise_error
+      end
+    end
+
+    it 'calling the run method, adds events to the queue' do
+      expect(redis).to receive(:keys).at_least(:once).and_return(['foo.bar'])
+      expect(redis).to receive(:lpop).at_least(:once).and_return('l1')
+
+      allow(redis).to receive(:connected?).and_return(connected.last)
+      allow(redis).to receive(:quit)
+
+      tt = Thread.new do
+        end_by = Time.now + 3
+        while accumulator.size < 1 and Time.now <= end_by
+          sleep 0.1
+        end
+        subject.do_stop
+      end
+
+      subject.run(accumulator)
+
+      tt.join
+
+      expect(accumulator.size).to be > 0
+    end
+
+    it 'multiple close calls, calls to redis once' do
+      subject.use_redis(redis)
+      allow(redis).to receive(:keys).at_least(:once).and_return(['foo.bar'])
+      allow(redis).to receive(:lpop).and_return('l1')
+      expect(redis).to receive(:connected?).and_return(connected.last)
+      quit_calls.each do |call|
+        expect(redis).to receive(call).at_most(:once)
+      end
+
+      subject.do_stop
+      connected.push(false) #can't use let block here so push to array
+      expect {subject.do_stop}.not_to raise_error
+      subject.do_stop
+    end
+  end
+
   context 'for the subscribe data_types' do
     def run_it_thread(inst)
       Thread.new(inst) do |subj|
@@ -396,7 +474,7 @@ def close_thread(inst, rt)
 
   describe LogStash::Inputs::Redis do
     context "when using data type" do
-      ["list", "channel", "pattern_channel"].each do |data_type|
+      ["list", "channel", "pattern_channel", "pattern_list"].each do |data_type|
         context data_type do
           it_behaves_like "an interruptible input plugin" do
             let(:config) { {'key' => 'foo', 'data_type' => data_type } }