chore-core 1.8.2 → 4.0.0

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

@@ -3,21 +3,26 @@ 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
 
+        # @param [Hash] opts Publisher options
         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+.
+        # Publishes a message to an SQS queue
+        #
+        # @param [String] queue_name Name of the SQS queue
+        # @param [Hash] job Job instance definition, will be encoded to JSON
+        #
+        # @return [struct Aws::SQS::Types::SendMessageResult]
         def publish(queue_name,job)
-          queue = self.queue(queue_name)
-          queue.send_message(encode_job(job))
+          queue = queue(queue_name)
+          queue.send_message(message_body: encode_job(job))
         end
 
         # Sets a flag that instructs the publisher to reset the connection the next time it's used
@@ -25,29 +30,33 @@ module Chore
           @@reset_next = true
         end
 
-        # Access to the configured SQS connection object
+        private
+
+        # SQS API client object
+        #
+        # @return [Aws::SQS::Client]
         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)
+          @sqs ||= Chore::Queues::SQS.sqs_client
         end
 
-        # Retrieves the SQS queue with the given +name+. The method will cache the results to prevent round trips on subsequent calls
+        # Retrieves the SQS queue object. 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
+        #
+        # @param [String] name Name of SQS queue
+        #
+        # @return [Aws::SQS::Queue]
         def queue(name)
-         if @@reset_next
-            AWS::Core::Http::ConnectionPool.pools.each do |p|
-              p.empty!
-            end
+          if @@reset_next
+            Aws.empty_connection_pools!
             @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]]
+
+          @sqs_queue_urls[name] ||= sqs.get_queue_url(queue_name: name).queue_url
+          @sqs_queues[name] ||= Aws::SQS::Queue.new(url: @sqs_queue_urls[name], client: sqs)
         end
       end
     end

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

@@ -11,29 +11,29 @@ module Chore
         @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.
-      # 
+      # 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 oldest message was added. If either case is true, the
+      # items in the batch will be executed.
+      #
       # Calling <tt>stop</tt> will cause the thread to finish it's current check, and exit
-      def schedule(batch_timeout=20)
+      def schedule(batch_timeout)
         @thread = Thread.new(batch_timeout) do |timeout|
-          Chore.logger.info "Batching timeout thread starting"
+          Chore.logger.info "Batching thread starting with #{batch_timeout} second timeout"
           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}"
+            begin
+              oldest_item = @batch.first
+              timestamp = oldest_item && oldest_item.created_at
+              Chore.logger.debug "Oldest message in batch: #{timestamp}, size: #{@batch.size}"
+              if timestamp && Time.now > (timestamp + timeout)
+                Chore.logger.debug "Batching timeout reached (#{timestamp + timeout}), current size: #{@batch.size}"
                 self.execute(true)
-                @last_message = nil
               end
-              sleep(1) 
+              sleep(1)
             rescue => e
               Chore.logger.error "Batcher#schedule raised an exception: #{e.inspect}"
             end
@@ -44,7 +44,6 @@ module Chore
       # Adds the +item+ to the current batch
       def add(item)
         @batch << item
-        @last_message = Time.now
         execute if ready?
       end
 

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

@@ -10,16 +10,16 @@ module Chore
       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.
+      # 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)
+        @consumer.consume do |message_id, message_receipt_handle, queue_name, queue_timeout, body, previous_attempts|
+          work = UnitOfWork.new(message_id, message_receipt_handle, queue_name, queue_timeout, body, previous_attempts, @consumer)
           @fetcher.manager.assign(work)
         end
       end

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

@@ -5,13 +5,14 @@ module Chore
       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 'batch_timeout', '--batch-timeout SIZE', Integer, 'Maximum number of seconds to wait until processing a message'
       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
+        @batcher.schedule(Chore.config.batch_timeout)
         @running = true
       end
 
@@ -22,7 +23,7 @@ module Chore
         Chore.logger.debug "Starting up consumer strategy: #{self.class.name}"
         threads = []
         Chore.config.queues.each do |queue|
-          Chore.config.threads_per_queue.times do 
+          Chore.config.threads_per_queue.times do
             if running?
               threads << start_consumer_thread(queue)
             end
@@ -31,7 +32,7 @@ module Chore
 
         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.
