shoryuken 2.1.1 → 3.1.3

data/bin/cli/sqs.rb

@@ -0,0 +1,204 @@
+require 'date'
+
+# rubocop:disable Metrics/AbcSize, Metrics/BlockLength
+module Shoryuken
+  module CLI
+    class SQS < Base
+      namespace :sqs
+
+      no_commands do
+        def normalize_dump_message(message)
+          message[:id] = message.delete(:message_id)
+          message[:message_body] = message.delete(:body)
+          message.delete(:receipt_handle)
+          message.delete(:md5_of_body)
+          message.delete(:md5_of_message_attributes)
+          message
+        end
+
+        def sqs
+          @_sqs ||= Aws::SQS::Client.new
+        end
+
+        def find_queue_url(queue_name)
+          sqs.get_queue_url(queue_name: queue_name).queue_url
+        rescue Aws::SQS::Errors::NonExistentQueue
+          fail_task "The specified queue #{queue_name} does not exist"
+        end
+
+        def batch_delete(url, messages)
+          messages.to_a.flatten.each_slice(10) do |batch|
+            sqs.delete_message_batch(
+              queue_url: url,
+              entries: batch.map { |message| { id: message.message_id, receipt_handle: message.receipt_handle } }
+            ).failed.any? do |failure|
+              say(
+                "Could not delete #{failure.id}, code: #{failure.code}, message: #{failure.message}, sender_fault: #{failure.sender_fault}",
+                :yellow
+              )
+            end
+          end
+        end
+
+        def batch_send(url, messages)
+          messages.to_a.flatten.map(&method(:normalize_dump_message)).each_slice(10) do |batch|
+            sqs.send_message_batch(queue_url: url, entries: batch).failed.any? do |failure|
+              say "Could not requeue #{failure.id}, code: #{failure.code}", :yellow
+            end
+          end
+        end
+
+        def find_all(url, limit, &block)
+          count = 0
+          batch_size = limit > 10 ? 10 : limit
+
+          loop do
+            n = limit - count
+            batch_size = n if n < batch_size
+
+            messages = sqs.receive_message(
+              queue_url: url,
+              max_number_of_messages: batch_size,
+              message_attribute_names: ['All']
+            ).messages
+
+            messages.each { |m| yield m }
+
+            count += messages.size
+
+            break if count >= limit
+            break if messages.empty?
+          end
+
+          count
+        end
+
+        def list_and_print_queues(urls)
+          attrs = %w(QueueArn ApproximateNumberOfMessages ApproximateNumberOfMessagesNotVisible LastModifiedTimestamp)
+
+          entries = urls.map { |u| sqs.get_queue_attributes(queue_url: u, attribute_names: attrs).attributes }.map do |q|
+            [
+              q['QueueArn'].split(':').last,
+              q['ApproximateNumberOfMessages'],
+              q['ApproximateNumberOfMessagesNotVisible'],
+              Time.at(q['LastModifiedTimestamp'].to_i)
+            ]
+          end
+
+          entries.unshift(['Queue', 'Messages Available', 'Messages Inflight', 'Last Modified'])
+
+          print_table(entries)
+        end
+
+        def dump_file(path, queue_name)
+          File.join(path, "#{queue_name}-#{Date.today}.jsonl")
+        end
+      end
+
+      desc 'ls [QUEUE-NAME-PREFIX]', 'Lists queues'
+      method_option :watch,          aliases: '-w',  type: :boolean,              desc: 'watch queues'
+      method_option :watch_interval,                 type: :numeric, default: 10, desc: 'watch interval'
+      def ls(queue_name_prefix = '')
+        trap('SIGINT', 'EXIT') # expect ctrl-c from loop
+
+        urls = sqs.list_queues(queue_name_prefix: queue_name_prefix).queue_urls
+
+        loop do
+          list_and_print_queues(urls)
+
+          break unless options[:watch]
+
+          sleep options[:watch_interval]
+          puts
+        end
+      end
+
+      desc 'dump QUEUE-NAME', 'Dumps messages from a queue into a JSON lines file'
+      method_option :number, aliases: '-n', type: :numeric, default: Float::INFINITY, desc: 'number of messages to dump'
+      method_option :path,   aliases: '-p', type: :string,  default: './',            desc: 'path to save the dump file'
+      method_option :delete, aliases: '-d', type: :boolean, default: true,            desc: 'delete from the queue'
+      def dump(queue_name)
+        path = dump_file(options[:path], queue_name)
+
+        fail_task "File #{path} already exists" if File.exist?(path)
+
+        url = find_queue_url(queue_name)
+
+        messages = []
+
+        file = nil
+
+        count = find_all(url, options[:number]) do |m|
+          file ||= File.open(path, 'w')
+
+          file.puts(JSON.dump(m.to_h))
+
+          messages << m if options[:delete]
+        end
+
+        batch_delete(url, messages) if options[:delete]
+
+        if count.zero?
+          say "Queue #{queue_name} is empty", :yellow
+        else
+          say "Dump saved in #{path} with #{count} messages", :green
+        end
+      ensure
+        file.close if file
+      end
+
+      desc 'requeue QUEUE-NAME PATH', 'Requeues messages from a dump file'
+      def requeue(queue_name, path)
+        fail_task "Path #{path} not found" unless File.exist?(path)
+
+        messages = File.readlines(path).map { |line| JSON.parse(line, symbolize_names: true) }
+
+        batch_send(find_queue_url(queue_name), messages)
+
+        say "Requeued #{messages.size} messages from #{path} to #{queue_name}", :green
+      end
+
+      desc 'mv QUEUE-NAME-SOURCE QUEUE-NAME-TARGET', 'Moves messages from one queue (source) to another (target)'
+      method_option :number, aliases: '-n', type: :numeric, default: Float::INFINITY, desc: 'number of messages to move'
+      method_option :delete, aliases: '-d', type: :boolean, default: true,            desc: 'delete from the queue'
+      def mv(queue_name_source, queue_name_target)
+        url_source = find_queue_url(queue_name_source)
+        messages = []
+
+        count = find_all(url_source, options[:number]) do |m|
+          messages << m
+        end
+
+        batch_send(find_queue_url(queue_name_target), messages.map(&:to_h))
+        batch_delete(url_source, messages) if options[:delete]
+
+        if count.zero?
+          say "Queue #{queue_name_source} is empty", :yellow
+        else
+          say "Moved #{count} messages from #{queue_name_source} to #{queue_name_target}", :green
+        end
+      end
+
+      desc 'purge QUEUE-NAME', 'Deletes the messages in a queue'
+      def purge(queue_name)
+        sqs.purge_queue(queue_url: find_queue_url(queue_name))
+
+        say "Purge request sent for #{queue_name}. The message deletion process takes up to 60 seconds", :yellow
+      end
+
+      desc 'create QUEUE-NAME', 'Create a queue'
+      def create(queue_name)
+        queue_url = sqs.create_queue(queue_name: queue_name).queue_url
+
+        say "Queue #{queue_name} was successfully created. Queue URL #{queue_url}", :green
+      end
+
+      desc 'delete QUEUE-NAME', 'delete a queue'
+      def delete(queue_name)
+        sqs.delete_queue(queue_url: find_queue_url(queue_name))
+
+        say "Queue #{queue_name} was successfully delete", :green
+      end
+    end
+  end
+end

