famoseagle-sweat_shop 0.3.0 → 0.5.0

data/README.markdown

@@ -1,14 +1,14 @@
 # SweatShop
 
 SweatShop provides an api to background resource intensive tasks. Much of the api design was copied from Workling, with a few tweaks.
-Currently, it runs kestrel, but it can support any number of queues.
+Currently, it runs rabbitmq and kestrel, but it can support any number of queues.
 
 ## Installing 
 
-   gem install sweat_shop
-   freeze in your rails directory
-   cd vendor/gems/sweat_shop
-   rake setup
+    gem install sweat_shop
+    freeze in your rails directory
+    cd vendor/gems/sweat_shop
+    rake setup
 
 ## Writing workers
 
@@ -36,11 +36,15 @@ queues.
 
 ## Running the queue
 
-SweatShop has been tested with Kestrel, but it will also work with Starling. You can install and start kestrel following the instructions here:
+SweatShop has been tested with Rabbit and Kestrel, but it will also work with Starling. Please use the following resources to install the server:
 
+Kestrel:
 http://github.com/robey/kestrel/tree/master
 
-config/sweatshop.yml specifies the machine address of the queue (default localhost:22133).
+Rabbit:
+http://github.com/ezmobius/nanite/tree/master
+
+config/sweatshop.yml specifies the machine address of the queue (default localhost:5672). You can also specify the queue type with the queue param.
 
 ## Running the workers
 
@@ -59,6 +63,7 @@ If you would like to run SweatShop as a daemon on a linux machine, use the initd
 
     i_can_daemonize
     memcache (for kestrel)
+    carrot (for rabbit)
 
 # LICENSE
 

data/VERSION.yml

@@ -1,4 +1,4 @@
 --- 
 :patch: 0
 :major: 0
-:minor: 3
+:minor: 5

data/config/defaults.yml

@@ -1,4 +1,5 @@
 # default options
 servers:
-  - localhost:22133
+  - localhost:5672
 enable: true
+queue: rabbit

data/config/sweatshop.yml

@@ -1,12 +1,12 @@
 development:
   servers:
-    - localhost:22133
+    - localhost:5672
   enable: true
 test:
   servers:
-    - localhost:22133
-  enable: true
+    - localhost:5672
+  enable: false
 production:
   servers:
-    - localhost:22133
+    - localhost:5672
   enable: true

data/lib/message_queue/base.rb

@@ -0,0 +1,17 @@
+module MessageQueue
+  class Base
+    attr_reader :servers
+    def queue_size(queue);    end
+    def enqueue(queue, data); end
+    def dequeue(queue);       end
+    def confirm(queue);       end
+    def subscribe(queue);     end
+    def delete(queue);        end
+    def client;               end
+    def stop;                 end
+
+    def subscribe?
+      false
+    end
+  end
+end

data/lib/message_queue/kestrel.rb

@@ -0,0 +1,34 @@
+module MessageQueue
+  class Kestrel < Base
+    attr_reader :client, :servers
+
+    def initialize(opts)
+      @servers = opts[:servers]
+    end
+
+    def queue_size(queue)
+      size  = 0
+      stats = client.stats
+      servers.each do |server|
+        size += stats[server]["queue_#{queue}_items"].to_i
+      end
+      size
+    end
+
+    def enqueue(queue, data)
+      client.set(queue, data)
+    end
+
+    def dequeue(queue)
+      client.get("#{queue}/open")
+    end
+
+    def confirm(queue)
+      client.get("#{queue}/close")
+    end
+
+    def client
+      @client ||= MemCache.new(servers) 
+    end
+  end
+end

data/lib/message_queue/rabbit-async.rb

@@ -0,0 +1,83 @@
+require 'mq'
+module MessageQueue
+  class Rabbit < Base
+    attr_accessor :em_thread
+
+    def initialize(opts={})
+      @servers = opts[:servers]
+      @info = {}
+      @host, @port = @servers.first.split(':')
+      @port = @port.to_i
+    end
+
+    def queue_size(queue)
+      num = 0
+      client.queue(queue).status{|messages, consumers| num = messages}
+      num
+    end
+
+    def enqueue(queue, data)
+      client.queue(queue, :durable => true).publish(Marshal.dump(data), :persistent => true)
+    end
+
+    def dequeue(queue)
+      client.queue(queue).pop do |info, task|
+        @info[queue] = info
+        return Marshal.load(task)
+      end
+    end
+
+    def confirm(queue)
+      if @info[queue] 
+        @info[queue].ack
+        @info[queue] = nil
+      end
+    end
+
+    def client
+      @client ||= begin 
+        start_em
+        if servers
+          MQ.new(AMQP.connect(:host => @host, :port => @port)) 
+        else
+          MQ.new
+        end
+      end
+    end
+
+    def start_em
+      if em_thread.nil? and not EM.reactor_running?
+        self.em_thread = Thread.new{EM.run}
+        ['INT', 'TERM'].each do |sig|
+          old = trap(sig) do 
+            stop
+            old.call
+          end
+        end
+      end
+    end
+
+    def subscribe?
+      true
+    end
+
+    def subscribe(queue, &block)
+      AMQP.start(:host => @host, :port => @port) do
+        mq = MQ.new
+        mq.send(AMQP::Protocol::Basic::Qos.new(:prefetch_size => 0, :prefetch_count => 1, :global => false))
+        mq.queue(queue, :durable => true).subscribe(:ack => true) do |info, task|
+          if task
+            @info[queue] = info
+            task = Marshal.load(task)
+            block.call(task)
+          end
+        end
+      end
+    end
+
+    def stop
+      em_thread.join(0.15) unless em_thread.nil?
+      AMQP.stop{ EM.stop }
+    end
+  end
+end

