tom_queue 0.0.1.dev

data/spec/tom_queue/queue_manager_spec.rb

@@ -0,0 +1,218 @@
+require 'tom_queue/helper'
+
+describe TomQueue::QueueManager do
+
+  let(:manager) { TomQueue::QueueManager.new("test-#{Time.now.to_f}") }
+  let(:channel) { TomQueue.bunny.create_channel }
+
+  describe "basic creation" do
+  
+    it "should be a thing" do
+      defined?(TomQueue::QueueManager).should be_true
+    end  
+
+    it "should be created with a name-prefix" do
+      manager.prefix.should =~ /^test-[\d.]+$/
+    end
+
+    it "should default the prefix to TomQueue.default_prefix if available" do
+      TomQueue.default_prefix = "test-#{Time.now.to_f}"
+      TomQueue::QueueManager.new.prefix.should == TomQueue.default_prefix
+    end
+
+    it "should raise an ArgumentError if no prefix is specified and no default is available" do
+      TomQueue.default_prefix = nil
+      lambda {
+        TomQueue::QueueManager.new
+      }.should raise_exception(ArgumentError, /prefix is required/)
+    end
+
+    it "should use the TomQueue.bunny object" do
+      manager.bunny.should == TomQueue.bunny
+    end
+
+    it "should stick to the same bunny object, even if TomQueue.bunny changes" do
+      manager
+      TomQueue.bunny = "A FAKE RABBIT"
+      manager.bunny.should be_a(Bunny::Session)
+    end
+  end
+
+  describe "AMQP configuration" do
+
+    TomQueue::PRIORITIES.each do |priority|
+      it "should create a queue for '#{priority}' priority" do
+        manager.queues[priority].name.should == "#{manager.prefix}.balance.#{priority}"
+        # Declare the queue, if the parameters don't match the brokers existing channel, then bunny will throw an
+        # exception.
+        channel.queue("#{manager.prefix}.balance.#{priority}", :durable => true, :auto_delete => false, :exclusive => false)
+      end
+    end
+
+    it "should create a single durable topic exchange" do
+      manager.exchange.name.should == "#{manager.prefix}.work"
+      # Now we declare it again on the broker, which will raise an exception if the parameters don't match
+      channel.topic("#{manager.prefix}.work", :durable => true, :auto_delete => false)
+    end
+
+  end
+
+  describe "QueueManager message publishing" do
+
+    it "should forward the payload directly" do
+      manager.publish("foobar")
+      manager.pop.ack!.payload.should == "foobar"
+    end
+
+    it "should return nil" do
+      manager.publish("some work").should be_nil
+    end
+
+    it "should raise an exception if the payload isn't a string" do
+      lambda {
+        manager.publish({"some" => {"structured_data" => true}})
+      }.should raise_exception(ArgumentError, /must be a string/)
+    end
+
+    describe "deferred execution" do
+
+      it "should allow a run-at time to be specified" do
+        manager.publish("future", :run_at => Time.now + 2.2)
+      end
+
+      it "should throw an ArgumentError exception if :run_at isn't a Time object" do
+        lambda { 
+          manager.publish("future", :run_at => "around 10pm ?")
+        }.should raise_exception(ArgumentError, /must be a Time object/)
+      end
+
+      it "should write the run_at time in the message headers as an ISO-8601 timestamp, with 4-digits of decimal precision" do
+        execution_time = Time.now - 1.0
+        manager.publish("future", :run_at => execution_time)
+        manager.pop.ack!.headers[:headers]['run_at'].should == execution_time.iso8601(4)
+      end
+
+      it "should default to :run_at the current time" do
+        manager.publish("future")
+        future_time = Time.now
+        Time.parse(manager.pop.ack!.headers[:headers]['run_at']).should < future_time
+      end
+    end
+
+    describe "message priorities" do
+      it "should have an array of priorities, in the correct order" do
+        TomQueue::PRIORITIES.should be_a(Array)
+        TomQueue::PRIORITIES.should == [
+          TomQueue::HIGH_PRIORITY,
+          TomQueue::NORMAL_PRIORITY,
+          TomQueue::LOW_PRIORITY,
+          TomQueue::BULK_PRIORITY
+        ]
+      end
+
+      it "should allow the message priority to be set" do
+        manager.publish("foobar", :priority => TomQueue::BULK_PRIORITY)
+      end
+
+      it "should throw an ArgumentError if an unknown priority value is used" do
+        lambda {
+          manager.publish("foobar", :priority => "VERY BLOODY IMPORTANT")
+        }.should raise_exception(ArgumentError, /unknown priority level/)
+      end
+
+      it "should write the priority in the message header as 'job_priority'" do
+        manager.publish("foobar", :priority => TomQueue::BULK_PRIORITY)
+        manager.pop.ack!.headers[:headers]['job_priority'].should == TomQueue::BULK_PRIORITY
+      end
+
+      it "should default to normal priority" do
+        manager.publish("foobar")
+        manager.pop.ack!.headers[:headers]['job_priority'].should == TomQueue::NORMAL_PRIORITY
+      end
+    end
+
+    TomQueue::PRIORITIES.each do |priority|
+      it "should publish #{priority} priority messages to the single exchange, with routing key set to '#{priority}'" do
+        manager.publish("foo", :priority => priority)
+        manager.pop.ack!.response.tap do |resp|
+          resp.exchange.should == "#{manager.prefix}.work"
+          resp.routing_key.should == priority
+        end
+      end
+    end
+
+  end
+
+
+  describe "QueueManager - deferred message handling" do
+
+    describe "when popping a message" do
+      it "should ensure a deferred manager with the same prefix is running" do
+        manager.publish("work")
+        TomQueue::DeferredWorkManager.instance(manager.prefix).should_receive(:ensure_running)
+        manager.pop
+      end
+    end
+
+    describe "when publishing a deferred message" do
+      it "should not publish to the normal AMQP queue" do
+        manager.publish("work", :run_at => Time.now + 0.1)
+        manager.queues.values.find { |q| channel.basic_get(q.name).first }.should be_nil
+      end
+      it "should call #handle_deferred on the appropriate deferred work manager" do
+        TomQueue::DeferredWorkManager.instance(manager.prefix).should_receive(:handle_deferred)
+        manager.publish("work", :run_at => Time.now + 0.1)
+      end
+      it "should pass the original payload" do
+        TomQueue::DeferredWorkManager.instance(manager.prefix).should_receive(:handle_deferred).with("work", anything)
+        manager.publish("work", :run_at => Time.now + 0.1)
+      end
+      it "should pass the original options" do
+        run_time = Time.now + 0.5
+        TomQueue::DeferredWorkManager.instance(manager.prefix).should_receive(:handle_deferred).with(anything, hash_including(:priority => TomQueue::NORMAL_PRIORITY, :run_at => run_time))
+        manager.publish("work", :run_at => run_time)
+      end
+    end
+
+  end
+
+  describe "QueueManager#pop - work popping" do
+    before do
+      manager.publish("foo")
+      manager.publish("bar")
+    end
+
+    it "should not have setup a consumer before the first call" do
+      manager.queues.values.each do |queue|
+        queue.status[:consumer_count].should == 0
+      end
+    end
+
+    it "should not leave any running consumers for immediate messages" do
+      manager.pop.ack!
+      manager.queues.values.each do |queue|
+        queue.status[:consumer_count].should == 0
+      end
+    end
+
+    it "should not leave any running consumers after it has waited for a message " do
+      manager.pop.ack!
+      manager.pop.ack!
+      Thread.new { sleep 0.1; manager.publish("baz") }
+      manager.pop.ack!
+      manager.queues.values.each do |queue|
+        queue.status[:consumer_count].should == 0
+      end
+    end
+
+    it "should return a QueueManager::Work instance" do
+      manager.pop.ack!.should be_a(TomQueue::Work)
+    end
+
+    it "should return the message at the head of the queue" do
+      manager.pop.ack!.payload.should == "foo"
+      manager.pop.ack!.payload.should == "bar"
+    end
+  end
+
+end

