beetle 0.1

data/lib/beetle/configuration.rb

@@ -0,0 +1,31 @@
+module Beetle
+  class Configuration
+    # default logger (defaults to <tt>Logger.new(STDOUT)</tt>)
+    attr_accessor :logger
+    # number of seconds after which keys are removed form the message deduplication store (defaults to <tt>3.days</tt>)
+    attr_accessor :gc_threshold
+    # the machines where the deduplication store lives (defaults to <tt>"localhost:6379"</tt>)
+    attr_accessor :redis_hosts
+    # redis database number to use for the message deduplication store (defaults to <tt>4</tt>)
+    attr_accessor :redis_db
+    # list of amqp servers to use (defaults to <tt>"localhost:5672"</tt>)
+    attr_accessor :servers
+    # the virtual host to use on the AMQP servers
+    attr_accessor :vhost
+    # the AMQP user to use when connecting to the AMQP servers
+    attr_accessor :user
+    # the password to use when connectiong to the AMQP servers
+    attr_accessor :password
+
+    def initialize #:nodoc:
+      self.logger = Logger.new(STDOUT)
+      self.gc_threshold = 3.days
+      self.redis_hosts = "localhost:6379"
+      self.redis_db = 4
+      self.servers = "localhost:5672"
+      self.vhost = "/"
+      self.user = "guest"
+      self.password = "guest"
+    end
+  end
+end

data/lib/beetle/deduplication_store.rb

