sweatshop 1.4.0

data/History.txt

@@ -0,0 +1,6 @@
+=== 1.0.0 / 2009-01-13
+
+* 1 major enhancement
+
+  * Birthday!
+

data/LICENSE

@@ -0,0 +1,20 @@
+Copyright (c) 2009 Amos Elliston
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

data/README.markdown

@@ -0,0 +1,103 @@
+# 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 rabbitmq and kestrel, but it can support any number of queues.
+
+## Installing
+
+    gem install sweat_shop
+    freeze in your gems directory (add config.gem 'sweat_shop' to your environment)
+    cd vendor/gems/sweat_shop
+    rake setup
+
+## Writing workers
+
+Put `email_worker.rb` into app/workers and sublcass `SweatShop::Worker`:
+
+    class EmailWorker < SweatShop::Worker
+      def send_mail(to)
+        user = User.find_by_id(to)
+        Mailer.deliver_welcome(to)
+      end
+    end
+
+Then, anywhere in your app you can execute:
+
+    EmailWorker.async_send_mail(1)
+
+The `async` signifies that this task will be placed on a queue to be serviced by the EmailWorker possibly on another machine. You can also
+call:
+
+    EmailWorker.send_mail(1)
+
+That will do the work immediately, without placing the task on the queue. You can also define a `queue_group` at the top of the file
+which will allow you to split workers out into logical groups. This is important if you have various machines serving different
+queues.
+
+## Running the queue
+
+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
+
+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.
+
+## Rabbit cluster support
+
+The following example configuration shows support for Rabbit clusters
+within a queue group:
+
+        default:
+          queue: rabbit
+          cluster:
+            - hostA:5672
+            - hostB:5672
+          user: 'guest'
+          pass: 'guest'
+          vhost: '/'
+        enable: true
+
+Sweatshop will attempt to connect to each server listed under
+"cluster" in order, until it either manages to establish a connection
+or until it runs out of servers.
+
+If you only have a single Rabbit server, you can omit the "cluster"
+option and just add "host" and "port" (or host: localhost:5672) options, as shown below:
+
+       default:
+         queue: rabbit
+         host: localhost
+         port: 5672
+         user: 'guest'
+         pass: 'guest'
+         vhost: '/'
+       enable: true
+
+
+## Running the workers
+
+Assuming you ran `rake setup` in Rails, you can type:
+
+    script/sweatshop
+
+By default, the script will run all workers defined in the app/workers dir. Every task will be processed on each queue using a round-robin algorithm. You can also add the `-d` flag which will put the worker in daemon mode. The daemon also takes other params.  Add a `-h` for more details.
+
+    script/sweatshop -d
+    script/sweatshop -d stop
+
+If you would like to run SweatShop as a daemon on a linux machine, use the initd.sh script provided in the sweat_shop/script dir.
+
+# REQUIREMENTS
+
+    memcache (for kestrel)
+    carrot (for rabbit)
+
+# LICENSE
+
+Copyright (c) 2009 Amos Elliston, Geni.com; Published under The MIT License, see License

data/Rakefile

@@ -0,0 +1,43 @@
+require 'rake'
+require 'rake/testtask'
+require 'rake/rdoctask'
+require 'rcov/rcovtask'
+
+begin
+  require 'jeweler'
+  Jeweler::Tasks.new do |s|
+    s.name = "sweatshop"
+    s.summary = %Q{Sweatshop is a simple asynchronous worker queue build on top of rabbitmq/ampq}
+    s.email = "amos@geni.com"
+    s.homepage = "http://github.com/famoseagle/sweat-shop"
+    s.description = "See summary"
+    s.authors = ["Amos Elliston"]
+    s.files =  FileList["[A-Z]*", "install.rb", "{lib,test,config,script}/**/*"]
+    s.add_dependency('famoseagle-carrot', '= 0.7.0')
+  end
+  Jeweler::GemcutterTasks.new
+rescue LoadError
+  puts "Jeweler not available. Install it with: sudo gem install technicalpickles-jeweler -s http://gems.github.com"
+end
+
+Rake::TestTask.new
+
+Rake::RDocTask.new do |rdoc|
+  rdoc.rdoc_dir = 'rdoc'
+  rdoc.title    = 'new_project'
+  rdoc.options << '--line-numbers' << '--inline-source'
+  rdoc.rdoc_files.include('README*')
+  rdoc.rdoc_files.include('lib/**/*.rb')
+end
+
+Rcov::RcovTask.new do |t|
+  t.libs << 'test'
+  t.test_files = FileList['test/**/*_test.rb']
+  t.verbose = true
+end
+
+task :default => :test
+
+task :setup do
+  require File.dirname(__FILE__) + '/install'
+end

