chore-core 1.5.2

data/lib/chore/queues/sqs/consumer.rb

@@ -0,0 +1,121 @@
+require 'aws/sqs'
+require 'chore/duplicate_detector'
+
+AWS.eager_autoload! AWS::Core
+AWS.eager_autoload! AWS::SQS
+
+module Chore
+  module Queues
+    module SQS
+      # SQS Consumer for Chore. Requests messages from SQS and passes them to be worked on. Also controls
+      # deleting completed messages within SQS.
+      class Consumer < Chore::Consumer
+        # Initialize the reset at on class load
+        @@reset_at = Time.now
+
+        Chore::CLI.register_option 'aws_access_key', '--aws-access-key KEY', 'Valid AWS Access Key'
+        Chore::CLI.register_option 'aws_secret_key', '--aws-secret-key KEY', 'Valid AWS Secret Key'
+        Chore::CLI.register_option 'dedupe_servers', '--dedupe-servers SERVERS', 'List of mememcache compatible server(s) to use for storing SQS Message Dedupe cache'
+        Chore::CLI.register_option 'queue_polling_size', '--queue_polling_size NUM', Integer, 'Amount of messages to grab on each request' do |arg|
+          raise ArgumentError, "Cannot specify a queue polling size greater than 10" if arg > 10
+        end
+
+        def initialize(queue_name, opts={})
+          super(queue_name, opts)
+        end
+
+        # Sets a flag that instructs the publisher to reset the connection the next time it's used
+        def self.reset_connection!
+          @@reset_at = Time.now
+        end
+
+        # Begins requesting messages from SQS, which will invoke the +&handler+ over each message
+        def consume(&handler)
+          while running?
+            begin
+              messages = handle_messages(&handler)
+              sleep 1 if messages.empty?
+            rescue AWS::SQS::Errors::NonExistentQueue => e
+              Chore.logger.error "You specified a queue that does not exist. You must create the queue before starting Chore. Shutting down..."
+              raise Chore::TerribleMistake  
+            rescue => e
+              Chore.logger.error { "SQSConsumer#Consume: #{e.inspect} #{e.backtrace * "\n"}" }
+            end
+          end
+        end
+
+        # Rejects the given message from SQS by +id+. Currently a noop
+        def reject(id)
+
+        end
+
+        # Deletes the given message from SQS by +id+
+        def complete(id)
+          Chore.logger.debug "Completing (deleting): #{id}"
+          queue.batch_delete([id])
+        end
+
+        private
+
+        # Requests messages from SQS, and invokes the provided +&block+ over each one. Afterwards, the :on_fetch
+        # hook will be invoked, per message
+        def handle_messages(&block)
+          msg = queue.receive_messages(:limit => sqs_polling_amount, :attributes => [:receive_count])
+
+          messages = *msg
+          messages.each do |message|
+            block.call(message.handle, queue_name, queue_timeout, message.body, message.receive_count - 1) unless duplicate_message?(message)
+            Chore.run_hooks_for(:on_fetch, message.handle, message.body)
+          end
+          messages
+        end
+
+        # Checks if the given message has already been received within the timeout window for this queue
+        def duplicate_message?(message)
+          dupe_detector.found_duplicate?(message)
+        end
+
+        # Returns the instance of the DuplicateDetector used to ensure unique messages.
+        # Will create one if one doesn't already exist
+        def dupe_detector
+          @dupes ||= DuplicateDetector.new({:servers => Chore.config.dedupe_servers,
+                                            :dupe_on_cache_failure => Chore.config.dupe_on_cache_failure})
+        end
+
+        # Retrieves the SQS queue with the given +name+. The method will cache the results to prevent round trips on
+        # subsequent calls. If <tt>reset_connection!</tt> has been called, this will result in the connection being
+        # re-initialized, as well as clear any cached results from prior calls
+        def queue
+          if !@sqs_last_connected || (@@reset_at && @@reset_at >= @sqs_last_connected)
+            AWS::Core::Http::ConnectionPool.pools.each do |p|
+              p.empty!
+            end
+            @sqs = nil
+            @sqs_last_connected = Time.now
+            @queue = nil
+          end
+          @queue_url ||= sqs.queues.url_for(@queue_name)
+          @queue ||= sqs.queues[@queue_url]
+        end
+
+        # The visibility timeout of the queue for this consumer
+        def queue_timeout
+          @queue_timeout ||= queue.visibility_timeout
+        end
+
+        # Access to the configured SQS connection object
+        def sqs
+          @sqs ||= AWS::SQS.new(
+            :access_key_id => Chore.config.aws_access_key,
+            :secret_access_key => Chore.config.aws_secret_key,
+            :logger => Chore.logger,
+            :log_level => :debug)
+        end
+
+        def sqs_polling_amount
+          Chore.config.queue_polling_size || 10
+        end
+      end
+    end
+  end
+end