@@ -0,0 +1,152 @@
+module Beetle
+  # The deduplication store is used internally by Beetle::Client to store information on
+  # the status of message processing. This includes:
+  # * how often a message has already been seen by some consumer
+  # * whether a message has been processed successfully
+  # * how many attempts have been made to execute a message handler for a given message
+  # * how long we should wait before trying to execute the message handler after a failure
+  # * how many exceptions have been raised during previous execution attempts
+  # * how long we should wait before trying to perform the next execution attempt
+  # * whether some other process is already trying to execute the message handler
+  #
+  # It also provides a method to garbage collect keys for expired messages.
+  class DeduplicationStore
+    # creates a new deduplication store
+    def initialize(hosts = "localhost:6379", db = 4)
+      @hosts = hosts
+      @db = db
+    end
+
+    # get the Redis instance
+    def redis
+      @redis ||= find_redis_master
+    end
+
+    # list of key suffixes to use for storing values in Redis.
+    KEY_SUFFIXES = [:status, :ack_count, :timeout, :delay, :attempts, :exceptions, :mutex, :expires]
+
+    # build a Redis key out of a message id and a given suffix
+    def key(msg_id, suffix)
+      "#{msg_id}:#{suffix}"
+    end
+
+    # list of keys which potentially exist in Redis for the given message id
+    def keys(msg_id)
+      KEY_SUFFIXES.map{|suffix| key(msg_id, suffix)}
+    end
+
+    # extract message id from a given Redis key
+    def msg_id(key)
+      key =~ /^(msgid:[^:]*:[-0-9a-f]*):.*$/ && $1
+    end
+
+    # garbage collect keys in Redis (always assume the worst!)
+    def garbage_collect_keys(now = Time.now.to_i)
+      keys = redis.keys("msgid:*:expires")
+      threshold = now + Beetle.config.gc_threshold
+      keys.each do |key|
+        expires_at = redis.get key
+        if expires_at && expires_at.to_i < threshold
+          msg_id = msg_id(key)
+          redis.del(keys(msg_id))
+        end
+      end
+    end
+
+    # unconditionally store a key <tt>value></tt> with given <tt>suffix</tt> for given <tt>msg_id</tt>.
+    def set(msg_id, suffix, value)
+      with_failover { redis.set(key(msg_id, suffix), value) }
+    end
+
+    # store a key <tt>value></tt> with given <tt>suffix</tt> for given <tt>msg_id</tt> if it doesn't exists yet.
+    def setnx(msg_id, suffix, value)
+      with_failover { redis.setnx(key(msg_id, suffix), value) }
+    end
+
+    # store some key/value pairs if none of the given keys exist.
+    def msetnx(msg_id, values)
+      values = values.inject({}){|h,(k,v)| h[key(msg_id, k)] = v; h}
+      with_failover { redis.msetnx(values) }
+    end
+
+    # increment counter for key with given <tt>suffix</tt> for given <tt>msg_id</tt>. returns an integer.
+    def incr(msg_id, suffix)
+      with_failover { redis.incr(key(msg_id, suffix)) }
+    end
+
+    # retrieve the value with given <tt>suffix</tt> for given <tt>msg_id</tt>. returns a string.
+    def get(msg_id, suffix)
+      with_failover { redis.get(key(msg_id, suffix)) }
+    end
+
+    # delete key with given <tt>suffix</tt> for given <tt>msg_id</tt>.
+    def del(msg_id, suffix)
+      with_failover { redis.del(key(msg_id, suffix)) }
+    end
+
+    # delete all keys associated with the given <tt>msg_id</tt>.
+    def del_keys(msg_id)
+      with_failover { redis.del(keys(msg_id)) }
+    end
+
+    # check whether key with given suffix exists for a given <tt>msg_id</tt>.
+    def exists(msg_id, suffix)
+      with_failover { redis.exists(key(msg_id, suffix)) }
+    end
+
+    # flush the configured redis database. useful for testing.
+    def flushdb
+      with_failover { redis.flushdb }
+    end
+
+    # performs redis operations by yielding a passed in block, waiting for a new master to
+    # show up on the network if the operation throws an exception. if a new master doesn't
+    # appear after 120 seconds, we raise an exception.
+    def with_failover #:nodoc:
+      tries = 0
+      begin
+        yield
+      rescue Exception => e
+        Beetle::reraise_expectation_errors!
+        logger.error "Beetle: redis connection error '#{e}'"
+        if (tries+=1) < 120
+          @redis = nil
+          sleep 1
+          logger.info "Beetle: retrying redis operation"
+          retry
+        else
+          raise NoRedisMaster.new(e.to_s)
+        end
+      end
+    end
+
+    # find the master redis instance
+    def find_redis_master
+      masters = []
+      redis_instances.each do |redis|
+        begin
+          masters << redis if redis.info[:role] == "master"
+        rescue Exception => e
+          logger.error "Beetle: could not determine status of redis instance #{redis.server}"
+        end
+      end
+      raise NoRedisMaster.new("unable to determine a new master redis instance") if masters.empty?
+      raise TwoRedisMasters.new("more than one redis master instances") if masters.size > 1
+      logger.info "Beetle: configured new redis master #{masters.first.server}"
+      masters.first
+    end
+
+    # returns the list of redis instances
+    def redis_instances
+      @redis_instances ||= @hosts.split(/ *, */).map{|s| s.split(':')}.map do |host, port|
+         Redis.new(:host => host, :port => port, :db => @db)
+      end
+    end
+
+    # returns the configured logger
+    def logger
+      Beetle.config.logger
+    end
+
+  end
+end

data/lib/beetle/handler.rb