@@ -48,21 +49,22 @@ module Chore
         @running
       end
 
-      private 
+      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|
+            consumer.consume do |message_id, message_receipt_handle, 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}"}
+              Chore.logger.debug { "Got message: #{message_id}"}
 
-              work = UnitOfWork.new(id, queue_name, queue_timeout, body, previous_attempts, consumer)
+              work = UnitOfWork.new(message_id, message_receipt_handle, queue_name, queue_timeout, body, previous_attempts, consumer)
+              Chore.run_hooks_for(:consumed_from_source, work)
               @batcher.add(work)
             end
           rescue Chore::TerribleMistake

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

@@ -0,0 +1,120 @@
+module Chore
+  module Strategy
+    class ThrottledConsumerStrategy #:nodoc:
+      def initialize(fetcher)
+        @fetcher = fetcher
+        @queue = SizedQueue.new(Chore.config.num_workers)
+        @return_queue = Queue.new
+        @max_queue_size = Chore.config.num_workers
+        @consumers_per_queue = Chore.config.threads_per_queue
+        @running = true
+        @consumers = []
+      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.info "TCS: Starting up: #{self.class.name}"
+        threads = []
+        Chore.config.queues.each do |consume_queue|
+          Chore.logger.info "TCS: Starting #{@consumers_per_queue} threads for Queue #{consume_queue}"
+          @consumers_per_queue.times do
+            next unless running?
+            threads << consume(consume_queue)
+          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 "TCS: Shutting down fetcher: #{self.class.name}"
+          @running = false
+          @consumers.each do |consumer|
+            Chore.logger.info "TCS: Stopping consumer: #{consumer.object_id}"
+            @queue.clear
+            @return_queue.clear
+            consumer.stop
+          end
+        end
+      end
+
+      # Returns whether or not the ThreadedConsumerStrategy is running or not
+      def running?
+        @running
+      end
+
+      # return upto number_of_free_workers work objects
+      def provide_work(no_free_workers)
+        work_units = []
+        free_workers = [no_free_workers, @queue.size + @return_queue.size].min
+        while free_workers > 0
+          # Drain from the return queue first, then the consumer thread queue
+          queue = @return_queue.empty? ? @queue : @return_queue
+          work_units << queue.pop
+          free_workers -= 1
+        end
+        work_units
+      end
+
+      # Gives work back to the queue in case it couldn't be assigned
+      #
+      # This will go into a separate queue so that it will be prioritized
+      # over other work that hasn't been attempted yet.  It also avoids
+      # a deadlock where @queue is full and the master is waiting to return
+      # work that it couldn't assign.
+      def return_work(work_units)
+        work_units.each do |work|
+          @return_queue.push(work)
+        end
+      end
+
+      private
+
+      def consume(consume_queue)
+        consumer = Chore.config.consumer.new(consume_queue)
+        @consumers << consumer
+        start_consumer_thread(consumer)
+      end
+
+      # Starts a consumer thread for polling the given +consume_queue+.
+      # If <tt>stop!<tt> is called, the threads will shut themsevles down.
+      def start_consumer_thread(consumer)
+        t = Thread.new(consumer) do |th|
+          begin
+            create_work_units(th)
+          rescue Chore::TerribleMistake => e
+            Chore.logger.error 'Terrible mistake, shutting down Chore'
+            Chore.logger.error "#{e.inspect} at #{e.backtrace}"
+            @fetcher.manager.shutdown!
+          end
+        end
+        t
+      end
+
+      def create_work_units(consumer)
+        consumer.consume do |message_id, message_receipt_handle, queue, timeout, body, previous_attempts|
+          # Note: The unit of work object contains a consumer object that when
+          # used to consume from SQS, would have a mutex (that comes as a part
+          # of the AWS sdk); When sending these objects across from one process
+          # to another, we cannot send this across (becasue of the mutex). To
+          # work around this, we simply ignore the consumer object when creating
+          # the unit of work object, and when the worker recieves the work
+          # object, it assigns it a consumer object.
+          # (to allow for communication back to the queue it was consumed from)
+          work = UnitOfWork.new(message_id, message_receipt_handle, queue, timeout, body, previous_attempts)
+          Chore.run_hooks_for(:consumed_from_source, work)
+          @queue.push(work) if running?
+          Chore.run_hooks_for(:added_to_queue, work)
+        end
+      end
+    end # ThrottledConsumerStrategy
+  end
+end # Chore