data/VERSION.yml

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

data/bin/sweatd

@@ -0,0 +1,3 @@
+#!/usr/bin/env ruby
+
+require File.dirname(__FILE__) + '/../lib/sweat_shop/sweatd'

data/config/defaults.yml

@@ -0,0 +1,8 @@
+# default options
+default: 
+  queue: rabbit
+  host: localhost:5672
+  user: 'guest'
+  pass: 'guest'
+  vhost: '/'
+enable: true

data/config/sweatshop.yml

@@ -0,0 +1,15 @@
+development:
+  default:
+    queue: rabbit
+    host: localhost:5672
+  enable: true
+test:
+  default:
+    queue: rabbit
+    host: localhost:5672
+  enable: true
+production:
+  default:
+    queue: rabbit
+    host: localhost:5672
+  enable: true

data/install.rb

@@ -0,0 +1,20 @@
+require 'fileutils'
+root = File.dirname(__FILE__)
+
+unless defined?(RAILS_ROOT)
+  search_paths = %W{/.. /../.. /../../..}
+  search_paths.each do |path|
+    if File.exist?(root + path + '/config/environment.rb') 
+      RAILS_ROOT = root + path 
+      break
+    end
+  end
+end
+return unless RAILS_ROOT
+
+config = 'sweatshop.yml'
+script = 'sweatshop'
+
+FileUtils.cp(File.join(root, 'config', config), File.join(RAILS_ROOT, 'config', config))
+FileUtils.cp(File.join(root, 'script', script), File.join(RAILS_ROOT, 'script', script))
+puts "\n\ninstalled #{ [config, script].join(", ") } \n\n"

data/lib/message_queue/base.rb

@@ -0,0 +1,17 @@
+module MessageQueue
+  class Base
+    attr_reader :opts
+    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,32 @@
+module MessageQueue
+  class Kestrel < Base
+    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.rb

@@ -0,0 +1,108 @@
+require 'carrot'
+module MessageQueue
+  class Rabbit < Base
+
+    def initialize(opts={})
+      @opts = opts
+    end
+
+    def delete(queue)
+      send_command do
+        client.queue(queue).delete
+      end
+    end
+
+    def queue_size(queue)
+      send_command do
+        client.queue(queue).message_count
+      end
+    end
+
+    def enqueue(queue, data)
+      send_command do
+        client.queue(queue, :durable => true).publish(Marshal.dump(data), :persistent => true)
+      end
+    end
+
+    def dequeue(queue)
+      send_command do
+        task = client.queue(queue).pop(:ack => true)
+        return unless task
+        Marshal.load(task)
+      end
+    end
+
+    def confirm(queue)
+      send_command do
+        client.queue(queue).ack
+      end
+    end
+
+    def send_command(&block)
+      retried = false
+      begin
+        block.call
+      rescue Carrot::AMQP::Server::ServerDown => e
+        if not retried
+          puts "Error #{e.message}. Retrying..."
+          @client = nil
+          retried = true
+          retry
+        else
+          raise e
+        end
+      end
+    end
+
+    def client
+      return @client if @client
+
+      if @opts['cluster']
+        @opts['cluster'].each_with_index do |server, i|
+          host, port = server.split(':')
+          begin
+            @client = Carrot.new(
+              :host   => host,
+              :port   => port.to_i,
+              :user   => @opts['user'],
+              :pass   => @opts['pass'],
+              :vhost  => @opts['vhost'],
+              :insist => @opts['insist']
+            )
+            return @client
+          rescue Carrot::AMQP::Server::ServerDown => e
+            if i == (@opts['cluster'].size-1)
+              raise e
+            else
+              next
+            end
+          end
+        end
+      else
+        if @opts['host'] =~ /:/
+          host, port = @opts['host'].split(':')
+        else
+          host = @opts['host']
+          port = @opts['port']
+        end
+        @client = Carrot.new(
+          :host   => host,
+          :port   => port.to_i,
+          :user   => @opts['user'],
+          :pass   => @opts['pass'],
+          :vhost  => @opts['vhost'],
+          :insist => @opts['insist']
+        )
+      end
+      @client
+    end
+
+    def client=(client)
+      @client = client
+    end
+
+    def stop
+      client.stop
+    end
+  end
+end