data/spec/tom_queue/sorted_array_spec.rb

@@ -0,0 +1,160 @@
+require 'tom_queue/helper'
+
+describe Range, 'tomqueue_binary_search' do
+
+  it "should return nil for an empty range" do
+    (0...0).tomqueue_binary_search.should be_nil
+  end
+
+  describe "for a single item range" do
+    let(:range) { 5...6 }
+
+    it "should yield the index" do
+      range.tomqueue_binary_search do |index|
+        @index = index
+      end
+      @index.should == 5
+    end
+
+    it "should return 0 if the yield returned -1" do
+      range.tomqueue_binary_search { |index| -1 }.should == 5
+    end
+
+    it "should return 1 if the yield returned +1" do
+      range.tomqueue_binary_search { |index| +1 }.should == 6
+    end
+
+    it "should return 0 if the yield returned 0" do
+      range.tomqueue_binary_search { |index| 0 }.should == 5
+    end
+  end
+
+  describe "for two item range" do
+    let(:range) { 7..8 }
+    
+    it "should yield the lower number" do
+      range.tomqueue_binary_search do |index|
+        @index = index
+        0
+      end
+      @index.should == 7
+    end
+
+    it "should return the lower number if the block returns -1" do
+      range.tomqueue_binary_search { |i| -1 }.should == 7
+    end
+    it "should return the lower number if the block returns 0" do
+      range.tomqueue_binary_search { |i| 0 }.should == 7
+    end
+
+    it "should yield the second number if the block returns +1" do
+      range.tomqueue_binary_search do |i| 
+        if i == 7
+          1
+        elsif i == 8
+          @yielded = true
+          0
+        end
+      end
+      @yielded.should be_true
+    end
+  end
+
+  describe "for a three item range" do
+    let(:range) { 7..9 }
+
+    it "should yield the mid-point" do
+      range.tomqueue_binary_search do |index|
+        @index = index
+        0
+      end
+      @index.should == 8
+    end
+
+    it "should return the mid-point if the block returns 0" do
+      range.tomqueue_binary_search { |index| 0 }.should == 8
+    end
+
+    it "should recurse to the right on +1" do
+      @yielded = []
+      range.tomqueue_binary_search { |index| @yielded << index; 1 }.should == 10
+      @yielded.should == [8,9]
+    end
+
+    it "should recurse to the left on -1" do
+      @yielded = []
+      range.tomqueue_binary_search { |index| @yielded << index; -1 }.should == 7
+      @yielded.should == [8,7]
+    end
+
+  end
+
+  describe "acceptance 1" do
+    let(:range) { 0...100 }
+    let(:value) { 43 }
+
+    before do
+      @yielded = []
+      @result = range.tomqueue_binary_search { |i| @yielded << i; value <=> i }
+    end
+
+    it "should get the correct result" do
+      @result.should == value
+    end
+
+    it "should yield the correct values" do
+      @yielded.should == [49, 24, 36, 42, 45, 43]
+    end
+  end
+
+  describe "acceptance 2" do
+    let(:range) { 0..3 }
+    let(:value) { 3 }
+
+    before do
+      @yielded = []
+      @result = range.tomqueue_binary_search { |i| @yielded << i; value <=> i }
+    end
+
+    it "should get the correct result" do
+      @result.should == value
+    end
+
+    it "should yield the correct values" do
+      @yielded.should == [1,2,3]
+    end
+  end
+
+
+end
+
+
+
+describe TomQueue::SortedArray do
+
+  let(:array) { TomQueue::SortedArray.new }
+
+  it "should insert in sorted order" do
+    array << 4
+    array << 5
+    array << 2
+    array << 1
+    array << 3
+    array.should == [1,2,3,4,5]
+  end
+
+  it "should work for all permutations of insertion" do
+    numbers = [0,1,2,3,4,5,6]
+    numbers.permutation.each do |permutation|
+      array = TomQueue::SortedArray.new
+      permutation.each do |i|
+        array << i
+      end
+      array.should == numbers
+    end
+  end
+
+  it "should return itself when inserting" do
+    (array << 3).should == array
+  end
+end