@@ -0,0 +1,95 @@
+module Beetle
+  # Instances of class Handler are created by the message processing logic in class
+  # Message. There should be no need to ever create them in client code, except for
+  # testing purposes.
+  #
+  # Most applications will define Handler subclasses and override the process, error and
+  # failure methods.
+  class Handler
+    # the Message instance which caused the handler to be created
+    attr_reader :message
+
+    def self.create(block_or_handler, opts={}) #:nodoc:
+      if block_or_handler.is_a? Handler
+        block_or_handler
+      elsif block_or_handler.is_a?(Class) && block_or_handler.ancestors.include?(Handler)
+        block_or_handler.new
+      else
+        new(block_or_handler, opts)
+      end
+    end
+
+    # optionally capture processor, error and failure callbacks
+    def initialize(processor=nil, opts={}) #:notnew:
+      @processor = processor
+      @error_callback = opts[:errback]
+      @failure_callback = opts[:failback]
+    end
+
+    # called when a message should be processed. if the message was caused by an RPC, the
+    # return value will be sent back to the caller. calls the initialized processor proc
+    # if a processor proc was specified when creating the Handler instance. calls method
+    # process if no proc was given. make sure to call super if you override this method in
+    # a subclass.
+    def call(message)
+      @message = message
+      if @processor
+        @processor.call(message)
+      else
+        process
+      end
+    end
+
+    # called for message processing if no processor was specfied when the handler instance
+    # was created
+    def process
+      logger.info "Beetle: received message #{message.inspect}"
+    end
+
+    # should not be overriden in subclasses
+    def process_exception(exception) #:nodoc:
+      if @error_callback
+        @error_callback.call(message, exception)
+      else
+        error(exception)
+      end
+    rescue Exception
+      Beetle::reraise_expectation_errors!
+    end
+
+    # should not be overriden in subclasses
+    def process_failure(result) #:nodoc:
+      if @failure_callback
+        @failure_callback.call(message, result)
+      else
+        failure(result)
+      end
+    rescue Exception
+      Beetle::reraise_expectation_errors!
+    end
+
+    # called when handler execution raised an exception and no error callback was
+    # specified when the handler instance was created
+    def error(exception)
+      logger.error "Beetle: handler execution raised an exception: #{exception}"
+    end
+
+    # called when message processing has finally failed (i.e., the number of allowed
+    # handler execution attempts or the number of allowed exceptions has been reached) and
+    # no failure callback was specified when this handler instance was created.
+    def failure(result)
+      logger.error "Beetle: handler has finally failed"
+    end
+
+    # returns the configured Beetle logger
+    def logger
+      Beetle.config.logger
+    end
+
+    # returns the configured Beetle logger
+    def self.logger
+      Beetle.config.logger
+    end
+
+  end
+end

data/lib/beetle/message.rb

