netinlet-sweat_shop 1.1.0

data/lib/sweat_shop/metaid.rb

@@ -0,0 +1,5 @@
+def metaclass; class << self; self; end; end
+def meta_eval(&blk); metaclass.instance_eval(&blk); end
+def meta_def(name, &blk)
+  meta_eval { define_method name, &blk }
+end

data/lib/sweat_shop/sweatd.rb

@@ -0,0 +1,76 @@
+require File.dirname(__FILE__) + '/daemoned'
+
+module SweatShop
+  class Sweatd
+    include Daemoned
+    queues     = []
+    groups     = []
+    rails_root = nil
+    start_cmd  = "ruby #{__FILE__} start #{ARGV.reject{|a| a == 'start'}.join(' ')}"
+
+    arg '--workers=Worker,Worker', 'Workers to service (Default is all)' do |value|
+      queues = value.split(',')
+    end
+
+    arg '--groups=GROUP,GROUP', 'Groups of queues to service' do |value|
+      groups = value.split(',').collect{|g| g.to_sym}
+    end
+
+    arg '--worker-file=WORKERFILE', 'Worker file to load'  do |value|
+      require value
+    end
+
+    arg '--worker-dir=WORKERDIR', 'Directory containing workers'  do |value|
+      Dir.glob(value + '*.rb').each{|worker| require worker}
+    end
+
+    arg '--rails=DIR', 'Pass in RAILS_ROOT to run this daemon in a rails environment' do |value|
+      rails_root = value
+    end
+
+    sig(:term, :int) do
+      puts "Shutting down sweatd..."
+      SweatShop.stop
+    end
+
+    sig(:hup) do
+      puts "Received HUP"
+      SweatShop.stop
+      remove_pid!
+      puts "Restarting sweatd with #{start_cmd}..."
+      `#{start_cmd}`        
+    end
+    
+    before do
+      if rails_root
+        puts "Loading Rails..."
+        require rails_root + '/config/environment' 
+      end
+      require File.dirname(__FILE__) + '/../sweat_shop'
+    end
+
+    daemonize(:kill_timeout => 20) do
+      workers = []
+
+      if groups.any?
+        workers += SweatShop.workers_in_group(groups)
+      end
+
+      if queues.any?
+        workers += queues.collect{|q| Object.module_eval(q)}
+      end
+
+      if workers.any?
+        worker_str = workers.join(',')
+        puts "Starting #{worker_str}..." 
+        $0 = "Sweatd: #{worker_str}"
+        SweatShop.do_tasks(workers)
+      else
+        puts "Starting all workers..." 
+        $0 = 'Sweatd: all'
+        SweatShop.do_all_tasks
+      end
+    end
+
+  end
+end

data/lib/sweat_shop/worker.rb