data/spec/tom_queue/tom_queue_integration_spec.rb

@@ -0,0 +1,296 @@
+
+require 'net/http'
+require 'tom_queue/helper'
+
+describe TomQueue::QueueManager, "simple publish / pop" do
+
+  let(:manager) { TomQueue::QueueManager.new("test-#{Time.now.to_f}", 'manager') }
+  let(:consumer) { TomQueue::QueueManager.new(manager.prefix, 'consumer1') }
+  let(:consumer2) { TomQueue::QueueManager.new(manager.prefix, 'consumer2') }
+
+  it "should pop a previously published message" do
+    manager.publish('some work')
+    manager.pop.payload.should == 'some work'
+  end
+
+  it "should block on #pop until work is published" do
+    manager
+
+    Thread.new do
+      sleep 0.1
+      manager.publish('some work')
+    end
+
+    consumer.pop.payload.should == 'some work'
+  end
+
+  it "should work between objects (hello, rabbitmq)" do
+    manager.publish "work"
+    consumer.pop.payload.should == "work"
+  end
+
+  it "should load-balance work between multiple consumers" do
+    manager.publish "foo"
+    manager.publish "bar"
+
+    consumer.pop.payload.should == "foo"
+    consumer2.pop.payload.should == "bar"
+  end
+
+  it "should work for more than one message!" do
+    input, output = [], []
+    (0..9).each do |i|
+      input << i.to_s
+      manager.publish i.to_s
+    end
+
+    (input.size / 2).times do 
+      a = consumer.pop
+      b = consumer2.pop
+      output << a.ack!.payload
+      output << b.ack!.payload
+    end
+    output.should == input
+  end
+
+  it "should not drop messages when two different priorities arrive" do
+    manager.publish("1", :priority => TomQueue::BULK_PRIORITY)
+    manager.publish("2", :priority => TomQueue::NORMAL_PRIORITY)
+    manager.publish("3", :priority => TomQueue::HIGH_PRIORITY)
+    out = []
+    out << consumer.pop.ack!.payload
+    out << consumer.pop.ack!.payload
+    out << consumer.pop.ack!.payload
+    out.sort.should == ["1", "2", "3"]
+  end
+
+  it "should handle priority queueing, maintaining per-priority FIFO ordering" do
+    manager.publish("1", :priority => TomQueue::BULK_PRIORITY) 
+    manager.publish("2", :priority => TomQueue::NORMAL_PRIORITY)
+    manager.publish("3", :priority => TomQueue::HIGH_PRIORITY)
+
+    # 1,2,3 in the queue - 3 wins as it's highest priority
+    consumer.pop.ack!.payload.should == "3"
+
+    manager.publish("4", :priority => TomQueue::NORMAL_PRIORITY)
+
+    # 1,2,4 in the queue - 2 wins as it's highest (NORMAL) and first in    
+    consumer.pop.ack!.payload.should == "2"
+    
+    manager.publish("5", :priority => TomQueue::BULK_PRIORITY)
+
+    # 1,4,5 in the queue - we'd expect 4 (highest), 1 (first bulk), 5 (second bulk)
+    consumer.pop.ack!.payload.should == "4"
+    consumer.pop.ack!.payload.should == "1"
+    consumer.pop.ack!.payload.should == "5"
+  end
+
+  it "should handle priority queueing across two consumers" do
+    manager.publish("1", :priority => TomQueue::BULK_PRIORITY) 
+    manager.publish("2", :priority => TomQueue::HIGH_PRIORITY)
+    manager.publish("3", :priority => TomQueue::NORMAL_PRIORITY)
+    
+
+    # 1,2,3 in the queue - 3 wins as it's highest priority
+    order = []
+    order << consumer.pop.ack!.payload
+
+    manager.publish("4", :priority => TomQueue::NORMAL_PRIORITY)
+
+    # 1,2,4 in the queue - 2 wins as it's highest (NORMAL) and first in    
+    order << consumer.pop.ack!.payload
+
+    manager.publish("5", :priority => TomQueue::BULK_PRIORITY)
+
+    # 1,4,5 in the queue - we'd expect 4 (highest), 1 (first bulk), 5 (second bulk)
+    order << consumer.pop.ack!.payload
+    order << consumer2.pop.ack!.payload
+    order << consumer.pop.ack!.payload
+
+    order.should == ["2","3","4","1","5"]
+  end
+
+  it "should immediately run a high priority task, when there are lots of bulks" do
+    100.times do |i|
+      manager.publish("stuff #{i}", :priority => TomQueue::BULK_PRIORITY)
+    end
+
+    consumer.pop.ack!.payload.should == "stuff 0"
+    consumer2.pop.ack!.payload.should == "stuff 1"
+    consumer.pop.payload.should == "stuff 2"
+
+    manager.publish("HIGH1", :priority => TomQueue::HIGH_PRIORITY)
+    manager.publish("NORMAL1", :priority => TomQueue::NORMAL_PRIORITY)
+    manager.publish("HIGH2", :priority => TomQueue::HIGH_PRIORITY)
+    manager.publish("NORMAL2", :priority => TomQueue::NORMAL_PRIORITY)
+
+    consumer.pop.ack!.payload.should == "HIGH1"
+    consumer.pop.ack!.payload.should == "HIGH2"
+    consumer2.pop.ack!.payload.should == "NORMAL1"
+    consumer.pop.ack!.payload.should == "NORMAL2"
+
+    consumer2.pop.ack!.payload.should == "stuff 3"
+  end
+
+  it "should allow a message to be deferred for future execution" do
+    execution_time = Time.now + 0.2
+    manager.publish("future-work", :run_at => execution_time )
+
+    consumer.pop.ack!
+    Time.now.to_f.should > execution_time.to_f
+  end
+
+  describe "slow tests", :timeout => 100 do
+
+    class QueueConsumerThread
+
+      class WorkObject < Struct.new(:payload, :received_at, :run_at)
+        def <=>(other)
+          self.received_at.to_f <=> other.received_at.to_f
+        end
+      end
+
+      attr_reader :work, :thread
+
+      def initialize(manager, &work_proc)
+        @manager = manager
+        @work_proc = work_proc
+        @work = []
+      end
+
+      def thread_main
+        loop do
+          begin
+            work = @manager.pop
+            recv_time = Time.now
+
+            Thread.exit if work.payload == "done"
+            
+            work_obj = WorkObject.new(work.payload, recv_time, Time.parse(work.headers[:headers]['run_at']))
+            @work << work_obj 
+            @work_proc && @work_proc.call(work_obj)
+
+            work.ack!
+          rescue
+            p $!
+          end
+        end
+
+      end
+
+      def signal_shutdown(time=nil)
+        time ||= Time.now
+        @manager.publish("done", :run_at => time)
+      end
+      def start!
+        @thread ||= Thread.new(&method(:thread_main))
+        self
+      end
+    end
+
+
+    it "should work with lots of messages, without dropping and deliver FIFO" do
+      source_order = []
+
+      # Run both consumers, in parallel threads, so in some cases, 
+      # there should be a thread waiting for work
+      consumers = 16.times.collect do |i|
+        consumer = TomQueue::QueueManager.new(manager.prefix, "thread-#{i}")
+        QueueConsumerThread.new(consumer) { |work| sleep rand(0.5) }.start!
+      end
+
+      # Now publish some work
+      250.times do |i|
+        work = "work #{i}"
+        source_order << work
+        manager.publish(work)
+      end
+        
+      # Now publish a bunch of messages to cause the threads to exit the loop
+      consumers.each { |c| c.signal_shutdown }
+      consumers.each { |c| c.thread.join }
+
+      # Merge the list of received work together sorted by received time,
+      # and compare to the source list
+      sink_order = consumers.map { |c| c.work }.flatten.sort.map { |a| a.payload }
+
+      source_order.should == sink_order
+    end
+
+    it "should be able to drain the queue, block and resume when new work arrives" do
+      source_order = []
+
+      # Run both consumers, in parallel threads, so in some cases, 
+      # there should be a thread waiting for work
+      consumers = 10.times.collect do |i|
+        consumer = TomQueue::QueueManager.new(manager.prefix, "thread-#{i}")
+        QueueConsumerThread.new(consumer) { |work| sleep rand(0.5) }.start!
+      end
+
+      # This sleep gives the workers enough time to block on the first call to pop
+      sleep 0.1 until manager.queues[TomQueue::NORMAL_PRIORITY].status[:consumer_count] == consumers.size
+
+      # Now publish some work
+      50.times do |i|
+        work = "work #{i}"
+        source_order << work
+        manager.publish(work)
+      end
+
+      # Rough and ready - wait for the queue to empty
+      sleep 0.1 until manager.queues[TomQueue::NORMAL_PRIORITY].status[:message_count] == 0
+
+      # Now publish some more work
+      50.times do |i|
+        work = "work 2-#{i}"
+        source_order << work
+        manager.publish(work)
+      end
+
+      # Now publish a bunch of messages to cause the threads to exit the loop
+      consumers.each { |c| c.signal_shutdown }
+      consumers.each { |c| c.thread.join }
+
+      # Now merge all the consumers internal work arrays into one
+      # sorted by the received_at timestamps
+      sink_order = consumers.map { |c| c.work }.flatten.sort.map { |a| a.payload }
+
+      # Compare what the publisher did to what the workers did.
+      sink_order.should == source_order
+    end
+
+    it "should work with lots of deferred work on the queue, and still schedule all messages" do
+      # sit in a loop to pop it all off again
+      consumers = 5.times.collect do |i|
+        consumer = TomQueue::QueueManager.new(manager.prefix, "thread-#{i}")
+        QueueConsumerThread.new(consumer).start!
+      end
+
+      # Generate some work
+      max_run_at = Time.now
+      200.times do |i| 
+        run_at = Time.now + (rand * 6.0)
+        max_run_at = [max_run_at, run_at].max
+        manager.publish(JSON.dump(:id => i), :run_at => run_at)
+      end
+
+      # Shutdown the consumers again
+      consumers.each do |c|
+        c.signal_shutdown(max_run_at + 1.0)
+      end
+      consumers.each { |c| c.thread.join }
+
+      # Now make sure none of the messages were delivered too late!
+      total_size = 0
+      consumers.each do |c|
+        total_size += c.work.size
+        c.work.each do |work|
+          work.received_at.should < (work.run_at + 1.0)
+        end
+      end
+
+      total_size.should == 200
+    end
+  end
+
+end