data/bin/shoryuken

@@ -1,12 +1,58 @@
 #!/usr/bin/env ruby
 
-require_relative '../lib/shoryuken/cli'
-
-begin
-  Shoryuken::CLI.instance.run(ARGV)
-rescue => e
-  raise e if $DEBUG
-  STDERR.puts e.message
-  STDERR.puts e.backtrace.join("\n")
-  exit 1
+require 'rubygems'
+
+require 'thor'
+require 'aws-sdk-core'
+require_relative 'cli/base'
+require_relative 'cli/sqs'
+require_relative '../lib/shoryuken/runner'
+
+# rubocop:disable Metrics/AbcSize
+module Shoryuken
+  module CLI
+    class Runner < Base
+      default_task :start
+
+      register(Shoryuken::CLI::SQS, 'sqs', 'sqs COMMAND', 'SQS commands')
+
+      desc 'start', 'Starts shoryuken'
+      method_option :concurrency, aliases: '-c', type: :numeric, desc: 'Processor threads to use'
+      method_option :daemon,      aliases: '-d', type: :boolean, desc: 'Daemonize process'
+      method_option :queues,      aliases: '-q', type: :array,   desc: 'Queues to process with optional weights'
+      method_option :require,     aliases: '-r', type: :string,  desc: 'Dir or path of the workers'
+      method_option :config,      aliases: '-C', type: :string,  desc: 'Path to config file'
+      method_option :config_file,                type: :string,  desc: 'Path to config file (backwards compatibility)'
+      method_option :rails,       aliases: '-R', type: :boolean, desc: 'Load Rails'
+      method_option :logfile,     aliases: '-L', type: :string,  desc: 'Path to logfile'
+      method_option :pidfile,     aliases: '-P', type: :string,  desc: 'Path to pidfile'
+      method_option :verbose,     aliases: '-v', type: :boolean, desc: 'Print more verbose output'
+      method_option :delay,       aliases: '-D', type: :numeric, desc: 'Number of seconds to pause fetching from an empty queue'
+      def start
+        opts = options.to_h.symbolize_keys
+
+        if opts[:config_file]
+          say "[DEPRECATED] Please use --config instead of --config-file", :yellow
+        end
+
+        opts[:config_file] = opts.delete(:config) if opts[:config]
+
+        # Keep compatibility with old CLI queue format
+        opts[:queues] = options[:queues].map { |q| q.split(',') } if options[:queues]
+
+        if options[:daemon] && options[:logfile].nil?
+          fail_task "You should set a logfile if you're going to daemonize"
+        end
+
+        Shoryuken::Runner.instance.run(opts.freeze)
+      end
+
+      desc 'version', 'Prints version'
+      def version
+        say "Shoryuken #{Shoryuken::VERSION}"
+      end
+    end
+  end
 end
