Support list patterns

Author: Juuso Mäyränen

lib/logstash/inputs/redis.rb

@@ -18,6 +18,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 +52,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 +89,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 +223,73 @@ def list_runner(output_queue)
     end
   end
 
+  # private
+  def list_pattern_stop
+    return if @redis.nil? || !@redis.connected?
+
+    @redis.quit rescue nil
+    @redis = nil
+  end
+
+  def launch_worker(key)
+     Thread.new do
+      redis = new_redis_instance
+      begin
+        (0..MAX_ITEMS_PER_WORKER).each do
+          value = redis.lpop(key)
+          break if value.nil?
+          @events << value
+        end
+      ensure
+        redis.quit rescue nil
+      end
+    end
+  end
+
+  def clear_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
+    free_slots = @worker_thread_count - @worker_threads.length
+    return if free_slots == 0
+    keys = @redis.keys(@key)
+    keys.shuffle
+    keys = keys[0..free_slots - 1]
+    keys.each do |key|
+      @worker_threads[key] = launch_worker(key)
+    end
+  end
+
+  # private
+  def list_pattern_runner(output_queue)
+    begin
+      @events = Queue.new
+      @worker_threads = {}
+      @redis ||= connect
+      ensure_workers
+    rescue ::Redis::BaseError => e
+      @logger.warn("Redis connection problem", :exception => e)
+      @redis = nil
+    end
+    while !stop?
+      begin
+        clear_workers
+        ensure_workers
+        item = @events.pop(true)
+        next unless item
+        queue_event(item, output_queue)
+      rescue ThreadError => e
+        sleep(0.1)
+      end
+    end
+  end
+
   def list_batch_listener(redis, output_queue)
     begin
       results = redis.evalsha(@redis_script_sha, [@key], [@batch_count-1])