data/lib/chore/queues/sqs/publisher.rb

@@ -0,0 +1,55 @@
+require 'chore/publisher'
+
+module Chore
+  module Queues
+    module SQS
+      
+      # SQS Publisher, for writing messages to SQS from Chore
+      class Publisher < Chore::Publisher
+        @@reset_next = true
+
+        def initialize(opts={})
+          super
+          @sqs_queues = {}
+          @sqs_queue_urls = {}
+        end
+
+        # Takes a given Chore::Job instance +job+, and publishes it by looking up the +queue_name+.
+        def publish(queue_name,job)
+          queue = self.queue(queue_name)
+          queue.send_message(encode_job(job))
+        end
+
+        # Sets a flag that instructs the publisher to reset the connection the next time it's used
+        def self.reset_connection!
+          @@reset_next = true
+        end
+
+        # Access to the configured SQS connection object
+        def sqs
+          @sqs ||= AWS::SQS.new(
+            :access_key_id => Chore.config.aws_access_key,
+            :secret_access_key => Chore.config.aws_secret_key,
+            :logger => Chore.logger,
+            :log_level => :debug)
+        end
+
+        # Retrieves the SQS queue with the given +name+. The method will cache the results to prevent round trips on subsequent calls
+        # If <tt>reset_connection!</tt> has been called, this will result in the connection being re-initialized,
+        # as well as clear any cached results from prior calls
+        def queue(name)
+         if @@reset_next
+            AWS::Core::Http::ConnectionPool.pools.each do |p|
+              p.empty!
+            end
+            @sqs = nil
+            @@reset_next = false
+            @sqs_queues = {}
+          end
+          @sqs_queue_urls[name] ||= self.sqs.queues.url_for(name)
+          @sqs_queues[name] ||= self.sqs.queues[@sqs_queue_urls[name]]
+        end
+      end
+    end
+  end
+end

data/lib/chore/queues/sqs.rb

@@ -0,0 +1,38 @@
+module Chore
+  module Queues
+    module SQS
+      # Helper method to create queues based on the currently known list as provided by your configured Chore::Jobs
+      # This is meant to be invoked from a rake task, and not directly.
+      # These queues will be created with the default settings, which may not be ideal. 
+      # This is meant only as a convenience helper for testing, and not as a way to create production quality queues in SQS
+      def self.create_queues!
+        raise 'You must have atleast one Chore Job configured and loaded before attempting to create queues' unless Chore.prefixed_queue_names.length > 0
+        #This will raise an exception if AWS has not been configured by the project making use of Chore
+        sqs_queues = AWS::SQS.new.queues
+        Chore.prefixed_queue_names.each do |queue_name|
+          Chore.logger.info "Chore Creating Queue: #{queue_name}"
+          begin
+            sqs_queues.create(queue_name)
+          rescue AWS::SQS::Errors::QueueAlreadyExists
+            Chore.logger.info "exists with different config"
+          end
+        end
+        Chore.prefixed_queue_names
+      end
+
+      # Helper method to delete all known queues based on the list as provided by your configured Chore::Jobs
+      # This is meant to be invoked from a rake task, and not directly.
+      def self.delete_queues!
+        raise 'You must have atleast one Chore Job configured and loaded before attempting to create queues' unless Chore.prefixed_queue_names.length > 0
+        #This will raise an exception if AWS has not been configured by the project making use of Chore
+        sqs_queues = AWS::SQS.new.queues
+        Chore.prefixed_queue_names.each do |queue_name|
+          Chore.logger.info "Chore Deleting Queue: #{queue_name}"
+          url = sqs_queues.url_for(queue_name)
+          sqs_queues[url].delete
+        end
+        Chore.prefixed_queue_names
+      end
+    end
+  end
+end

data/lib/chore/railtie.rb

@@ -0,0 +1,18 @@
+module Chore
+  # Railtie for incorporating chores rake tasks into your Rails application
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      Dir[File.join(File.dirname(__FILE__),'tasks/*.task')].each { |f| load f }
+    end
+
+    config.after_initialize do
+      if Chore.configuring?
+        # Reset the logger on forks to avoid deadlocks
+        Rails.logger = Chore.logger
+        Chore.add_hook(:after_fork) do |worker|
+          Rails.logger = Chore.logger
+        end
+      end
+    end
+  end
+end