+
+Shoryuken::CLI::Runner.start

data/examples/default_worker.rb

@@ -4,6 +4,6 @@ class DefaultWorker
   shoryuken_options queue: 'default', auto_delete: true
 
   def perform(sqs_msg, body)
-    Shoryuken.logger.info("Received message: '#{body}'")
+    Shoryuken.logger.debug("Received message: #{body}")
   end
 end

data/lib/shoryuken/client.rb

@@ -1,31 +1,19 @@
 module Shoryuken
   class Client
     @@queues = {}
-    @@topics = {}
 
     class << self
       def queues(name)
         @@queues[name.to_s] ||= Shoryuken::Queue.new(sqs, name)
       end
 
-      def sns
-        @sns ||= Shoryuken::AwsConfig.sns
-      end
-
-      def sns_arn
-        @sns_arn ||= SnsArn
-      end
-
       def sqs
-        @sqs ||= Shoryuken::AwsConfig.sqs
+        @@sqs ||= Shoryuken.sqs_client
       end
 
-      def topics(name)
-        @@topics[name.to_s] ||= Topic.new(name, sns)
+      def sqs=(sqs)
+        @@sqs = sqs
       end
-
-      attr_accessor :account_id
-      attr_writer :sns, :sqs, :sqs_resource, :sns_arn
     end
   end
 end

data/lib/shoryuken/default_worker_registry.rb

@@ -14,13 +14,17 @@ module Shoryuken
 
     def fetch_worker(queue, message)
       worker_class = !message.is_a?(Array) &&
-        message.message_attributes &&
-        message.message_attributes['shoryuken_class'] &&
-        message.message_attributes['shoryuken_class'][:string_value]
+                     message.message_attributes &&
+                     message.message_attributes['shoryuken_class'] &&
+                     message.message_attributes['shoryuken_class'][:string_value]
 
