Support list patterns

Author: Juuso Mäyränen

lib/logstash/inputs/redis.rb

@@ -3,6 +3,7 @@
 require "logstash/inputs/base"
 require "logstash/inputs/threadable"
 require 'redis'
+require 'concurrent/atomics'
 
 # 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
@@ -18,6 +19,7 @@
 #
 module LogStash module Inputs class Redis < LogStash::Inputs::Threadable
   BATCH_EMPTY_SLEEP = 0.25
+  MAX_ITEMS_PER_WORKER = 1000
 
   config_name "redis"
 
@@ -51,14 +53,16 @@ module LogStash module Inputs class Redis < LogStash::Inputs::Threadable
   # 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
+  config :data_type, :validate => [ "list", "list_pattern", "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
+
   public
   # public API
   # use to store a proc that can provide a Redis instance or mock
@@ -86,6 +90,9 @@ def register
     if @data_type == 'list' || @data_type == 'dummy'
       @run_method = method(:list_runner)
       @stop_method = method(:list_stop)
+    elsif @data_type == 'list_pattern'
+      @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)
@@ -217,6 +224,86 @@ def list_runner(output_queue)
     end
   end
 
+  # private
+  def list_pattern_stop
+    return if @redis.nil? || !@redis.connected?
+
+    @redis.quit rescue nil
+    @redis = nil
+    @worker_threads.each do |key, thread|
+      if thread.alive?
+        thread.join(5)
+        @logger.warn("Thread #{key} joined")
+      end
+    end
+  end
+
+  def launch_worker(output_queue, key)
+     Thread.new do
+      redis = new_redis_instance
+      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
+      ensure
+        redis.quit rescue nil
+      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
+
+  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
+    end
+  end
+
+  def init_list_pattern_runner(output_queue)
+    @worker_threads ||= {}
+    @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)
+      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)
     begin
       results = redis.evalsha(@redis_script_sha, [@key], [@batch_count-1])