simpleworker 0.0.1 → 0.1.0

data/lib/simpleworker/scripts/reliable_queue.lua

@@ -0,0 +1,12 @@
+local task = redis.call('LPOP', KEYS[1])
+if not task then
+   return nil
+end
+
+local task_key = string.format("%s:active:%s:%s:%s", ARGV[1], ARGV[2], ARGV[3], task)
+redis.pcall('SADD', KEYS[2], task_key)
+redis.pcall('SET',
+           task_key,
+           task,
+           "EX", ARGV[4])
+return task

data/lib/simpleworker/ssh_worker.rb

@@ -1,33 +1,54 @@
 
 module SimpleWorker
-  class SshWorker
-    include AbstractWorker
 
-    attr_accessor :directory, :user, :host
+  #
+  # SshWorker.new(opts)
+  #
+  # where 'opts' is a Hash of options:
+  #
+  #  :user    => String name of user on remote host (default: `whoami`)
+  #  :host    => String name of remote host (default: `hostname`)
+  #  :port    => String port to connect on the remote host (default: 22)
+  #  :cmd     => String bash string to execute on the remote host in the users login shell, in the remote 'dirname', and with the JOBID environment variable set to the current jobid, (default: `bundle install; rake;`)
+  #  :dirname => String remote directory in the users home directory (default: dynamic jobid)
+  class SshWorker < AbstractListener
 