data/lib/chore/signal.rb

@@ -0,0 +1,175 @@
+module Chore
+  # Provides smarter signal handling capabilities than Ruby's built-in
+  # Signal class.  Specifically it runs callbacks in a separate thread since:
+  #   (1) Ruby 2.0 cannot obtain locks in the main Signal thread
+  #   (2) Doing so can result in deadlocks in Ruby 1.9.x.
+  # 
+  # Ruby's core implementation can be found at: http://ruby-doc.org/core-1.9.3/Signal.html
+  # 
+  # == Differences
+  # 
+  # There are a few important differences with the way signals trapped through
+  # this class behave than through Ruby's Signal class.
+  # 
+  # === Sequential processing
+  # 
+  # In Ruby, signals are interrupt-driven -- the thread is executing at the time
+  # will be interrupted at that point in the call stack and start executing the
+  # signal handler.  This increases the potential for deadlocks if mutexes are
+  # in use by both the thread and the signal handler.
+  # 
+  # In Chore, signal handlers are executed sequentially.  When a handler is
+  # started, it must complete before the next signal is processed.  These
+  # handlers are also executed in their own thread and, therefore, will compete
+  # for resources with the rest of the application.
+  # 
+  # === Forking
+  # 
+  # In Ruby, forking does not disrupt the ability to process signals.  Signals
+  # trapped in the master process will continue to be trapped in forked child
+  # processes.
+  # 
+  # In Chore, this is not the case.  When a process is forked, any trapped
+  # signals will no longer get processed.  This is because the thread that
+  # processes those incoming signals gets killed.
+  # 
+  # In order to process these signals, `Chore::Signal.reset` must be called,
+  # followed by additional calls to re-register those signal handlers.
+  # 
+  # == Signal ordering
+  # 
+  # It is important to note that in Ruby, signals are essentially processed
+  # as LIFO (Last-In, First-Out) since they are interrupt driven.  Similar
+  # behaviors is present in Chore's implementation.
+  # 
+  # Having LIFO behavior is the reason why this class uses a queue for
+  # tracking the list of incoming signals, instead of writing them out to a
+  # pipe.
+  class Signal
+    # The handlers registered for trapping certain signals.  Maps signal => handler.
+    @handlers = {}
+
+    # The set of incoming, unprocessed high-priority signals (such as QUIT / INT)
+    @primary_signals = []
+
+    # The set of incoming, unprocessed low-priority signals (such as CHLD)
+    @secondary_signals = []
+
+    # The priorities of signals to handle.  If not defined, the signal is
+    # considered high-priority.
+    PRIORITIES = {
+      'CHLD' => :secondary
+    }
+
+    # Stream used to track when signals are ready to be processed
+    @wake_in, @wake_out = IO.pipe
+
+    class << self
+      # Traps the given signal and runs the block when the signal is sent to
+      # this process.  This will run the block outside of the trap thread.
+      # 
+      # Only a single handler can be registered for a signal at any point.  If
+      # a signal has already been trapped, a warning will be generated and the
+      # previous handler for the signal will be returned.
+      # 
+      # See ::Signal#trap @ http://ruby-doc.org/core-1.9.3/Signal.html#method-c-trap
+      # for more information.
+      def trap(signal, command = nil, &block)
+        if command
+          # Command given for Ruby to interpret -- pass it directly onto Signal
+          @handlers.delete(signal)
+          ::Signal.trap(signal, command)
+        else
+          # Ensure we're listening for signals
+          listen
+
+          if @handlers[signal]
+            Chore.logger.debug "#{signal} signal has been overwritten:\n#{caller * "\n"}"
+          end
+
+          # Wrap handlers so they run in the listener thread
+          signals = PRIORITIES[signal] == :secondary ? @secondary_signals : @primary_signals
+          @handlers[signal] = block
+          ::Signal.trap(signal) do
+            signals << signal
+            wakeup
+          end
+        end
+      end
+
+      # Resets signals and handlers back to their defaults.  Any unprocessed
+      # signals will be discarded.
+      # 
+      # This should be called after forking a processing in order to ensure
+      # that signals continue to get processed.  *Note*, however, that new
+      # handlers must get registered after forking.
+      def reset
+        # Reset traps back to their default behavior.  Note that this *must*
+        # be done first in order to prevent trap handlers from being called
+        # while the wake pipe / listener are being reset.  If this is run
+        # out of order, then it's possible for those callbacks to hit errors.
+        @handlers.keys.each {|signal| trap(signal, 'DEFAULT')}
+
+        # Reset signals back to their empty state
+        @listener = nil
+        @primary_signals.clear
+        @secondary_signals.clear
+        @wake_out.close
+        @wake_in.close
+        @wake_in, @wake_out = IO.pipe
+      end
+
+      private
+      # Starts the thread that processes incoming signals
+      def listen
+        @listener ||= Thread.new do
+          on_wakeup do
+            while signal = next_signal
+              process(signal)
+            end
+          end
+        end
+      end
+
+      # Looks up what the next signal is to process.  Signals are typically
+      # processed LIFO (Last In, First Out), though primary signals are
+      # prioritized over secondary signals.
+      def next_signal
+        @primary_signals.pop || @secondary_signals.pop
+      end
+
+      # Waits until a wakeup signal is received.  When it is received, the
+      # provided block will be yielded to.
+      def on_wakeup
+        begin
+          while @wake_in.getc
+            yield
+          end
+        rescue IOError => e
+          # Ignore: listener has been stopped
+          Chore.logger.debug "Signal stream closed: #{e}\n#{e.backtrace * "\n"}"
+        end
+      end
+
+      # Wakes up the listener thread to indicate that signals are ready to be
+      # processed
+      def wakeup
+        @wake_out.write('.')
+      end
+
+      # Processes the given signal by running the handler in a separate
+      # thread.
+      def process(signal)
+        handler = @handlers[signal]
+        if handler
+          begin
+            handler.call
+          rescue => e
+            # Prevent signal handlers from killing the listener thread
+            Chore.logger.error "Failed to run #{signal} signal handler: #{e}\n#{e.backtrace * "\n"}"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/chore/strategies/consumer/batcher.rb

