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(:list_pattern_runner)
+      @stop_method = method(:list_pattern_stop)
     elsif @data_type == 'channel'
       @run_method = method(:channel_runner)
       @stop_method = method(:subscribe_stop)
@@ -193,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?
@@ -217,6 +250,88 @@ 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
+    reset_redis
+    reset_threadpool
+  end
+
+  # 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)
+    @current_workers.add(key)
+    @threadpool.post do
+      begin
+        worker_consume(output_queue, key)
+      ensure
+        @current_workers.delete(key)
+      end
+    end
+  end
+
+  # private
+  def ensure_workers(output_queue)
+    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
+
+  # private
+  def init_list_pattern_runner
+    init_threadpool
+    @redis ||= connect
+  end
+
+  # private
+  def list_pattern_runner(output_queue)
+    while !stop?
+      begin
+        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
+  end
+
   def list_batch_listener(redis, output_queue)
     begin
       results = redis.evalsha(@redis_script_sha, [@key], [@batch_count-1])

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 } }