data/lib/sweat_shop.rb

@@ -0,0 +1,179 @@
+require 'rubygems'
+require 'digest'
+require 'yaml'
+
+$:.unshift(File.dirname(__FILE__))
+require 'message_queue/base'
+require 'message_queue/rabbit'
+require 'message_queue/kestrel'
+require 'sweat_shop/worker'
+
+module SweatShop
+  extend self
+
+  def workers
+    @workers ||= []
+  end
+
+  def workers=(workers)
+    @workers = workers 
+  end
+
+  def workers_in_group(groups)
+    groups = [groups] unless groups.is_a?(Array)
+    if groups.include?(:all)
+      workers
+    else
+      workers.select do |worker|
+        groups.include?(worker.queue_group)
+      end
+    end
+  end
+
+  def do_tasks(workers)
+    if queue.subscribe?
+      EM.run do
+        workers.each do |worker|
+          worker.subscribe
+        end
+      end
+    else
+      loop do
+        wait = true
+        workers.each do |worker|
+          if task = worker.dequeue
+            worker.do_task(task)
+            wait = false
+          end
+        end
+        if stop?
+          workers.each do |worker|
+            worker.stop
+          end
+          queue.stop
+          exit 
+        end
+        sleep 1 if wait
+      end
+    end
+  end
+
+  def do_all_tasks
+    do_tasks(
+      workers_in_group(:all)
+    )
+  end
+
+  def do_default_tasks
+    do_tasks(
+      workers_in_group(:default)
+    )
+  end
+
+  def config
+    @config ||= begin
+      defaults = YAML.load_file(File.dirname(__FILE__) + '/../config/defaults.yml')
+      if defined?(RAILS_ROOT)
+        file = RAILS_ROOT + '/config/sweatshop.yml'
+        if File.exist?(file)
+          YAML.load_file(file)[RAILS_ENV || 'development']
+        else
+          defaults['enable'] = false
+          defaults
+        end
+      else
+        defaults
+      end
+    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]
+      all
+    end
+  end
+
+  def queue(type = 'default')
+    type = config[type] ? type : 'default'
+    return queues[type] if queues[type]
+
+    qconfig = config[type]
+    qtype   = qconfig['queue'] || 'rabbit'
+    queue   = constantize("MessageQueue::#{qtype.capitalize}")
+
+    queues[type] = queue.new(qconfig)
+  end
+
+  def queue=(queue, type = 'default')
+    queues[type] = queue
+  end
+
+  def queues
+    @queues ||= {}
+  end
+
+  def queue_groups
+    @queue_groups ||= workers.collect{|w| w.queue_group} << 'default'
+  end
+
+  def pp_sizes
+    max_width = workers.collect{|w| w.to_s.size}.max
+    puts '-' * (max_width + 10)
+    puts queue_sizes.collect{ |p| sprintf("%-#{max_width}s %2s", p.first, p.last) }.join("\n")
+    puts '-' * (max_width + 10)
+  end
+
+  def cluster_info
+    servers = []
+    queue_groups.each do |group|
+      qconfig = config[group]
+      next unless qconfig
+      next unless qconfig['cluster']
+      servers << qconfig['cluster']
+    end
+    servers.flatten!
+
+    servers.each do |server|
+      puts "\nQueue sizes on #{server}"
+      queue = MessageQueue::Rabbit.new('host' => server)
+      queue_groups.each do |group|
+        queues[group] = queue
+      end
+      pp_sizes
+      puts
+    end
+    @queues = {}
+    nil
+  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)
+  Dir.glob(RAILS_ROOT + '/app/workers/*.rb').each{|worker| require worker }
+end