@@ -0,0 +1,76 @@
+module Chore
+  module Strategy
+
+    # Handles holding jobs in memory until such time as the batch has become full, per the developers configured threshold,
+    # or enough time elapses that Chore determines to not wait any longer (20 seconds by default)
+    class Batcher
+      attr_accessor :callback
+      attr_accessor :batch
+
+      def initialize(size)
+        @size = size
+        @batch = []
+        @mutex = Mutex.new
+        @last_message = nil
+        @callback = nil
+        @running = true
+      end
+
+      # The main entry point of the Batcher, <tt>schedule</tt> begins a thread with the provided +batch_timeout+ 
+      # as the only argument. While the Batcher is running, it will attempt to check if either the batch is full, 
+      # or if the +batch_timeout+ has elapsed since the last batch was executed. If the batch is full, it will be executed.
+      # If the +batch_timeout+ has elapsed, as soon as the next message enters the batch, it will be executed.
+      # 
+      # Calling <tt>stop</tt> will cause the thread to finish it's current check, and exit
+      def schedule(batch_timeout=20)
+        @thread = Thread.new(batch_timeout) do |timeout|
+          Chore.logger.info "Batching timeout thread starting"
+          while @running do
+            begin 
+              Chore.logger.debug "Last message added to batch: #{@last_message}: #{@batch.size}"
+              if @last_message && Time.now > (@last_message + timeout)
+                Chore.logger.debug "Batching timeout reached (#{@last_message + timeout}), current size: #{@batch.size}"
+                self.execute(true)
+                @last_message = nil
+              end
+              sleep(1) 
+            rescue => e
+              Chore.logger.error "Batcher#schedule raised an exception: #{e.inspect}"
+            end
+          end
+        end
+      end
+
+      # Adds the +item+ to the current batch
+      def add(item)
+        @batch << item
+        @last_message = Time.now
+        execute if ready?
+      end
+
+      # Calls for the batch to be executed. If +force+ is set to true, the batch will execute even if it is not full yet
+      def execute(force = false)
+        batch = nil
+        @mutex.synchronize do
+          if force || ready?
+            batch = @batch.slice!(0...@size)
+          end
+        end
+
+        if batch && !batch.empty?
+          @callback.call(batch)
+        end
+      end
+
+      # Determines if the batch is ready to fire, by comparing it's size to the configured batch_size
+      def ready?
+        @batch.size >= @size
+      end
+
+      # Sets a flag which will begin shutting down the Batcher
+      def stop
+        @running = false
+      end
+    end
+  end
+end

data/lib/chore/strategies/consumer/single_consumer_strategy.rb