@@ -0,0 +1,152 @@
+require File.dirname(__FILE__) + '/metaid'
+require File.dirname(__FILE__) + '/../serializers/serializer'
+
+module SweatShop
+  class Worker
+    def self.inherited(subclass)
+      self.workers << subclass
+    end
+
+    def self.method_missing(method, *args, &block)
+      if method.to_s =~ /^async_(.*)/
+        method        = $1
+        expected_args = instance.method(method).arity
+        if expected_args != args.size 
+          raise ArgumentError.new("#{method} expects #{expected_args} arguments")
+        end
+        #return instance.send(method, *args) unless true == config['enable']
+
+        uid  = ::Digest::MD5.hexdigest("#{name}:#{method}:#{args}:#{Time.now.to_f}")
+        task = {:args => args, :method => method, :uid => uid, :queued_at => Time.now.to_i}
+
+        log("Putting #{uid} on #{queue_name}")
+        enqueue(task) if config['enable']
+
+        uid
+      elsif instance.respond_to?(method)
+        instance.send(method, *args)
+      else
+        super
+      end
+    end
+    
+    def self.serialize_with(serializer_name)
+      @serializer = SweatShop::Serializer.serializers[serializer_name]
+      raise RuntimeError.new("No Such Serializer ['#{serializer_name}']") if @serializer.nil?
+    end
+    
+    def self.serializer
+      @serializer ||= SweatShop::Serializer.default
+    end
+
+    def self.instance
+      @instance ||= new
+    end
+
+    def self.config
+      SweatShop.config
+    end
+
+    def self.queue_name
+      @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
+
+    def self.enqueue(task)
+      queue.enqueue(queue_name, task, serializer)
+    end
+
+    def self.dequeue
+      queue.dequeue(queue_name, serializer)
+    end
+
+    def self.confirm
+      queue.confirm(queue_name)
+    end
+
+    def self.subscribe
+      queue.subscribe(queue_name, serializer) do |task|
+        do_task(task)
+      end
+    end
+    
+    def self.do_tasks
+      while task = dequeue
+        do_task(task)
+      end
+    end
+
+    def self.do_task(task)
+      call_before_task(task)
+
+      queued_at = task[:queued_at] ? "(queued #{Time.at(task[:queued_at]).strftime('%Y/%m/%d %H:%M:%S')})" : ''
+      log("Dequeuing #{queue_name}::#{task[:method]} #{queued_at}")
+      task[:result] = instance.send(task[:method], *task[:args])
+
+      call_after_task(task)
+      confirm
+    end
+
+    def self.call_before_task(task)
+      superclass.call_before_task(task) if superclass.respond_to?(:call_before_task)
+      before_task.call(task) if before_task
+    end
+
+    def self.call_after_task(task)
+      superclass.call_after_task(task) if superclass.respond_to?(:call_after_task)
+      after_task.call(task) if after_task
+    end
+
+    def self.queue
+      SweatShop.queue(queue_group.to_s)
+    end
+
+    def self.workers
+      SweatShop.workers
+    end
+
+    def self.config
+      SweatShop.config
+    end
+
+    def self.log(msg)
+      SweatShop.log(msg)
+    end
+
+    def self.before_task(&block)
+      if block
+        @before_task = block
+      else
+        @before_task
+      end
+    end
+
+    def self.after_task(&block)
+      if block
+        @after_task = block
+      else
+        @after_task
+      end
+    end
+
+    def self.stop
+      instance.stop
+    end
+
+    # called before we exit -- subclass can implement this method 
+    def stop; end;
+
+
+    def self.queue_group(group=nil)
+      group ? meta_def(:_queue_group){ group } : _queue_group
+    end
+    queue_group :default
+  end
+end

data/test/hello_worker.rb

@@ -0,0 +1,13 @@
+require File.dirname(__FILE__) + '/../lib/sweat_shop'
+class HelloWorker < SweatShop::Worker
+  TEST_FILE = File.dirname(__FILE__) + '/test.txt' unless defined?(TEST_FILE)
+
+  def hello(name)
+    puts name
+    "Hi, #{name}"
+  end
+
+  after_task do |task|
+    File.open(TEST_FILE, 'w'){|f| f << task[:result]}
+  end
+end

data/test/test_functional_worker.rb

@@ -0,0 +1,35 @@
+require File.dirname(__FILE__) + '/../lib/sweat_shop'
+require File.dirname(__FILE__) + '/test_helper'
+require File.dirname(__FILE__) + '/hello_worker'
+class WorkerTest < Test::Unit::TestCase
+
+  def setup
+    File.delete(HelloWorker::TEST_FILE) if File.exist?(HelloWorker::TEST_FILE)
+  end
+
+  def teardown
+    File.delete(HelloWorker::TEST_FILE) if File.exist?(HelloWorker::TEST_FILE)
+  end
+
+  test "daemon" do
+    begin
+      SweatShop.queue = nil
+      SweatShop.logger = :silent
+
+      worker = File.expand_path(File.dirname(__FILE__) + '/hello_worker')
+      sweatd = "#{File.dirname(__FILE__)}/../lib/sweat_shop/sweatd.rb" 
+      uid = HelloWorker.async_hello('Amos')
+
+      `ruby #{sweatd} --worker-file #{worker} start`
+      `ruby #{sweatd} stop`
+
+      File.delete('sweatd.log') if File.exist?('sweatd.log')
+      assert_equal 'Hi, Amos', File.read(HelloWorker::TEST_FILE)
+    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
+  
+end

