Implement Client.push_batch

lib/sidekiq/client.rb

@@ -36,33 +36,52 @@ def self.registered_queues
     #   Sidekiq::Client.push('queue' => 'my_queue', 'class' => MyWorker, 'args' => ['foo', 1, :bat => 'bar'])
     #
     def self.push(item)
-      raise(ArgumentError, "Message must be a Hash of the form: { 'class' => SomeWorker, 'args' => ['bob', 1, :foo => 'bar'] }") unless item.is_a?(Hash)
-      raise(ArgumentError, "Message must include a class and set of arguments: #{item.inspect}") if !item['class'] || !item['args']
-      raise(ArgumentError, "Message must include a Sidekiq::Worker class, not class name: #{item['class'].ancestors.inspect}") if !item['class'].is_a?(Class) || !item['class'].respond_to?('get_sidekiq_options')
+      normed = normalize_item(item)
+      normed, payload = process_single(item['class'], normed)
 
-      worker_class = item['class']
-      item['class'] = item['class'].to_s
+      pushed = false
+      Sidekiq.redis do |conn|
+        if normed['at']
+          pushed = conn.zadd('schedule', normed['at'].to_s, payload)
+        else
+          _, pushed = conn.multi do
+            conn.sadd('queues', normed['queue'])
+            conn.rpush("queue:#{normed['queue']}", payload)
+          end
+        end
+      end
+      pushed ? normed['jid'] : nil
+    end
 
-      item = worker_class.get_sidekiq_options.merge(item)
-      item['retry'] = !!item['retry']
-      queue = item['queue']
-      item['jid'] = SecureRandom.hex(12)
+    ##
+    # Push a large number of jobs to Redis.  In practice this method is only
+    # useful if you are pushing tens of thousands of jobs or more.  This method
+    # basically cuts down on the redis round trip latency.
+    #
+    # Takes the same arguments as Client.push except that args is expected to be
+    # an Array of Arrays.  All other keys are duplicated for each job.  Each job
+    # is run through the client middleware pipeline and each job gets its own Job ID
+    # as normal.
+    #
+    # Returns the number of jobs pushed or nil if the pushed failed.  The number of jobs
+    # pushed can be less than the number given if the middleware stopped processing for one
+    # or more jobs.
+    def self.push_batch(items)
+      normed = normalize_item(items)
+      payloads = items['args'].map do |args|
+        _, payload = process_single(items['class'], normed.merge('args' => args, 'jid' => SecureRandom.hex(12)))
+        payload
+      end.compact
 
       pushed = false
-      Sidekiq.client_middleware.invoke(worker_class, item, queue) do
-        payload = Sidekiq.dump_json(item)
-        Sidekiq.redis do |conn|
-          if item['at']
-            pushed = conn.zadd('schedule', item['at'].to_s, payload)
-          else
-            _, pushed = conn.multi do
-              conn.sadd('queues', queue)
-              conn.rpush("queue:#{queue}", payload)
-            end
-          end
+      Sidekiq.redis do |conn|
+        _, pushed = conn.multi do
+          conn.sadd('queues', normed['queue'])
+          conn.rpush("queue:#{normed['queue']}", payloads)
         end
       end
-      pushed ? item['jid'] : nil
+
+      pushed ? payloads.size : nil
     end
 
     # Resque compatibility helpers.
@@ -82,5 +101,30 @@ def self.enqueue(klass, *args)
     def self.enqueue_to(queue, klass, *args)
       klass.client_push('queue' => queue, 'class' => klass, 'args' => args)
     end
+
+    private
+
+    def self.process_single(worker_class, item)
+      queue = item['queue']
+
+      Sidekiq.client_middleware.invoke(worker_class, item, queue) do
+        payload = Sidekiq.dump_json(item)
+        return item, payload
+      end
+    end
+
+    def self.normalize_item(item)
+      raise(ArgumentError, "Message must be a Hash of the form: { 'class' => SomeWorker, 'args' => ['bob', 1, :foo => 'bar'] }") unless item.is_a?(Hash)
+      raise(ArgumentError, "Message must include a class and set of arguments: #{item.inspect}") if !item['class'] || !item['args']
+      raise(ArgumentError, "Message must include a Sidekiq::Worker class, not class name: #{item['class'].ancestors.inspect}") if !item['class'].is_a?(Class) || !item['class'].respond_to?('get_sidekiq_options')
+
+      normalized_item = item.dup
+      normalized_item['class'] = normalized_item['class'].to_s
+      normalized_item['retry'] = !!normalized_item['retry']
+      normalized_item['jid'] = SecureRandom.hex(12)
+
+      item['class'].get_sidekiq_options.merge normalized_item
+    end
+
   end
 end

test/test_client.rb

@@ -99,6 +99,14 @@ class QueuedWorker
     end
   end
 
+  describe 'batch' do
+    it 'can push a large set of jobs at once' do
+      a = Time.now
+      count = Sidekiq::Client.push_batch('class' => QueuedWorker, 'args' => (1..1_000).to_a.map { |x| Array(x) })
+      assert_equal 1_000, count
+    end
+  end
+
   class BaseWorker
     include Sidekiq::Worker
     sidekiq_options 'retry' => 'base'