@@ -0,0 +1,34 @@
+module Chore
+  module Strategy
+
+    # Consumer strategy for requesting batches of work in a linear fashion. Ideally used for running
+    # a single Chore job locally in a development environment where performance or throughput may not matter.
+    # <tt>SingleConsumerStrategy</tt> will raise an exception if you're configured to listen for more than 1 queue
+    class SingleConsumerStrategy
+      def initialize(fetcher, opts={})
+        @fetcher = fetcher
+      end
+
+      # Begins fetching from the configured queue by way of the configured Consumer. This can only be used if you have a
+      # single queue which can be kept up with at a relatively low volume. If you have more than a single queue configured,
+      # it will raise an exception.
+      def fetch
+        Chore.logger.debug "Starting up consumer strategy: #{self.class.name}"
+        queues = Chore.config.queues
+        raise "When using SingleConsumerStrategy only one queue can be defined. Queues: #{queues}" unless queues.size == 1
+        
+        @consumer = Chore.config.consumer.new(queues.first)
+        @consumer.consume do |id,queue_name,queue_timeout,body,previous_attempts|
+          work = UnitOfWork.new(id, queue_name, queue_timeout, body, previous_attempts, @consumer)
+          @fetcher.manager.assign(work)
+        end
+      end
+
+      # Stops consuming messages from the queue
+      def stop!
+        Chore.logger.info "Shutting down fetcher: #{self.class.name.to_s}"
+        @consumer.stop
+      end
+    end
+  end
+end

data/lib/chore/strategies/consumer/threaded_consumer_strategy.rb

@@ -0,0 +1,81 @@
+require 'chore/strategies/consumer/batcher'
+module Chore
+  module Strategy
+    class ThreadedConsumerStrategy #:nodoc
+      attr_accessor :batcher
+
+      Chore::CLI.register_option 'batch_size', '--batch-size SIZE', Integer, 'Number of items to collect for a single worker to process'
+      Chore::CLI.register_option 'threads_per_queue', '--threads-per-queue NUM_THREADS', Integer, 'Number of threads to create for each named queue'
+
+      def initialize(fetcher)
+        @fetcher = fetcher
+        @batcher = Batcher.new(Chore.config.batch_size)
+        @batcher.callback = lambda { |batch| @fetcher.manager.assign(batch) }
+        @batcher.schedule
+        @running = true
+      end
+
+      # Begins fetching from queues by spinning up the configured +:threads_per_queue:+ count of threads
+      # for each queue you're consuming from.
+      # Once all threads are spun up and running, the threads are then joined.
+      def fetch
+        Chore.logger.debug "Starting up consumer strategy: #{self.class.name}"
+        threads = []
+        Chore.config.queues.each do |queue|
+          Chore.config.threads_per_queue.times do 
+            if running?
+              threads << start_consumer_thread(queue)
+            end
+          end
+        end
+
+        threads.each(&:join)
+      end
+      
+      # If the ThreadedConsumerStrategy is currently running <tt>stop!</tt> will begin signalling it to stop
+      # It will stop the batcher from forking more work, as well as set a flag which will disable it's own consuming
+      # threads once they finish with their current work.
+      def stop!
+        if running?
+          Chore.logger.info "Shutting down fetcher: #{self.class.name.to_s}"
+          @batcher.stop
+          @running = false
+        end
+      end
+
+      # Returns whether or not the ThreadedConsumerStrategy is running or not
+      def running?
+        @running
+      end
+
+      private 
+      # Starts a consumer thread for polling the given +queue+.
+      # If <tt>stop!<tt> is called, the threads will shut themsevles down.
+      def start_consumer_thread(queue)
+        t = Thread.new(queue) do |tQueue|
+          begin
+            consumer = Chore.config.consumer.new(tQueue)
+            consumer.consume do |id, queue_name, queue_timeout, body, previous_attempts|
+              # Quick hack to force this thread to end it's work
+              # if we're shutting down. Could be delayed due to the
+              # weird sometimes-blocking nature of SQS.
+              consumer.stop if !running?
+              Chore.logger.debug { "Got message: #{id}"}
+
+              work = UnitOfWork.new(id, queue_name, queue_timeout, body, previous_attempts, consumer)
+              @batcher.add(work)
+            end
+          rescue Chore::TerribleMistake
+            Chore.logger.error "I've made a terrible mistake... shutting down Chore"
+            self.stop!
+            @fetcher.manager.shutdown!
+          rescue => e
+            Chore.logger.error "ThreadedConsumerStrategy#consumer thread raised an exception: #{e.inspect} at #{e.backtrace}"
+          end
+        end
+        t
+      end
+
+    end #ThreadedConsumerStrategy
+  end
+end #Chore