-    def initialize
-      @script = %W[bash #{File.expand_path(File.dirname(__FILE__))}/bash/ssh-remoteworker]
-      @env = {}
+    attr_reader :user, :host, :port, :dirname, :cmd
+
+    def initialize(opts = {})
+      @user    = opts[:user] || `whoami`.strip
+      @host    = opts[:host] || `hostname`.strip
+      @port    = opts[:port] || '22'
+      @cmd     = opts[:cmd]  || 'bundle install; rake;'
+      @dirname = opts[:dirname]
+    end
+
+    # Destructive rsync to remote and start a process in the background on the remote server.
+    def on_start(jobid)
+      @dirname ||= jobid
+      @jobid = jobid
+
+      sync_to_remote
+      async_cmd
     end
 
-    def self.create(config = {})
-      worker = new
-      worker.script = config['script'] if config.has_key? 'script'
-      worker.cmd = config['cmd'] if config.has_key? 'cmd'
-      worker.user = config['user'] if config.has_key? 'user'
-      worker.host = config['host'] if config.has_key? 'host'
-      worker.directory = config['directory'] if config.has_key? 'directory'
-      worker
+    # Rsync from remote.
+    def on_stop
+      sync_from_remote
     end
 
     private
 
-    def set_process_env
-      env['user'] = user || `whoami`
-      env['host'] = host || 'localhost'
-      env['directory'] = directory || Dir.pwd
+    def sync_to_remote
+      `rsync -a --delete -e "ssh -p #{port}" "${PWD}/" "#{user}@#{host}:~/#{dirname}"`
+    end
+
+    def sync_from_remote
+      `rsync -a -e "ssh -p #{port}" "#{user}@#{host}:~/#{dirname}/" "${PWD}"`
+    end
 
-      super
+    def async_cmd
+      `ssh -p #{port} "#{user}@#{host}" "/bin/bash -lc 'cd ~/#{dirname}; export JOBID=#{@jobid}; #{cmd}' </dev/null >/dev/null 2>&1 &"`
     end
   end
 end

data/lib/simpleworker/task_queue.rb

@@ -0,0 +1,78 @@
+
+module SimpleWorker
+
+  # TaskQueue.new(redis, hostname, jobid, opts)
+  #
+  # where hostname is the machines hostname or a unique identifier
+  # and 'opts' is a Hash of options:
+  #
+  #   :namespace    => String prefix to keys in redis used by SimpleWorker (default: simpleworker)
+  #   :task_timeout => Fixnum time after which a task expires, this should be > timout set in Runner (default: 10 seconds)
+  class TaskQueue
+    include RedisSupport
+
+    DEFAULT_OPTIONS = {
+      :namespace    => 'simpleworker'
+    }
+
+    def initialize(redis, hostname, jobid, opts = {})
+      opts               = DEFAULT_OPTIONS.dup.merge(opts)
+      @redis             = redis
+      @namespace         = opts[:namespace]
+      @task_timeout      = opts[:task_timeout]
+      @jobid             = jobid
+      @hostname          = hostname
+      @active_key_prefix = "#{active_tasks_key}:#{@hostname}"
+      @task_timeout      = JSON.parse(@redis.get(config_key))['task_timeout']
+      load_lua_scripts
+    end
+
+    def fire_start
+      push_to_log('on_node_start', @hostname)
+    end
+
+    def fire_stop
+      push_to_log('on_node_stop', @hostname)
+    end
+
+    def fire_task_start(task)
+      push_to_log('on_task_start', @hostname, task)
+    end
+
+    def fire_task_stop(task)
+      push_to_log('on_task_stop', @hostname, task)
+    end
+
+    def expire_current_task
+      @redis.del "#{@active_key_prefix}:#{@current_task}" if @current_task
+      @current_task = nil
+    end
+
+    def fire_log_message(msg)
+      push_to_log('on_log', @hostname, msg)
+    end
+
+    def each_task
+      until pop.nil?
+        local_task = @current_task
+        fire_task_start local_task
+        yield local_task
+        fire_task_stop local_task
+      end
+    end
+
+    def pop
+      @redis.srem(active_tasks_key, "#{@active_key_prefix}:#{@current_task}") if @current_task
+      @current_task = @redis.evalsha(@reliable_queue_sha,
+                                     :keys => [tasks_key, active_tasks_key],
+                                     :argv => [namespace, jobid, @hostname, @task_timeout])
+      @current_task
+    end
+
+    private
+
+    def push_to_log(*args)
+      @redis.rpush(log_key, args.to_json)
+    end
+  end
+end

data/lib/simpleworker/version.rb

@@ -1,4 +1,4 @@
 
 module SimpleWorker
-  VERSION = '0.0.1'
+  VERSION = '0.1.0'
 end

data/lib/simpleworker.rb

@@ -1,11 +1,22 @@
 
 require 'childprocess'
 require 'yaml'
+require 'redis'
+require 'observer'
+require 'json'
+require 'securerandom'
+require 'logger'
 
 module SimpleWorker
 end
 
+require 'simpleworker/redis_support'
 require 'simpleworker/runner'
-require 'simpleworker/abstract_worker'
+require 'simpleworker/abstract_listener'
+require 'simpleworker/event_monitor'
+require 'simpleworker/event_server'
+require 'simpleworker/task_queue'
 require 'simpleworker/local_worker'
 require 'simpleworker/ssh_worker'
+require 'simpleworker/logging_listener'
+require 'simpleworker/retry_listener'

data/simpleworker.gemspec

@@ -1,4 +1,3 @@
-# -*- encoding: utf-8 -*-
 $:.push File.expand_path("../lib", __FILE__)
 require "simpleworker/version"
 
@@ -16,6 +15,7 @@ Gem::Specification.new do |s|
   s.rubyforge_project = "simpleworker"
 
   s.add_dependency 'childprocess', '>= 0.5.3'
+  s.add_dependency 'redis', '>= 0.0.1'
 
   s.add_development_dependency "rspec", "~> 2.5"
   s.add_development_dependency "rake", "~> 0.9.2"

data/spec/simpleworker/event_monitor_spec.rb

@@ -0,0 +1,92 @@
+require File.expand_path('../../spec_helper.rb', __FILE__)
+
+module SimpleWorker
+  describe EventMonitor do
+    let(:start_time)    { Time.now }
+    let(:event_monitor) { EventMonitor.new(start_time = start_time) }
+    let(:jobid)         { 'my_jobid' }
+    let(:hostname)      { 'my_hostname' }
+    let(:task)          { 'my_task' }
+
+    it 'can initialize event monitor' do
+      event_monitor = EventMonitor.new(start_time = start_time)
+
+      expect(event_monitor.start_time).to eq(start_time)
+      expect(event_monitor.latest_time).to eq(start_time)
+      expect(event_monitor.expired_tasks).to be_empty
+      expect(event_monitor.done?(1)).to be false
+      expect(event_monitor.done?(0)).to be true
+    end
+
+    it 'can start node' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+
+      expect(event_monitor.done?(0)).to be false
+    end
+
+    it 'can stop node' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_node_stop(hostname)
+
+      expect(event_monitor.done?(0)).to be true
+    end
+
+    it 'can start task' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_task_start(hostname, task)
+
+      expect(event_monitor.done?(0)).to be false
+    end
+
+    it 'can stop task' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_task_start(hostname, task)
+      event_monitor.on_task_stop(hostname, task)
+
+      expect(event_monitor.done?(0)).to be false
+    end
+
+    it 'can expire task' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_task_start(hostname, task)
+      event_monitor.on_task_expire(hostname, task)
+
+      expect(event_monitor.done?(0)).to be false
+    end
+
+    it 'can expire task and stop node' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_task_start(hostname, task)
+      event_monitor.on_task_expire(hostname, task)
+      event_monitor.on_node_stop(hostname)
+
+      expect(event_monitor.done?(0)).to be true
+    end
+
+    it 'can complete on happy path' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_task_start(hostname, task)
+      event_monitor.on_task_stop(hostname, task)
+      event_monitor.on_node_stop(hostname)
+
+      expect(event_monitor.done?(0)).to be true
+    end
+
+    it 'can complete with expired task' do
+      event_monitor.on_start(jobid)
+      event_monitor.on_node_start(hostname)
+      event_monitor.on_task_start(hostname, task)
+      event_monitor.on_task_expire(hostname, task)
+      event_monitor.on_node_stop(hostname)
+
+      expect(event_monitor.done?(0)).to be true
+    end
+  end
+end

data/spec/simpleworker/event_server_spec.rb

@@ -0,0 +1,24 @@
+require File.expand_path('../../spec_helper.rb', __FILE__)
+
+module SimpleWorker
+  describe EventServer do
+    let(:redis)            { double(Redis, :script => nil) }
+    let(:namespace)        { 'my_namespace' }
+    let(:jobid)            { 'my_jobid' }
+    let(:event_server)     { EventServer.new(redis, namespace, jobid) }
+    let(:log_key)          { "#{namespace}:log:#{jobid}" }
+    let(:active_tasks_key) { "#{namespace}:active:#{jobid}" }
+    let(:tasks_key)        { "#{namespace}:tasks:#{jobid}" }
+
+    it 'can pull events' do
+      expect(redis).to receive(:multi).and_return([[["event"].to_json], [["#{namespace}:active:#{jobid}:my_hostname:my_task"],["#{namespace}:active:#{jobid}:my2_hostname:my2_task"]], 0])
+
+      expect(event_server).to receive(:fire).with("event")
+      expect(event_server).to receive(:fire).with('on_task_expire', 'my_hostname', 'my_task')
+      expect(event_server).to receive(:fire).with('on_task_active', 'my2_hostname', 'my2_task')
+
+      result = event_server.pull_events
+      expect(result).to eq(1)
+    end
+  end
+end

data/spec/simpleworker/local_worker_spec.rb

@@ -2,32 +2,25 @@ require File.expand_path('../../spec_helper.rb', __FILE__)
 
 module SimpleWorker
   describe LocalWorker do
+    let(:mock_process) { double(ChildProcess, :io => double('io').as_null_object) }
     let(:mock_env)     { double(Hash) }
-    let(:mock_process) { double(ChildProcess, :start => nil, :stop => nil, :wait => nil, :io => double('io').as_null_object, :environment => mock_env) }
+    let(:local_worker) { LocalWorker.new }
+    let(:jobid)        { 'my_jobid' }
 
-    it 'can start local worker with defaults' do
-      ChildProcess.should_receive(:build).with('bash', end_with('simple-localworker')).and_return(mock_process)
+    it 'can start local worker' do
+      ChildProcess.should_receive(:build).and_return(mock_process)
+      mock_process.should_receive(:environment).and_return(mock_env)
+      mock_env.should_receive(:[]=).with('JOBID', jobid)
       mock_process.should_receive(:start)
-      mock_env.should_receive(:[]=).with('cmd', 'rake')
 
-      worker = LocalWorker.new
-      worker.cmd = 'rake'
-      worker.start
+      local_worker.on_start(jobid)
     end
 
-    it 'can start local worker with customization' do
-      ChildProcess.should_receive(:build).with('bash', end_with('simple-localworker')).and_return(mock_process)
-      mock_process.should_receive(:start)
-      mock_env.should_receive(:[]=).with('cmd', 'cuke')
-      mock_process.should_receive(:wait)
+    it 'can stop local worker' do
+      ChildProcess.should_receive(:build).and_return(mock_process)
       mock_process.should_receive(:stop)
 
-      worker = LocalWorker.new
-      worker.cmd = 'cuke'
-      worker.start
-      worker.wait
-      worker.stop
+      local_worker.on_stop
     end
   end
 end
-

data/spec/simpleworker/logging_listener_spec.rb

@@ -0,0 +1,44 @@
+require File.expand_path('../../spec_helper.rb', __FILE__)
+
+module SimpleWorker
+  describe LoggingListener do
+    let(:stdout)    { StringIO.new }
+    let(:listener) { LoggingListener.new stdout }
+
+    it 'logs all events' do
+      Time.stub(:now => Time.now)
+
+      jobid = 'my_jobid'
+      host  = 'my_host'
+      task  = 'my_task'
+
+      listener.on_start jobid
+      listener.on_node_start host
+      listener.on_task_start host, task
+      listener.on_task_active host, task
+      listener.on_task_expire host, task
+      listener.on_task_stop host, task
+      listener.on_node_stop host
+      listener.on_timeout
+      listener.on_log host, 'message'
+      listener.on_interrupted
+      listener.on_stop
+
+      timestamp = Time.now.strftime("%Y-%m-%d %H:%M:%S##{Process.pid}")
+
+      stdout.string.should == <<-OUTPUT
+I, [#{timestamp}]  INFO -- : start: my_jobid
+I, [#{timestamp}]  INFO -- : start node: my_host
+I, [#{timestamp}]  INFO -- : start host: my_host task: my_task
+I, [#{timestamp}]  INFO -- : active host: my_host task: my_task
+I, [#{timestamp}]  INFO -- : expire host: my_host task: my_task
+I, [#{timestamp}]  INFO -- : stop host: my_host task: my_task
+I, [#{timestamp}]  INFO -- : stop node: my_host
+I, [#{timestamp}]  INFO -- : timeout
+I, [#{timestamp}]  INFO -- : host: my_host message
+I, [#{timestamp}]  INFO -- : interrupted
+I, [#{timestamp}]  INFO -- : stop
+OUTPUT
+    end
+  end
+end

data/spec/simpleworker/retry_listener_spec.rb

@@ -0,0 +1,25 @@
+require File.expand_path('../../spec_helper.rb', __FILE__)
+
+module SimpleWorker
+  describe RetryListener do
+    let(:redis)     { double(Redis) }
+    let(:namespace) { 'my_namespace' }
+    let(:jobid)     { 'my_jobid' }
+    let(:hostname)  { 'my_hostname' }
+    let(:jobid)     { 'my_jobid' }
+
+    it 'does not retry if max retries is 0' do
+      listener = RetryListener.new(redis, 0, namespace, jobid)
+      redis.should_not_receive(:rpush)
+
+      listener.on_task_expire(hostname, jobid)
+    end
+
+    it 'does retry once if max retries is 1' do
+      listener = RetryListener.new(redis, 1, namespace, jobid)
+      redis.should_receive(:rpush).once
+
+      2.times { listener.on_task_expire(hostname, jobid) }
+    end
+  end
+end

data/spec/simpleworker/runner_spec.rb

@@ -2,18 +2,80 @@ require File.expand_path('../../spec_helper.rb', __FILE__)
 
 module SimpleWorker
   describe Runner do
-    let(:mock_worker) { double(AbstractWorker, :start => nil, :wait => nil, :stop => nil) }
-
-    it 'can run workers' do
-      mock_worker.should_receive(:start)
-      mock_worker.should_receive(:cmd=).with('cuke')
-      mock_worker.should_receive(:wait)
-      mock_worker.should_receive(:stop)
-      IO.should_receive(:read).with('simpleworker.yml').and_return('hi')
-      YAML.should_receive(:load).with('hi').and_return({'workers' => [{'type' => 'ssh'}]})
-      SshWorker.should_receive(:create).and_return(mock_worker)
-
-      Runner.run 'cuke'
+    let(:redis)              { double(Redis, :script => nil) }
+    let(:tasks)              { ['ask1', 'task2'] }
+    let(:hostname)           { 'my_hostname' }
+    let(:jobid)              { 'my_jobid' }
+    let(:runner)             { Runner.new(redis, tasks, :timeout => 120, :interval => 0) }
+    let(:mock_event_monitor) { double(EventMonitor, :update => nil) }
+    let(:mock_event_server)  { double(EventServer, :add_observer => nil) }
+    let(:mock_retry_listener) { double(RetryListener, :update => nil, :tasks => [] ) }
+    let(:listener)           { double(AbstractListener, :update => nil) }
+
+    before(:each) do
+      ::SecureRandom.should_receive(:hex).with(6).and_return(jobid)
+
+      EventServer.should_receive(:new).and_return(mock_event_server)
+      EventMonitor.should_receive(:new).and_return(mock_event_monitor)
+      RetryListener.should_receive(:new).and_return(mock_retry_listener)
+      redis.should_receive(:rpush).with("simpleworker:tasks:#{jobid}", tasks)
+      redis.should_receive(:set).with("simpleworker:config:my_jobid", {'task_timeout' => 10}.to_json)
+
+      redis.should_receive(:multi)
+    end
+
+    it 'can run' do
+      mock_event_monitor.should_receive(:done?).and_return(false, true)
+      mock_event_server.should_receive(:pull_events).and_return(2, 0)
+      mock_event_monitor.should_receive(:latest_time).and_return(Time.now)
+      runner.add_observer listener
+      listener.should_receive(:update).with("on_start", jobid)
+      listener.should_receive(:update).with("on_stop")
+
+      runner.run
+    end
+
+    it 'can timeout' do
+      mock_event_monitor.should_receive(:done?).and_return(false)
+      mock_event_server.should_receive(:pull_events).and_return(2, 2)
+      mock_event_monitor.should_receive(:latest_time).and_return(Time.now)
+
+      runner = Runner.new(redis, tasks, :timeout => 0, :interval => 0)
+      runner.add_observer listener
+      listener.should_receive(:update).with("on_timeout")
+
+      runner.run
+    end
+
+    it 'fires on_interrupted and shuts down if interrupted' do
+      runner.add_observer listener
+
+      mock_event_server.should_receive(:pull_events).and_raise(Interrupt)
+      listener.should_receive(:update).with("on_interrupted")
+
+      runner.run
+    end
+
+    it 'fires on_interrupted and shuts down if an error occurs' do
+      runner.add_observer listener
+
+      mock_event_server.should_receive(:pull_events).and_raise(StandardError)
+      listener.should_receive(:update).with("on_interrupted")
+
+      lambda { runner.run }.should raise_error(StandardError)
+    end
+
+    it 'continue to run if retries exist' do
+      mock_event_monitor.should_receive(:done?).twice.with(2).and_return(false)
+      mock_event_monitor.should_receive(:done?).with(0).and_return(true)
+      mock_event_server.should_receive(:pull_events).and_return(2, 2, 0)
+
+      mock_event_monitor.should_receive(:latest_time).twice.and_return(Time.now)
+      runner.add_observer listener
+      listener.should_receive(:update).with("on_start", jobid)
+      listener.should_receive(:update).with("on_stop")
+
+      runner.run
     end
   end
 end

data/spec/simpleworker/ssh_worker_spec.rb

@@ -2,42 +2,31 @@ require File.expand_path('../../spec_helper.rb', __FILE__)
 
 module SimpleWorker
   describe SshWorker do
-    let(:mock_env)     { double(Hash) }
-    let(:mock_process) { double(ChildProcess, :start => nil, :stop => nil, :wait => nil, :io => double('io').as_null_object, :environment => mock_env) }
+    let(:mock_process) { double(ChildProcess, :io => double('io').as_null_object) }
+    let(:jobid)        { 'my_jobid' }
+    let(:ssh_worker) { SshWorker.new(
+      :user    => 'jesg',
+      :host    => 'localhost',
+      :port    => '22',
+      :cmd     => 'my_cmd',
+      :dirname => 'work_dir')
+    }
 
-    it 'can start ssh worker with defaults' do
-      ChildProcess.should_receive(:build).with('bash', end_with('ssh-remoteworker')).and_return(mock_process)
-      mock_process.should_receive(:start)
-      mock_env.should_receive(:[]=).once.with('cmd', 'rake')
-      mock_env.should_receive(:[]=).once.with('user', `whoami`)
-      mock_env.should_receive(:[]=).once.with('host', 'localhost')
-      mock_env.should_receive(:[]=).once.with('directory', Dir.pwd)
+    it 'can start ssh worker' do
+      expected_sync_cmd ='rsync -a --delete -e "ssh -p 22" "${PWD}/" "jesg@localhost:~/work_dir"'
+      expected_async_cmd = %q[ssh -p 22 "jesg@localhost" "/bin/bash -lc 'cd ~/work_dir; export JOBID=my_jobid; my_cmd' </dev/null >/dev/null 2>&1 &"]
+      Kernel.stub('`')
+      ssh_worker.should_receive(:sync_to_remote)
+      ssh_worker.should_receive(:async_cmd)
 
-      worker = SshWorker.new
-      worker.cmd = 'rake'
-      worker.start
+      ssh_worker.on_start(jobid)
     end
 
-    it 'can start ssh worker with customization' do
-      ChildProcess.should_receive(:build).with('bash', end_with('ssh-remoteworker')).and_return(mock_process)
-      mock_process.should_receive(:start)
-      mock_env.should_receive(:[]=).once.with('cmd', 'cuke')
-      mock_env.should_receive(:[]=).once.with('user', 'bob')
-      mock_env.should_receive(:[]=).once.with('host', 'foo')
-      mock_env.should_receive(:[]=).once.with('directory', 'bar')
+    it 'can stop ssh worker' do
+      Kernel.stub('`')
+      ssh_worker.should_receive(:sync_from_remote)
 
-      mock_process.should_receive(:wait)
-      mock_process.should_receive(:stop)
-
-      worker = SshWorker.new
-      worker.cmd = 'cuke'
-      worker.user = 'bob'
-      worker.host = 'foo'
-      worker.directory = 'bar'
-      worker.start
-      worker.wait
-      worker.stop
+      ssh_worker.on_stop
     end
   end
 end
-