Support list patterns

Author: Juuso Mäyränen

lib/logstash/inputs/redis.rb

@@ -3,6 +3,8 @@
 require "logstash/inputs/base"
 require "logstash/inputs/threadable"
 require 'redis'
+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
@@ -49,16 +51,30 @@ 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", "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 => {}
 
+  # 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
   # use to store a proc that can provide a Redis instance or mock
@@ -77,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}"
 
@@ -86,6 +111,9 @@ def register
     if @data_type == 'list' || @data_type == 'dummy'
       @run_method = method(:list_runner)
       @stop_method = method(:list_stop)
+    elsif @data_type == 'pattern_list'
+      @run_method = method(:pattern_list_runner)
+      @stop_method = method(:pattern_list_stop)
     elsif @data_type == 'channel'
       @run_method = method(:channel_runner)
       @stop_method = method(:subscribe_stop)
@@ -94,8 +122,6 @@ def register
       @stop_method = method(:subscribe_stop)
     end
 
-    @list_method = batched? ? method(:list_batch_listener) : method(:list_single_listener)
-
     @identity = "#{@redis_url} #{@data_type}:#{@key}"
     @logger.info("Registering Redis", :identity => @identity)
   end # def register
@@ -119,7 +145,7 @@ def batched?
 
   # private
   def is_list_type?
-    @data_type == 'list'
+    @data_type == 'list' || @data_type == 'pattern_list'
   end
 
   # private
@@ -193,15 +219,21 @@ 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)
+    @list_method = batched? ? method(:list_batch_listener) : method(:list_single_listener)
     while !stop?
       begin
         @redis ||= connect
@@ -217,16 +249,113 @@ def list_runner(output_queue)
     end
   end
 
-  def list_batch_listener(redis, output_queue)
+  #private
+  def reset_threadpool
+    return if @threadpool.nil?
+    @threadpool.shutdown
+    @threadpool.wait_for_termination
+    @threadpool = nil
+  end
+
+  # private
+  def pattern_list_stop
+    reset_redis
+    reset_threadpool
+  end
+
+  # private
+  def pattern_list_process_item(redis, output_queue, key)
+    if stop?
+      @logger.debug("Breaking from thread #{key} as it was requested to stop")
+      return false
+    end
+    value = redis.lpop(key)
+    return false if value.nil?
+    queue_event(value, output_queue)
+    true
+  end
+
+  # private
+  def pattern_list_single_processor(redis, output_queue, key)
+    (0...@pattern_list_max_items).each do
+      break unless pattern_list_process_item(redis, output_queue, key)
+    end
+  end
+
+  # private
+  def pattern_list_batch_processor(redis, output_queue, key)
+    items_left = @pattern_list_max_items
+    while items_left > 0
+      limit = [items_left, @batch_count].min
+      processed = process_batch(redis, output_queue, key, limit, 0)
+      if processed.zero? || processed < limit
+        return
+      end
+      items_left -= processed
+    end
+  end
+
+  # private
+  def pattern_list_worker_consume(output_queue, key)
     begin
-      results = redis.evalsha(@redis_script_sha, [@key], [@batch_count-1])
-      results.each do |item|
-        queue_event(item, output_queue)
+      redis ||= connect
+      @pattern_list_processor.call(redis, output_queue, key)
+    rescue ::Redis::BaseError => e
+      @logger.warn("Redis connection problem in thread for key #{key}. Sleeping a while before exiting thread.", :exception => e)
+      sleep 1
+      return
+    ensure
+      redis.quit rescue nil
+    end
+  end
+
+  # private
+  def threadpool_capacity?
+    @threadpool.remaining_capacity > 0
+  end
+
+  # private
+  def pattern_list_launch_worker(output_queue, key)
+    @current_workers.add(key)
+    @threadpool.post do
+      begin
+        pattern_list_worker_consume(output_queue, key)
+      ensure
+        @current_workers.delete(key)
       end
+    end
+  end
 
-      if results.size.zero?
-        sleep BATCH_EMPTY_SLEEP
+  # private
+  def pattern_list_ensure_workers(output_queue)
+    return unless threadpool_capacity?
+    redis_runner do
+      @redis.keys(@key).shuffle.each do |key|
+        next if @current_workers.include?(key)
+        pattern_list_launch_worker(output_queue, key)
+        break unless threadpool_capacity?
       end
+    end
+  end
+
+  # private
+  def pattern_list_runner(output_queue)
+    @pattern_list_processor = batched? ? method(:pattern_list_batch_processor) : method(:pattern_list_single_processor)
+    while !stop?
+      init_threadpool if @threadpool.nil?
+      pattern_list_ensure_workers(output_queue)
+      sleep(@pattern_list_threadpool_sleep)
+    end
+  end
+
+  def process_batch(redis, output_queue, key, batch_size, sleep_time)
+    begin
+      results = redis.evalsha(@redis_script_sha, [key], [batch_size-1])
+      results.each do |item|
+        queue_event(item, output_queue)
+      end
+      sleep sleep_time if results.size.zero? && sleep_time > 0
+      results.size
 
       # Below is a commented-out implementation of 'batch fetch'
       # using pipelined LPOP calls. This in practice has been observed to
@@ -255,6 +384,10 @@ def list_batch_listener(redis, output_queue)
     end
   end
 
+  def list_batch_listener(redis, output_queue)
+    process_batch(redis, output_queue, @key, @batch_count, BATCH_EMPTY_SLEEP)
+  end
+
   def list_single_listener(redis, output_queue)
     item = redis.blpop(@key, 0, :timeout => 1)
     return unless item # from timeout or other conditions

spec/inputs/redis_spec.rb

@@ -63,6 +63,48 @@ 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 => "pattern_list"
+          batch_count => 1
+        }
+      }
+    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
+
+  it "should read events from a list pattern using batch_count (default 125)" do
+    key_base = SecureRandom.hex
+    conf = <<-CONFIG
+      input {
+        redis {
+          type => "blah"
+          key => "#{key}.*"
+          data_type => "pattern_list"
+          batch_count => 125
+        }
+      }
+    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 +306,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 +496,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 } }