refactored code

Fixes logstash-plugins#14

Author: guyboertje

lib/logstash/inputs/redis.rb

@@ -1,7 +1,7 @@
 # encoding: utf-8
+require "logstash/namespace"
 require "logstash/inputs/base"
 require "logstash/inputs/threadable"
-require "logstash/namespace"
 
 # 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
@@ -15,7 +15,9 @@
 # `batch_count` note: If you use the `batch_count` setting, you *must* use a Redis version 2.6.0 or
 # newer. Anything older does not support the operations used by batching.
 #
-class LogStash::Inputs::Redis < LogStash::Inputs::Threadable
+module Logstash module Inputs class Redis < LogStash::Inputs::Threadable
+# class LogStash::Inputs::Redis < LogStash::Inputs::Threadable
+
  config_name "redis"
 
  default :codec, "json"
@@ -57,14 +59,33 @@ class LogStash::Inputs::Redis < LogStash::Inputs::Threadable
  config :batch_count, :validate => :number, :default => 1
 
  public
+ # public API
+
+ REDIS_INPUT_POISON_MSG = '}-{ poison message from teardown }-{'.freeze
+
+ # use to store a proc that can provide a redis instance or mock
+ def add_external_redis_builder(builder) #callable
+ @redis_builder = builder
+ self
+ end
+
+ # use to apply an instance directly and bypass the builder
+ def use_redis(instance)
+ @redis = instance
+ self
+ end
+
+ def new_redis_instance
+ @redis_builder.call
+ end
+
  def register
  require 'redis'
- @redis = nil
  @redis_url = "redis://#{@password}@#{@host}:#{@port}/#{@db}"
 
  # TODO remove after setting key and data_type to true
  if @queue
- if @key or @data_type
+ if @key || @data_type
  raise RuntimeError.new(
  "Cannot specify queue parameter and key or data_type"
  )
@@ -73,38 +94,79 @@ def register
  @data_type = 'list'
  end
 
- if not @key or not @data_type
+ if !@key || !@data_type
  raise RuntimeError.new(
  "Must define queue, or key and data_type parameters"
  )
  end
  # end TODO
 
- @logger.info("Registering Redis", :identity => identity)
+ @redis_builder ||= method(:internal_redis_builder)
+
+ # just switch on data_type once
+ if @data_type == 'list' || @data_type == 'dummy'
+ @run_method = method(:list_runner)
+ @teardown_method = method(:list_teardown)
+ elsif @data_type == 'channel'
+ @run_method = method(:channel_runner)
+ @teardown_method = method(:subscribe_teardown)
+ elsif @data_type == 'pattern_channel'
+ @run_method = method(:pattern_channel_runner)
+ @teardown_method = method(:subscribe_teardown)
+ end
+
+ # TODO(sissel, boertje): set @identity directly when @name config option is removed.
+ @identity = @name != 'default' ? @name : "#{@redis_url} #{@data_type}:#{@key}"
+ @logger.info("Registering Redis", :identity => @identity)
  end # def register
 
- # A string used to identify a Redis instance in log messages
- # TODO(sissel): Use instance variables for this once the @name config
- # option is removed.
- private
- def identity
- @name || "#{@redis_url} #{@data_type}:#{@key}"
+ def run(output_queue)
+ @run_method.call(output_queue)
+ rescue LogStash::ShutdownSignal
+ # ignore and quit
+ end # def run
+
+ def teardown
+ @shutdown_requested = true
+ @teardown_method.call
  end
 
+ # private methods -----------------------------
  private
- def connect
- redis = Redis.new(
+
+ def batched?
+ @batch_count > 1
+ end
+
+ # private
+ def is_list_type?
+ @data_type == 'list'
+ end
+
+ # private
+ def redis_params
+ {
  :host => @host,
  :port => @port,
  :timeout => @timeout,
  :db => @db,
  :password => @password.nil? ? nil : @password.value
- )
- load_batch_script(redis) if @data_type == 'list' && (@batch_count > 1)
- return redis
+ }
+ end
+
+ # private
+ def internal_redis_builder
+ ::Redis.new(redis_params)
+ end
+
+ # private
+ def connect
+ redis = new_redis_instance
+ load_batch_script(redis) if batched? && is_list_type?
+ redis
  end # def connect
 
- private
+ # private
  def load_batch_script(redis)
  #A Redis Lua EVAL script to fetch a count of keys
  #in case count is bigger than current items in queue whole queue will be returned without extra nil values
@@ -126,7 +188,7 @@ def load_batch_script(redis)
  @redis_script_sha = redis.script(:load, redis_script)
  end
 
- private
+ # private
  def queue_event(msg, output_queue)
  begin
  @codec.decode(msg) do |event|
@@ -141,18 +203,51 @@ def queue_event(msg, output_queue)
  end
  end
 
- private
+ # private
+ def shutting_down?
+ @shutdown_requested
+ end
+
+ # private
+ def running?
+ !@shutdown_requested
+ end
+
+ # private
+ def list_teardown
+ return if @redis.nil? || !@redis.connected?
+
+ @redis.quit rescue nil
+ @redis = nil
+ end
+
+ # private
+ def list_runner(output_queue)
+ while running?
+ begin
+ @redis ||= connect
+ list_listener(@redis, output_queue)
+ rescue ::Redis::BaseError => e
+ @logger.warn("Redis connection problem", :exception => e)
+ # Reset the redis variable to trigger reconnect
+ @redis = nil
+ sleep 1
+ end
+ end
+ end
+
+ # private
  def list_listener(redis, output_queue)
 
  item = redis.blpop(@key, 0, :timeout => 1)
  return unless item # from timeout or other conditions
 
  # blpop returns the 'key' read from as well as the item result
  # we only care about the result (2nd item in the list).