data/test/test_helper.rb

@@ -0,0 +1,14 @@
+require 'test/unit'
+
+class << Test::Unit::TestCase
+  def test(name, &block)
+    test_name = "test_#{name.gsub(/[\s\W]/,'_')}"
+    raise ArgumentError, "#{test_name} is already defined" if self.instance_methods.include? test_name
+    define_method test_name, &block
+  end
+
+  def xtest(name, &block)
+    # no-op, an empty test method is defined to prevent "no tests in testcase" errors when all tests are disabled
+    define_method(:test_disabled) { assert true }
+  end
+end

data/test/test_json_serializer.rb

@@ -0,0 +1,50 @@
+require File.dirname(__FILE__) + '/test_helper'
+require File.dirname(__FILE__) + '/../lib/sweat_shop'
+
+class JsonSerializerTest < Test::Unit::TestCase
+  class UnderTest
+    attr_accessor :name, :address, :city, :state, :zip, :id
+    def ==(other)
+      ![:name, :address, :city, :state, :zip, :id].map{|a| return self.send(a) == other.send(a)}.include?(false)
+    end
+    
+    def self.json_create(params)
+      obj = new
+      for key, value in params
+        next if key == 'json_class'
+        obj.instance_variable_set "@#{key}", value
+      end
+      obj
+    end
+  end
+
+  def setup
+    @under_test = UnderTest.new
+    @under_test.id = 87
+    @under_test.name = "JsonTest"
+    @under_test.address = "555 Rock Ridge Road"
+    @under_test.city = "Rock Ridge"
+    @under_test.state = "Texas"
+    @under_test.zip = "90210"
+  end
+  
+  test "should properly serialize a simple data structure" do
+    dump = SweatShop::Serializers::JsonSerializer.serialize([{:foo => "bar"}, 23, 87, %w{doug cathy connor tommy}])
+    assert_equal '[{"foo":"bar"},23,87,["doug","cathy","connor","tommy"]]', dump
+  end
+
+  
+  test "should be able to serialize an arbitrary class" do
+    dump = SweatShop::Serializers::JsonSerializer.serialize(@under_test)
+    assert_equal @under_test, JSON.parse(dump)  # deserialize the dump because hash keys are getting set in different order from time to time & making tests fail
+  end
+  
+  test "should be able to deserialize a simple data structure" do
+    assert_equal [{"foo" => "bar"}, 23, 87, %w{doug cathy connor tommy}], SweatShop::Serializers::JsonSerializer.deserialize('[{"foo":"bar"},23,87,["doug","cathy","connor","tommy"]]')
+  end
+  
+  test "should be able to deserialize an arbitrary class" do
+    assert_equal @under_test, SweatShop::Serializers::JsonSerializer.deserialize('{"city":"Rock Ridge","name":"JsonTest","zip":"90210","json_class":"JsonSerializerTest::UnderTest","id":87,"address":"555 Rock Ridge Road","state":"Texas"}')
+  end
+
+end

data/test/test_marshal_serializer.rb

