introduce a multi_queue object

Author: hone

lib/resque.rb

@@ -13,6 +13,7 @@
 require 'resque/worker'
 require 'resque/plugin'
 require 'resque/queue'
+require 'resque/multi_queue'
 require 'resque/coder'
 require 'resque/multi_json_coder'
 

lib/resque/multi_queue.rb

@@ -0,0 +1,56 @@
+require 'redis'
+require 'redis-namespace'
+require 'thread'
+require 'mutex_m'
+
+module Resque
+  ###
+  # Holds multiple queues, allowing you to pop the first available job
+  class MultiQueue
+    include Mutex_m
+
+    ###
+    # Create a new MultiQueue using the +queues+ from the +redis+ connection
+    def initialize(queues, redis)
+      super()
+
+      @queues = {}
+      @redis  = redis
+
+      queues.each do |queue|
+        key = @redis.is_a?(Redis::Namespace) ? "#{@redis.namespace}:" : ""
+        key += queue.redis_name
+        @queues[key] = queue
+      end
+    end
+
+    # Pop an item off one of the queues.  This method will block until an item
+    # is available.
+    #
+    # Pass +true+ for a non-blocking pop.  If nothing is read on a non-blocking
+    # pop, a ThreadError is raised.
+    def pop(non_block = false)
+      if non_block
+        synchronize do
+          value = nil
+
+          @queues.values.each do |queue|
+            begin
+              return queue.pop(true)
+            rescue ThreadError
+            end
+          end
+
+          raise ThreadError
+        end
+      else
+        queue_names = @queues.values.map {|queue| queue.redis_name }
+        synchronize do
+          value = @redis.blpop(*(queue_names + [1])) until value
+          queue_name, payload = value
+          @queues[queue_name].coder.decode(payload)
+        end
+      end
+    end
+  end
+end

lib/resque/queue.rb

@@ -9,7 +9,7 @@ module Resque
   class Queue
     include Mutex_m
 
-    attr_reader :name, :redis_name
+    attr_reader :name, :redis_name, :coder
 
     ###
     # Create a new Queue object with +name+ on +redis+ connection, and using

test/multi_queue_test.rb

@@ -0,0 +1,40 @@
+require "test_helper"
+
+describe "Resque::MulitQueue" do
+  let(:redis) { Resque.redis }
+  let(:coder) { Resque::MultiJsonCoder.new }
+
+  before do
+    redis.flushall
+  end
+
+  it "blocks on pop" do
+    foo   = Resque::Queue.new 'foo', redis, coder
+    bar   = Resque::Queue.new 'bar', redis, coder
+    queue = Resque::MultiQueue.new([foo, bar], redis)
+    t     = Thread.new { queue.pop }
+
+    job = Resque::Job.new(:bar, :class => 'GoodJob')
+    bar << job
+    assert_equal coder.decode(job.to_json), t.join.value
+  end
+
+  it "nonblocking pop works" do
+    foo   = Resque::Queue.new 'foo', redis, coder
+    bar   = Resque::Queue.new 'bar', redis, coder
+    queue = Resque::MultiQueue.new([foo, bar], redis)
+
+    job = Resque::Job.new(:bar, :class => 'GoodJob')
+    bar << job
+    assert_equal coder.decode(job.to_json), queue.pop(true)
+  end
+
+  it "blocks forever on pop" do
+    foo   = Resque::Queue.new 'foo', redis, coder
+    bar   = Resque::Queue.new 'bar', redis, coder
+    queue = Resque::MultiQueue.new([foo, bar], redis)
+    assert_raises Timeout::Error do
+      Timeout::timeout(2) { queue.pop }
+    end
+  end
+end