-      worker_class = (worker_class.constantize rescue nil) || @workers[queue]
+      worker_class = begin
+                       worker_class.constantize
+                     rescue
+                       @workers[queue]
+                     end
 
-      worker_class.new
+      worker_class.new if worker_class
     end
 
     def queues
@@ -30,7 +34,7 @@ module Shoryuken
     def register_worker(queue, clazz)
       if (worker_class = @workers[queue])
         if worker_class.get_shoryuken_options['batch'] == true || clazz.get_shoryuken_options['batch'] == true
-          fail ArgumentError, "Could not register #{clazz} for '#{queue}', "\
+          fail ArgumentError, "Could not register #{clazz} for #{queue}, "\
             "because #{worker_class} is already registered for this queue, "\
             "and Shoryuken doesn't support a batchable worker for a queue with multiple workers"
         end

data/lib/shoryuken/environment_loader.rb

@@ -24,14 +24,12 @@ module Shoryuken
     end
 
     def load
-      load_rails if options[:rails]
+      load_rails if Shoryuken.options[:rails]
       prefix_active_job_queue_names
       parse_queues
       require_workers
-      initialize_aws
       validate_queues
       validate_workers
-      patch_deprecated_workers
     end
 
     private
@@ -44,24 +42,18 @@ module Shoryuken
     def config_file_options
       return {} unless (path = options[:config_file])
 
-      fail ArgumentError, "The supplied config file '#{path}' does not exist" unless File.exist?(path)
+      fail ArgumentError, "The supplied config file #{path} does not exist" unless File.exist?(path)
 
-      YAML.load(ERB.new(IO.read(path)).result).deep_symbolize_keys
-    end
-
-    # DEPRECATED: Please use configure_server and configure_client in
-    # https://github.com/phstc/shoryuken/blob/a81637d577b36c5cf245882733ea91a335b6602f/lib/shoryuken.rb#L82
-    # Please delete this method afert next release (v2.0.12 or later)
-    def initialize_aws
-      unless Shoryuken.options[:aws].to_h.empty?
-        Shoryuken.logger.warn { '[DEPRECATION] aws in shoryuken.yml is deprecated. Please use configure_server and configure_client in your initializer' }
+      if (result = YAML.load(ERB.new(IO.read(path)).result))
+        result.deep_symbolize_keys
+      else
+        {}
       end
-      Shoryuken::AwsConfig.setup(Shoryuken.options[:aws])
     end
 
     def initialize_logger
-      Shoryuken::Logging.initialize_logger(options[:logfile]) if options[:logfile]
-      Shoryuken.logger.level = Logger::DEBUG if options[:verbose]
+      Shoryuken::Logging.initialize_logger(Shoryuken.options[:logfile]) if Shoryuken.options[:logfile]
+      Shoryuken.logger.level = Logger::DEBUG if Shoryuken.options[:verbose]
     end
 
     def load_rails
@@ -77,62 +69,68 @@ module Shoryuken
         ::Rails::Application.initializer 'shoryuken.eager_load' do
           ::Rails.application.config.eager_load = true
         end
-        require 'shoryuken/extensions/active_job_adapter' if defined?(::ActiveJob)
+        require 'shoryuken/extensions/active_job_adapter' if Shoryuken.active_job?
         require File.expand_path('config/environment.rb')
       end
     end
 
     def merge_cli_defined_queues
-      cli_defined_queues = options.delete(:queues) || []
+      cli_defined_queues = options[:queues].to_a
 
       cli_defined_queues.each do |cli_defined_queue|
-        Shoryuken.options[:queues].delete_if do |config_file_queue|
-          config_file_queue[0] == cli_defined_queue[0]
-        end
+        # CLI defined queues override config_file defined queues
+        Shoryuken.options[:queues].delete_if { |config_file_queue| config_file_queue[0] == cli_defined_queue[0] }
 
         Shoryuken.options[:queues] << cli_defined_queue
       end
     end
 