data/lib/chore/strategies/worker/forked_worker_strategy.rb

@@ -3,6 +3,7 @@ require 'chore/signal'
 module Chore
   module Strategy
     class ForkedWorkerStrategy #:nodoc:
+      include Util
       attr_accessor :workers
 
       def initialize(manager, opts={})
@@ -63,6 +64,9 @@ module Chore
           pid = nil
           Chore.run_hooks_for(:around_fork,w) do
             pid = fork do
+              work.each do | item |
+                Chore.run_hooks_for(:fetched_off_internal_q, item)
+              end
               after_fork(w)
               Chore.run_hooks_for(:within_fork,w) do
                 Chore.run_hooks_for(:after_fork,w)
@@ -132,7 +136,7 @@ module Chore
       def after_fork(worker)
         # Immediately swap out the process name so that it doesn't look like
         # the master process
-        procline("Started:#{Time.now}")
+        procline("#{Chore.config.worker_procline}:Started:#{Time.now}")
 
         clear_child_signals
         trap_child_signals(worker)
@@ -202,11 +206,6 @@ module Chore
         Kernel.fork(&block)
       end
 
-      def procline(str)
-        Chore.logger.info str
-        $0 = "chore-#{Chore::VERSION}:#{str}"
-      end
-
       def signal_children(sig, pids_to_signal = pids)
         pids_to_signal.each do |pid|
           begin

data/lib/chore/strategies/worker/helpers/ipc.rb

@@ -0,0 +1,87 @@
+require 'socket'
+
+module Chore
+  module Strategy
+    module Ipc #:nodoc:
+      BIG_ENDIAN = 'L>'.freeze
+      MSG_BYTES = 4
+      READY_MSG = 'R'
+
+      def create_master_socket
+        File.delete socket_file if File.exist? socket_file
+        UNIXServer.new(socket_file).tap do |socket|
+          socket_options(socket)
+        end
+      end
+
+      def child_connection(socket)
+        socket.accept
+      end
+
+      # Sending a message to a socket (must be a connected socket)
+      def send_msg(socket, msg)
+        raise 'send_msg cannot send empty messages' if msg.nil? || msg.size.zero?
+        message = Marshal.dump(msg)
+        encoded_size = [message.size].pack(BIG_ENDIAN)
+        encoded_message = "#{encoded_size}#{message}"
+        socket.send encoded_message, 0
+      end
+
+      # read a message from socket (must be a connected socket)
+      def read_msg(socket)
+        encoded_size = socket.recv(MSG_BYTES, Socket::MSG_PEEK)
+        return if encoded_size.nil? || encoded_size == ''
+
+        size = encoded_size.unpack(BIG_ENDIAN).first
+        encoded_message = socket.recv(MSG_BYTES + size)
+        Marshal.load(encoded_message[MSG_BYTES..-1])
+      rescue Errno::ECONNRESET => ex
+        Chore.logger.info "IPC: Connection was closed on socket #{socket}"
+        raise ex
+      end
+
+      def add_worker_socket
+        UNIXSocket.new(socket_file).tap do |socket|
+          socket_options(socket)
+        end
+      end
+
+      def clear_ready(socket)
+        _ = socket.gets
+      end
+
+      def signal_ready(socket)
+        socket.puts READY_MSG
+      rescue Errno::EPIPE => ex
+        Chore.logger.info 'IPC: Connection was shutdown by master'
+        raise ex
+      end
+
+      def select_sockets(sockets, self_pipe = nil, timeout = 0.5)
+        all_socks = [sockets, self_pipe].flatten.compact
+        IO.select(all_socks, nil, all_socks, timeout)
+      end
+
+      def delete_socket_file
+        File.unlink(socket_file)
+      rescue
+        nil
+      end
+
+      # Used for unit tests
+      def ipc_help
+        :available
+      end
+
+      private
+
+      def socket_file
+        "./prefork_worker_sock-#{Process.pid}"
+      end
+
+      def socket_options(socket)
+        socket.setsockopt(:SOCKET, :REUSEADDR, true)
+      end
+    end
+  end
+end