data/lib/message_queue/rabbit.rb

@@ -0,0 +1,42 @@
+require 'carrot'
+module MessageQueue
+  class Rabbit < Base
+
+    def initialize(opts={})
+      @servers = opts[:servers]
+      @info = {}
+      @host, @port = @servers.first.split(':')
+      @port = @port.to_i
+    end
+
+    def delete(queue)
+      client.queue(queue).delete
+    end
+
+    def queue_size(queue)
+      client.queue(queue).message_count
+    end
+
+    def enqueue(queue, data)
+      client.queue(queue, :durable => true).publish(Marshal.dump(data), :persistent => true)
+    end
+
+    def dequeue(queue)
+      task = client.queue(queue).pop(:ack => true)
+      return unless task
+      Marshal.load(task)
+    end
+
+    def confirm(queue)
+      client.queue(queue).ack
+    end
+
+    def client
+      @client ||= Carrot.new(:host => @host, :port => @port) 
+    end
+
+    def stop
+      client.stop
+    end
+  end
+end

data/lib/sweat_shop.rb

@@ -3,7 +3,9 @@ require 'digest'
 require 'yaml'
 
 $:.unshift(File.dirname(__FILE__))
-require 'kestrel'
+require 'message_queue/base'
+require 'message_queue/rabbit'
+require 'message_queue/kestrel'
 require 'sweat_shop/worker'
 
 module SweatShop
@@ -29,16 +31,27 @@ module SweatShop
   end
 
   def do_tasks(workers)
-    loop do
-      wait = true
-      workers.each do |worker|
-        if task = worker.dequeue
-          worker.do_task(task)
-          wait = false
+    if queue.subscribe?
+      EM.run do
+        workers.each do |worker|
+          worker.subscribe
         end
       end
-      exit if stop?
-      sleep 1 if wait
+    else
+      loop do
+        wait = true
+        workers.each do |worker|
+          if task = worker.dequeue
+            worker.do_task(task)
+            wait = false
+          end
+        end
+        if stop?
+          queue.stop
+          exit 
+        end
+        sleep 1 if wait
+      end
     end
   end
 
@@ -54,14 +67,6 @@ module SweatShop
     )
   end
 
-  def stop
-    @stop = true
-  end
-
-  def stop?
-    @stop
-  end
-
   def config
     @config ||= begin
       defaults = YAML.load_file(File.dirname(__FILE__) + '/../config/defaults.yml')
@@ -79,6 +84,15 @@ module SweatShop
     end
   end
 
+  def stop
+    @stop = true
+    queue.stop if queue.subscribe?
+  end
+
+  def stop?
+    @stop
+  end
+
   def queue_sizes
     workers.inject([]) do |all, worker|
       all << [worker, worker.queue_size]
@@ -94,12 +108,33 @@ module SweatShop
   end
 
   def queue
-    @queue ||= Kestrel.new(:servers => config['servers'])
+    @queue ||= begin 
+      queue = config['queue'] || 'rabbit'
+      queue = constantize("MessageQueue::#{queue.capitalize}")
+      queue.new(:servers => config['servers'])
+    end
   end
 
   def queue=(queue)
     @queue = queue
   end
+
+  def log(msg)
+    return if logger == :silent
+    logger ? logger.debug(msg) : puts(msg)
+  end
+
+  def logger
+    @logger
+  end
+
+  def logger=(logger)
+    @logger = logger
+  end
+
+  def constantize(str)
+    Object.module_eval("#{str}", __FILE__, __LINE__)
+  end
 end
 
 if defined?(RAILS_ROOT)

data/lib/sweat_shop/worker.rb

@@ -2,8 +2,6 @@ require File.dirname(__FILE__) + '/metaid'
 
 module SweatShop
   class Worker
-    @@logger = nil
-
     def self.inherited(subclass)
       self.workers << subclass
     end
@@ -42,6 +40,10 @@ module SweatShop
       @queue_name ||= self.to_s
     end
 
+    def self.delete_queue
+      queue.delete(queue_name)
+    end
+
     def self.queue_size
       queue.queue_size(queue_name)
     end