-    def prefix_active_job_queue_names
-      return unless defined? ::ActiveJob
-      return unless Shoryuken.active_job_queue_name_prefixing
-
+    def prefix_active_job_queue_name(queue_name, weight)
       queue_name_prefix = ::ActiveJob::Base.queue_name_prefix
       queue_name_delimiter = ::ActiveJob::Base.queue_name_delimiter
 
       # See https://github.com/rails/rails/blob/master/activejob/lib/active_job/queue_name.rb#L27
+      name_parts = [queue_name_prefix.presence, queue_name]
+      prefixed_queue_name = name_parts.compact.join(queue_name_delimiter)
+      [prefixed_queue_name, weight]
+    end
+
+    def prefix_active_job_queue_names
+      return unless Shoryuken.active_job?
+      return unless Shoryuken.active_job_queue_name_prefixing
+
       Shoryuken.options[:queues].to_a.map! do |queue_name, weight|
-        name_parts = [queue_name_prefix.presence, queue_name]
-        prefixed_queue_name = name_parts.compact.join(queue_name_delimiter)
-        [prefixed_queue_name, weight]
+        prefix_active_job_queue_name(queue_name, weight)
       end
-    end
 
-    def parse_queue(queue, weight = nil)
-      [weight.to_i, 1].max.times { Shoryuken.queues << queue }
-    end
+      Shoryuken.options[:groups].to_a.map! do |group, options|
+        if options[:queues]
+          options[:queues].map! do |queue_name, weight|
+            prefix_active_job_queue_name(queue_name, weight)
+          end
+        end
 
-    def parse_queues
-      Shoryuken.options[:queues].to_a.each do |queue_and_weight|
-        parse_queue(*queue_and_weight)
+        [group, options]
       end
     end
 
-    def patch_deprecated_workers
-      Shoryuken.worker_registry.queues.each do |queue|
-        Shoryuken.worker_registry.workers(queue).each do |worker_class|
-          if worker_class.instance_method(:perform).arity == 1
-            Shoryuken.logger.warn { "[DEPRECATION] #{worker_class.name}#perform(sqs_msg) is deprecated. Please use #{worker_class.name}#perform(sqs_msg, body)" }
+    def parse_queue(queue, weight, group)
+      Shoryuken.add_queue(queue, [weight.to_i, 1].max, group)
+    end
 
-            worker_class.class_eval do
-              alias_method :deprecated_perform, :perform
+    def parse_queues
+      if Shoryuken.options[:queues].to_a.any?
+        Shoryuken.add_group('default', Shoryuken.options.fetch(:concurrency, 25))
 
-              def perform(sqs_msg, body = nil)
-                deprecated_perform(sqs_msg)
-              end
-            end
-          end
+        Shoryuken.options[:queues].to_a.each do |queue, weight|
+          parse_queue(queue, weight, 'default')
+        end
+      end
+
+      Shoryuken.options[:groups].to_a.each do |group, options|
+        Shoryuken.add_group(group, options.fetch(:concurrency, 25))
+        options[:queues].to_a.each do |queue, weight|
+          parse_queue(queue, weight, group)
         end
       end
     end
@@ -150,29 +148,34 @@ module Shoryuken
     end
 
     def validate_queues
-      Shoryuken.logger.warn { 'No queues supplied' } if Shoryuken.queues.empty?
+      return Shoryuken.logger.warn { 'No queues supplied' } if Shoryuken.ungrouped_queues.empty?
 
       non_existent_queues = []
 
-      Shoryuken.queues.uniq.each do |queue|
+      Shoryuken.ungrouped_queues.uniq.each do |queue|
         begin
-          Shoryuken::Client.queues queue
-        rescue Aws::SQS::Errors::NonExistentQueue
+          Shoryuken::Client.queues(queue)
+        rescue Aws::Errors::NoSuchEndpointError, Aws::SQS::Errors::NonExistentQueue
           non_existent_queues << queue
         end
       end
 
-      fail ArgumentError, "The specified queue(s) #{non_existent_queues} do not exist" if non_existent_queues.any?
+      return if non_existent_queues.none?
+
+      fail(
+        ArgumentError,
+        "The specified queue(s) #{non_existent_queues.join(', ')} do not exist.\nTry 'shoryuken sqs create QUEUE-NAME' for creating a queue with default settings"
+      )
     end
 
     def validate_workers