@@ -0,0 +1,336 @@
+require "timeout"
+
+module Beetle
+  # Instances of class Message are created when a scubscription callback fires. Class
+  # Message contains the code responsible for message deduplication and determining if it
+  # should retry executing the message handler after a handler has crashed (or forcefully
+  # aborted).
+  class Message
+    # current message format version
+    FORMAT_VERSION = 1
+    # flag for encoding redundant messages
+    FLAG_REDUNDANT = 1
+    # default lifetime of messages
+    DEFAULT_TTL = 1.day
+    # forcefully abort a running handler after this many seconds.
+    # can be overriden when registering a handler.
+    DEFAULT_HANDLER_TIMEOUT = 300.seconds
+    # how many times we should try to run a handler before giving up
+    DEFAULT_HANDLER_EXECUTION_ATTEMPTS = 1
+    # how many seconds we should wait before retrying handler execution
+    DEFAULT_HANDLER_EXECUTION_ATTEMPTS_DELAY = 10.seconds
+    # how many exceptions should be tolerated before giving up
+    DEFAULT_EXCEPTION_LIMIT = 0
+
+    # server from which the message was received
+    attr_reader :server
+    # name of the queue on which the message was received
+    attr_reader :queue
+    # the AMQP header received with the message
+    attr_reader :header
+    # the uuid of the message
+    attr_reader :uuid
+    # message payload
+    attr_reader :data
+    # the message format version of the message
+    attr_reader :format_version
+    # flags sent with the message
+    attr_reader :flags
+    # unix timestamp after which the message should be considered stale
+    attr_reader :expires_at
+    # how many seconds the handler is allowed to execute
+    attr_reader :timeout
+    # how long to wait before retrying the message handler
+    attr_reader :delay
+    # how many times we should try to run the handler
+    attr_reader :attempts_limit
+    # how many exceptions we should tolerate before giving up
+    attr_reader :exceptions_limit
+    # exception raised by handler execution
+    attr_reader :exception
+    # value returned by handler execution
+    attr_reader :handler_result
+
+    def initialize(queue, header, body, opts = {})
+      @queue  = queue
+      @header = header
+      @data   = body
+      setup(opts)
+      decode
+    end
+
+    def setup(opts) #:nodoc:
+      @server           = opts[:server]
+      @timeout          = opts[:timeout]    || DEFAULT_HANDLER_TIMEOUT
+      @delay            = opts[:delay]      || DEFAULT_HANDLER_EXECUTION_ATTEMPTS_DELAY
+      @attempts_limit   = opts[:attempts]   || DEFAULT_HANDLER_EXECUTION_ATTEMPTS
+      @exceptions_limit = opts[:exceptions] || DEFAULT_EXCEPTION_LIMIT
+      @attempts_limit   = @exceptions_limit + 1 if @attempts_limit <= @exceptions_limit
+      @store            = opts[:store]
+    end
+
+    # extracts various values form the AMQP header properties
+    def decode #:nodoc:
+      amqp_headers = header.properties
+      @uuid = amqp_headers[:message_id]
+      headers = amqp_headers[:headers]
+      @format_version = headers[:format_version].to_i
+      @flags = headers[:flags].to_i
+      @expires_at = headers[:expires_at].to_i
+    end
+
+    # build hash with options for the publisher
+    def self.publishing_options(opts = {}) #:nodoc:
+      flags = 0
+      flags |= FLAG_REDUNDANT if opts[:redundant]
+      expires_at = now + (opts[:ttl] || DEFAULT_TTL)
+      opts = opts.slice(*PUBLISHING_KEYS)
+      opts[:message_id] = generate_uuid.to_s
+      opts[:headers] = {
+        :format_version => FORMAT_VERSION.to_s,
+        :flags => flags.to_s,
+        :expires_at => expires_at.to_s
+      }
+      opts
+    end
+
+    # unique message id. used to form various keys in the deduplication store.
+    def msg_id
+      @msg_id ||= "msgid:#{queue}:#{uuid}"
+    end
+
+    # current time (UNIX timestamp)
+    def now #:nodoc:
+      Time.now.to_i
+    end
+
+    # current time (UNIX timestamp)
+    def self.now #:nodoc:
+      Time.now.to_i
+    end
+
+    # a message has expired if the header expiration timestamp is msaller than the current time
+    def expired?
+      @expires_at < now
+    end
+
+    # generate uuid for publishing
+    def self.generate_uuid
+      UUID4R::uuid(1)
+    end
+
+    # whether the publisher has tried sending this message to two servers
+    def redundant?
+      @flags & FLAG_REDUNDANT == FLAG_REDUNDANT
+    end
+
+    # whether this is a message we can process without accessing the deduplication store
+    def simple?
+      !redundant? && attempts_limit == 1
+    end
+
+    # store handler timeout timestamp in the deduplication store
+    def set_timeout!
+      @store.set(msg_id, :timeout, now + timeout)
+    end
+
+    # handler timed out?
+    def timed_out?
+      (t = @store.get(msg_id, :timeout)) && t.to_i < now
+    end
+
+    # reset handler timeout in the deduplication store
+    def timed_out!
+      @store.set(msg_id, :timeout, 0)
+    end
+
+    # message handling completed?
+    def completed?
+      @store.get(msg_id, :status) == "completed"
+    end
+
+    # mark message handling complete in the deduplication store
+    def completed!
+      @store.set(msg_id, :status, "completed")
+      timed_out!
+    end
+
+    # whether we should wait before running the handler
+    def delayed?
+      (t = @store.get(msg_id, :delay)) && t.to_i > now
+    end
+
+    # store delay value in the deduplication store
+    def set_delay!
+      @store.set(msg_id, :delay, now + delay)
+    end
+
+    # how many times we already tried running the handler
+    def attempts
+      @store.get(msg_id, :attempts).to_i
+    end
+
+    # record the fact that we are trying to run the handler
+    def increment_execution_attempts!
+      @store.incr(msg_id, :attempts)
+    end
+
+    # whether we have already tried running the handler as often as specified when the handler was registered
+    def attempts_limit_reached?
+      (limit = @store.get(msg_id, :attempts)) && limit.to_i >= attempts_limit
+    end
+
+    # increment number of exception occurences in the deduplication store
+    def increment_exception_count!
+      @store.incr(msg_id, :exceptions)
+    end
+
+    # whether the number of exceptions has exceeded the limit set when the handler was registered
+    def exceptions_limit_reached?
+      @store.get(msg_id, :exceptions).to_i > exceptions_limit
+    end
+
+    # have we already seen this message? if not, set the status to "incomplete" and store
+    # the message exipration timestamp in the deduplication store.
+    def key_exists?
+      old_message = 0 == @store.msetnx(msg_id, :status =>"incomplete", :expires => @expires_at)
+      if old_message
+        logger.debug "Beetle: received duplicate message: #{msg_id} on queue: #{@queue}"
+      end
+      old_message
+    end
+
+    # aquire execution mutex before we run the handler (and delete it if we can't aquire it).
+    def aquire_mutex!
+      if mutex = @store.setnx(msg_id, :mutex, now)
+        logger.debug "Beetle: aquired mutex: #{msg_id}"
+      else
+        delete_mutex!
+      end
+      mutex
+    end
+
+    # delete execution mutex
+    def delete_mutex!
+      @store.del(msg_id, :mutex)
+      logger.debug "Beetle: deleted mutex: #{msg_id}"
+    end
+
+    # process this message and do not allow any exception to escape to the caller
+    def process(handler)
+      logger.debug "Beetle: processing message #{msg_id}"
+      result = nil
+      begin
+        result = process_internal(handler)
+        handler.process_exception(@exception) if @exception
+        handler.process_failure(result) if result.failure?
+      rescue Exception => e
+        Beetle::reraise_expectation_errors!
+        logger.warn "Beetle: exception '#{e}' during processing of message #{msg_id}"
+        logger.warn "Beetle: backtrace: #{e.backtrace.join("\n")}"
+        result = RC::InternalError
+      end
+      result
+    end
+
+    private
+
+    def process_internal(handler)
+      if expired?
+        logger.warn "Beetle: ignored expired message (#{msg_id})!"
+        ack!
+        RC::Ancient
+      elsif simple?
+        ack!
+        run_handler(handler) == RC::HandlerCrash ? RC::AttemptsLimitReached : RC::OK
+      elsif !key_exists?
+        set_timeout!
+        run_handler!(handler)
+      elsif completed?
+        ack!
+        RC::OK
+      elsif delayed?
+        logger.warn "Beetle: ignored delayed message (#{msg_id})!"
+        RC::Delayed
+      elsif !timed_out?
+        RC::HandlerNotYetTimedOut
+      elsif attempts_limit_reached?
+        ack!
+        logger.warn "Beetle: reached the handler execution attempts limit: #{attempts_limit} on #{msg_id}"
+        RC::AttemptsLimitReached
+      elsif exceptions_limit_reached?
+        ack!
+        logger.warn "Beetle: reached the handler exceptions limit: #{exceptions_limit} on #{msg_id}"
+        RC::ExceptionsLimitReached
+      else
+        set_timeout!
+        if aquire_mutex!
+          run_handler!(handler)
+        else
+          RC::MutexLocked
+        end
+      end
+    end
+
+    def run_handler(handler)
+      Timeout::timeout(@timeout) { @handler_result = handler.call(self) }
+      RC::OK
+    rescue Exception => @exception
+      Beetle::reraise_expectation_errors!
+      logger.debug "Beetle: message handler crashed on #{msg_id}"
+      RC::HandlerCrash
+    ensure
+      ActiveRecord::Base.clear_active_connections! if defined?(ActiveRecord)
+    end
+
+    def run_handler!(handler)
+      increment_execution_attempts!
+      case result = run_handler(handler)
+      when RC::OK
+        completed!
+        ack!
+        result
+      else
+        handler_failed!(result)
+      end
+    end
+
+    def handler_failed!(result)
+      increment_exception_count!
+      if attempts_limit_reached?
+        ack!
+        logger.debug "Beetle: reached the handler execution attempts limit: #{attempts_limit} on #{msg_id}"
+        RC::AttemptsLimitReached
+      elsif exceptions_limit_reached?
+        ack!
+        logger.debug "Beetle: reached the handler exceptions limit: #{exceptions_limit} on #{msg_id}"
+        RC::ExceptionsLimitReached
+      else
+        delete_mutex!
+        timed_out!
+        set_delay!
+        result
+      end
+    end
+
+    def logger
+      @logger ||= self.class.logger
+    end
+
+    def self.logger
+      Beetle.config.logger
+    end
+
+    # ack the message for rabbit. deletes all keys associated with this message in the
+    # deduplication store if we are sure this is the last message with the given msg_id.
+    def ack!
+      #:doc:
+      logger.debug "Beetle: ack! for message #{msg_id}"
+      header.ack
+      return if simple? # simple messages don't use the deduplication store
+      if !redundant? || @store.incr(msg_id, :ack_count) == 2
+        @store.del_keys(msg_id)
+      end
+    end
+  end
+end