- queue_event(item[1], output_queue)
+ queue_event(item.last, output_queue)
 
  # If @batch_count is 1, there's no need to continue.
- return if @batch_count == 1
+ return if !batched?
 
  begin
  redis.evalsha(@redis_script_sha, [@key], [@batch_count-1]).each do |item|
@@ -173,7 +268,7 @@ def list_listener(redis, output_queue)
  #queue_event(item, output_queue) if item
  #end
  # --- End commented out implementation of 'batch fetch'
- rescue Redis::CommandError => e
+ rescue ::Redis::CommandError => e
  if e.to_s =~ /NOSCRIPT/ then
  @logger.warn("Redis may have been restarted, reloading Redis batch EVAL script", :exception => e);
  load_batch_script(redis)
@@ -184,15 +279,64 @@ def list_listener(redis, output_queue)
  end
  end
 
- private
- def channel_listener(redis, output_queue)
- redis.subscribe @key do |on|
+ # private
+ def subscribe_teardown
+ return if @redis.nil? || !@redis.connected?
+
+ # the underlying client in redis is blocked on subscribe
+ # need new redis instance, send poison message, causes
+ # an unsubscribe
+ publish_poison_message
+
+ # if its a SubscribedClient then:
+ # a) the poison_message did not work
+ # b) it does not have a disconnect method (yet)
+ if @redis.client.is_a?(::Redis::Client)
+ @redis.client.disconnect
+ end
+ @redis = nil
+ end
+
+ # private
+ def publish_poison_message
+ new_redis_instance.publish(@key, REDIS_INPUT_POISON_MSG)
+ end
+
+ # private
+ def redis_runner
+ begin
+ @redis ||= connect
+ yield
+ rescue ::Redis::BaseError => e
+ @logger.warn("Redis connection problem", :exception => e)
+ # Reset the redis variable to trigger reconnect
+ @redis = nil
+ sleep 1
+ retry
+ end
+ end
+
+ # private
+ def channel_runner(output_queue)
+ redis_runner do
+ channel_listener(output_queue)
+ end
+ end
+
+ # private
+ def channel_listener(output_queue)
+ @redis.subscribe(@key) do |on|
  on.subscribe do |channel, count|
  @logger.info("Subscribed", :channel => channel, :count => count)
  end
 
  on.message do |channel, message|
- queue_event message, output_queue
+ if message != REDIS_INPUT_POISON_MSG
+ queue_event(message, output_queue)
+ end
+ if shutting_down?
+ @redis.unsubscribe(@key)
+ end
  end
 
  on.unsubscribe do |channel, count|
@@ -201,15 +345,26 @@ def channel_listener(redis, output_queue)
  end
  end
 
- private
- def pattern_channel_listener(redis, output_queue)
- redis.psubscribe @key do |on|
+ def pattern_channel_runner(output_queue)
+ redis_runner do
+ pattern_channel_listener(output_queue)
+ end
+ end
+
+ # private
+ def pattern_channel_listener(output_queue)
+ @redis.psubscribe @key do |on|
  on.psubscribe do |channel, count|
  @logger.info("Subscribed", :channel => channel, :count => count)
  end
 
- on.pmessage do |ch, event, message|
- queue_event message, output_queue
+ on.pmessage do |pattern, channel, message|
+ if message != REDIS_INPUT_POISON_MSG
+ queue_event(message, output_queue)
+ end
+ if shutting_down?
+ @redis.punsubscribe(@key)
+ end
  end
 
  on.punsubscribe do |channel, count|
@@ -218,51 +373,6 @@ def pattern_channel_listener(redis, output_queue)
  end
  end
 
- # Since both listeners have the same basic loop, we've abstracted the outer
- # loop.
- private
- def listener_loop(listener, output_queue)
- while !@shutdown_requested
- begin
- @redis ||= connect
- self.send listener, @redis, output_queue
- rescue Redis::BaseError => e
- @logger.warn("Redis connection problem", :exception => e)
- # Reset the redis variable to trigger reconnect
- @redis = nil
- sleep 1
- end
- end
- end # listener_loop
+# end
 
- public
- def run(output_queue)
- if @data_type == 'list'
- listener_loop :list_listener, output_queue
- elsif @data_type == 'channel'
- listener_loop :channel_listener, output_queue
- else
- listener_loop :pattern_channel_listener, output_queue
- end
- rescue LogStash::ShutdownSignal
- # ignore and quit
- end # def run
-
- public
- def teardown
- @shutdown_requested = true
-
- if @redis
- if @data_type == 'list'
- @redis.quit rescue nil
- elsif @data_type == 'channel'
- @redis.unsubscribe rescue nil
- @redis.connection.disconnect
- elsif @data_type == 'pattern_channel'
- @redis.punsubscribe rescue nil
- @redis.connection.disconnect
- end
- @redis = nil
- end
- end
-end # class LogStash::Inputs::Redis
+end end end # Redis Inputs LogStash