@@ -57,6 +59,12 @@ module SweatShop
     def self.confirm
       queue.confirm(queue_name)
     end
+
+    def self.subscribe
+      queue.subscribe(queue_name) do |task|
+        do_task(task)
+      end
+    end
     
     def self.do_tasks
       while task = dequeue
@@ -85,21 +93,20 @@ module SweatShop
       after_task.call(task) if after_task
     end
 
-    def self.workers
-      SweatShop.workers
+    def self.queue
+      SweatShop.queue
     end
 
-    def self.log(msg)
-      return if logger == :silent
-      logger ? logger.debug(msg) : puts(msg)
+    def self.workers
+      SweatShop.workers
     end
 
-    def self.logger
-      @@logger
+    def self.config
+      SweatShop.config
     end
 
-    def self.logger=(logger)
-      @@logger = logger
+    def self.log(msg)
+      SweatShop.log(msg)
     end
 
     def self.before_task(&block)
@@ -118,10 +125,6 @@ module SweatShop
       end
     end
 
-    def self.queue
-      SweatShop.queue
-    end
-
     def self.queue_group(group=nil)
       group ? meta_def(:_queue_group){ group } : _queue_group
     end

data/test/test_functional_worker.rb

@@ -12,11 +12,10 @@ class WorkerTest < Test::Unit::TestCase
     File.delete(HelloWorker::TEST_FILE) if File.exist?(HelloWorker::TEST_FILE)
   end
 
-  # remove 'x' and start kestrel to run
   test "daemon" do
     begin
       SweatShop.queue = nil
-      SweatShop::Worker.logger = :silent
+      SweatShop.logger = :silent
 
       worker = File.expand_path(File.dirname(__FILE__) + '/hello_worker')
       sweatd = "#{File.dirname(__FILE__)}/../lib/sweat_shop/sweatd.rb" 
@@ -27,8 +26,10 @@ class WorkerTest < Test::Unit::TestCase
 
       File.delete('sweatd.log') if File.exist?('sweatd.log')
       assert_equal 'Hi, Amos', File.read(HelloWorker::TEST_FILE)
-    rescue MemCache::MemCacheError => e
-      puts "\n\n*** Start kestrel on localhost to run all functional tests. ***\n\n"
+    rescue Exception => e
+      puts e.message
+      puts e.backtrace.join("\n")
+      fail "\n\n*** Functional test failed, is the rabbit server running on localhost? ***\n"
     end
   end
   

data/test/test_sweatshop.rb

@@ -1,10 +1,8 @@
-require File.dirname(__FILE__) + '/../../../memcache/lib/memcache_mock'
 require File.dirname(__FILE__) + '/test_helper'
 require File.dirname(__FILE__) + '/../lib/sweat_shop'
 
 class SweatShopTest < Test::Unit::TestCase
   SweatShop.workers = []
-  SweatShop.queue   = Kestrel.new(:client => MemCacheMock.new)
 
   class HelloWorker < SweatShop::Worker
     def hello(name)
@@ -28,7 +26,7 @@ class SweatShopTest < Test::Unit::TestCase
   end
 
   test "uid" do
-    SweatShop::Worker.logger = :silent
+    SweatShop.logger = :silent
     uid = HelloWorker.async_hello('Amos')
     assert_not_nil uid
   end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification 
 name: famoseagle-sweat_shop
 version: !ruby/object:Gem::Version 
-  version: 0.3.0
+  version: 0.5.0
 platform: ruby
 authors: 
 - Amos Elliston
@@ -9,7 +9,7 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2009-03-18 00:00:00 -07:00
+date: 2009-04-12 00:00:00 -07:00
 default_executable: 
 dependencies: []
 
@@ -27,7 +27,11 @@ files:
 - Rakefile
 - README.markdown
 - VERSION.yml
-- lib/kestrel.rb
+- lib/message_queue
+- lib/message_queue/base.rb
+- lib/message_queue/kestrel.rb
+- lib/message_queue/rabbit-async.rb
+- lib/message_queue/rabbit.rb
 - lib/sweat_shop
 - lib/sweat_shop/metaid.rb
 - lib/sweat_shop/sweatd.rb

data/lib/kestrel.rb

@@ -1,29 +0,0 @@
-class Kestrel
-  attr_reader :client, :servers
-
-  def initialize(opts)
-    @servers = opts[:servers]
-    @client  = opts[:client] || MemCache.new(@servers) 
-  end
-
-  def queue_size(queue)
-    size  = 0
-    stats = client.stats
-    servers.each do |server|
-      size += stats[server]["queue_#{queue}_items"].to_i
-    end
-    size
-  end
-
-  def enqueue(queue, data)
-    client.set(queue, data)
-  end
-
-  def dequeue(queue)
-    client.get("#{queue}/open")
-  end
-
-  def confirm(queue)
-    client.get("#{queue}/close")
-  end
-end