-      all_queues = Shoryuken.queues
+      return if Shoryuken.active_job?
+
+      all_queues = Shoryuken.ungrouped_queues
       queues_with_workers = Shoryuken.worker_registry.queues
 
-      unless defined?(::ActiveJob)
-        (all_queues - queues_with_workers).each do |queue|
-          Shoryuken.logger.warn { "No worker supplied for '#{queue}'" }
-        end
+      (all_queues - queues_with_workers).each do |queue|
+        Shoryuken.logger.warn { "No worker supplied for #{queue}" }
       end
     end
   end

data/lib/shoryuken/fetcher.rb

@@ -1,73 +1,43 @@
 module Shoryuken
   class Fetcher
-    include Celluloid
     include Util
 
     FETCH_LIMIT = 10
 
-    def initialize(manager)
-      @manager = manager
-    end
-
-    def receive_messages(queue, limit)
-      # AWS limits the batch size by 10
-      limit = limit > FETCH_LIMIT ? FETCH_LIMIT : limit
+    attr_reader :group
 
-      options = (Shoryuken::AwsConfig.options[:receive_message] || {}).dup
-      options[:max_number_of_messages] = limit
-      options[:message_attribute_names] = %w(All)
-      options[:attribute_names] = %w(All)
-
-      Shoryuken::Client.queues(queue).receive_messages options
+    def initialize(group)
+      @group = group
     end
 
-    def fetch(queue, available_processors)
-      watchdog('Fetcher#fetch died') do
-        started_at = Time.now
-
-        logger.debug { "Looking for new messages in '#{queue}'" }
+    def fetch(queue, limit)
+      started_at = Time.now
 
-        begin
-          batch = Shoryuken.worker_registry.batch_receive_messages?(queue)
-          limit = batch ? FETCH_LIMIT : available_processors
+      logger.debug { "Looking for new messages in #{queue}" }
 
-          if (sqs_msgs = Array(receive_messages(queue, limit))).any?
-            logger.debug { "Found #{sqs_msgs.size} messages for '#{queue}'" }
+      sqs_msgs = Array(receive_messages(queue, [FETCH_LIMIT, limit].min))
 
-            if batch
-              @manager.async.assign(queue, patch_sqs_msgs!(sqs_msgs))
-            else
-              sqs_msgs.each { |sqs_msg| @manager.async.assign(queue, sqs_msg) }
-            end
+      logger.info { "Found #{sqs_msgs.size} messages for #{queue.name}" } unless sqs_msgs.empty?
+      logger.debug { "Fetcher for #{queue} completed in #{elapsed(started_at)} ms" }
 
-            @manager.async.rebalance_queue_weight!(queue)
-          else
-            logger.debug { "No message found for '#{queue}'" }
-
-            @manager.async.pause_queue!(queue)
-          end
+      sqs_msgs
+    end
 
-          logger.debug { "Fetcher for '#{queue}' completed in #{elapsed(started_at)} ms" }
-        rescue => ex
-          logger.error { "Error fetching message: #{ex}" }
-          logger.error { ex.backtrace.first }
-        end
+    private
 
-        @manager.async.dispatch
-      end
+    def receive_messages(queue, limit)
+      # AWS limits the batch size by 10
+      limit = limit > FETCH_LIMIT ? FETCH_LIMIT : limit
 
-    end
+      options = Shoryuken.sqs_client_receive_message_opts[group].to_h.dup
 
-    private
+      options[:max_number_of_messages]  = limit
+      options[:message_attribute_names] = %w(All)
+      options[:attribute_names]         = %w(All)
 
-    def patch_sqs_msgs!(sqs_msgs)
-      sqs_msgs.instance_eval do
-        def message_id
-          "batch-with-#{size}-messages"
-        end
-      end
+      options.merge!(queue.options)
 
-      sqs_msgs
+      Shoryuken::Client.queues(queue.name).receive_messages(options)
     end
   end
 end