@@ -0,0 +1,42 @@
+require File.dirname(__FILE__) + '/test_helper'
+require File.dirname(__FILE__) + '/../lib/sweat_shop'
+
+class MarshalSerializerTest < Test::Unit::TestCase
+
+  class UnderTest
+    attr_accessor :name, :address, :city, :state, :zip, :id
+    def ==(other)
+      ![:name, :address, :city, :state, :zip, :id].map{|a| return self.send(a) == other.send(a)}.include?(false)
+    end
+    
+  end
+
+  def setup
+    @under_test = UnderTest.new
+    @under_test.id = 87
+    @under_test.name = "MarshalTest"
+    @under_test.address = "555 Rock Ridge Road"
+    @under_test.city = "Rock Ridge"
+    @under_test.state = "Texas"
+    @under_test.zip = "90210"
+  end
+
+  test "should properly serialize a simple data structure" do
+    dump = SweatShop::Serializers::MarshalSerializer.serialize([{:foo => "bar"}, 23, 87, %w{doug cathy connor tommy}])
+    assert_equal Marshal.dump([{:foo => "bar"}, 23, 87, %w{doug cathy connor tommy}]), dump
+  end
+  
+  test "should be able to serialize an arbitrary class" do
+    dump = SweatShop::Serializers::MarshalSerializer.serialize(@under_test)
+    assert_equal Marshal.dump(@under_test), dump
+  end
+  
+  test "should be able to deserialize a simple data structure" do
+    assert_equal [{:foo => "bar"}, 23, 87, %w{doug cathy connor tommy}], SweatShop::Serializers::MarshalSerializer.deserialize(Marshal.dump([{:foo => "bar"}, 23, 87, %w{doug cathy connor tommy}]))
+  end
+  
+  test "should be able to deserialize an arbitrary class" do
+    assert_equal @under_test, SweatShop::Serializers::MarshalSerializer.deserialize(Marshal.dump(@under_test))
+  end
+
+end

data/test/test_serializer.rb

@@ -0,0 +1,95 @@
+require File.dirname(__FILE__) + '/test_helper'
+require File.dirname(__FILE__) + '/../lib/sweat_shop'
+
+class SerializerTest < Test::Unit::TestCase
+
+  # no serializer specified
+  class DefaultSerialWorker < SweatShop::Worker
+    def hello(name)
+      "Hi, #{name}"
+    end
+  end
+
+  # Marshall specified serializer
+  class MarshalWorker < SweatShop::Worker
+    serialize_with :marshal
+    def hello(name)
+      "Hi, #{name}"
+    end
+  end
+
+  # JSON specified serializer
+  class JsonWorker < SweatShop::Worker
+    serialize_with :json
+    def hello(name)
+      "Hi, #{name}"
+    end
+  end
+  
+  # YAML specified serializer
+  class YamlWorker < SweatShop::Worker
+    serialize_with :yaml
+    def hello(name)
+      "Hi, #{name}"
+    end
+  end
+  
+  class NoNameSerializer < SweatShop::Serializer
+    def self.serialize(payload)
+      "x"
+    end
+    
+    def self.deserialize(payload)
+      "y"
+    end
+  end
+  
+  class NamedSerializer < SweatShop::Serializer
+    serializer_name :silly
+    def self.serialize(payload)
+      "x"
+    end
+    
+    def self.deserialize(payload)
+      "y"
+    end
+    
+  end
+    
+  test "should use default serializer if none is specified" do
+    assert_equal SweatShop::Serializers::MarshalSerializer, DefaultSerialWorker.serializer
+  end
+  
+  test "should mark marshal as the default serializer out of the box" do
+    assert_equal SweatShop::Serializer.default, SweatShop::Serializers::MarshalSerializer
+  end
+  test "should be able to specify the default serializer" do
+    SweatShop::Serializer.default = :json
+    assert_equal SweatShop::Serializer.default, SweatShop::Serializers::JsonSerializer
+    SweatShop::Serializer.default = :marshal # set it back so tests don't break
+  end
+  
+  test "should be able to specify json as the serializer on a per class basis" do
+    assert_equal JsonWorker.serializer, SweatShop::Serializers::JsonSerializer
+  end
+  
+  test "should be able to specify marshal as the serializer on a per class basis" do
+    assert_equal MarshalWorker.serializer, SweatShop::Serializers::MarshalSerializer
+  end
+
+  test "should be able to specify yaml as the serializer on a per class basis" do
+    assert_equal YamlWorker.serializer, SweatShop::Serializers::YamlSerializer
+  end
+
+  test "should provide a default short name based on class name" do
+    assert_equal :no_name, NoNameSerializer.get_name
+  end
+  
+  test "should accept a custom short name" do
+    assert_equal :silly, NamedSerializer.get_name
+  end
+  
+  test "should register all available serializers" do
+    assert_equal ["json", "marshal", "no_name", "silly", "yaml"], SweatShop::Serializer.serializers.keys.map{|k| k.to_s}.sort
+  end
+end