Bira-sweat_shop 1.0.0

Sign up to get free protection for your applications and to get access to all the features.
data/History.txt ADDED
@@ -0,0 +1,6 @@
1
+ === 1.0.0 / 2009-01-13
2
+
3
+ * 1 major enhancement
4
+
5
+ * Birthday!
6
+
data/LICENSE ADDED
@@ -0,0 +1,20 @@
1
+ Copyright (c) 2009 Amos Elliston
2
+
3
+ Permission is hereby granted, free of charge, to any person obtaining
4
+ a copy of this software and associated documentation files (the
5
+ "Software"), to deal in the Software without restriction, including
6
+ without limitation the rights to use, copy, modify, merge, publish,
7
+ distribute, sublicense, and/or sell copies of the Software, and to
8
+ permit persons to whom the Software is furnished to do so, subject to
9
+ the following conditions:
10
+
11
+ The above copyright notice and this permission notice shall be
12
+ included in all copies or substantial portions of the Software.
13
+
14
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
15
+ EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
16
+ MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
17
+ NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
18
+ LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
19
+ OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
20
+ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
data/README.markdown ADDED
@@ -0,0 +1,70 @@
1
+ # SweatShop
2
+
3
+ SweatShop provides an api to background resource intensive tasks. Much of the api design was copied from Workling, with a few tweaks.
4
+ Currently, it runs rabbitmq and kestrel, but it can support any number of queues.
5
+
6
+ ## Installing
7
+
8
+ gem install sweat_shop
9
+ freeze in your rails directory
10
+ cd vendor/gems/sweat_shop
11
+ rake setup
12
+
13
+ ## Writing workers
14
+
15
+ Put `email_worker.rb` into app/workers and sublcass `SweatShop::Worker`:
16
+
17
+ class EmailWorker
18
+ def send_mail(to)
19
+ user = User.find_by_id(to)
20
+ Mailer.deliver_welcome(to)
21
+ end
22
+ end
23
+
24
+ Then, anywhere in your app you can execute:
25
+
26
+ EmailWorker.async_send_mail(1)
27
+
28
+ 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
29
+ call:
30
+
31
+ EmailWorker.send_mail(1)
32
+
33
+ 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
34
+ which will allow you to split workers out into logical groups. This is important if you have various machines serving different
35
+ queues.
36
+
37
+ ## Running the queue
38
+
39
+ SweatShop has been tested with Rabbit and Kestrel, but it will also work with Starling. Please use the following resources to install the server:
40
+
41
+ Kestrel:
42
+ http://github.com/robey/kestrel/tree/master
43
+
44
+ Rabbit:
45
+ http://github.com/ezmobius/nanite/tree/master
46
+
47
+ config/sweatshop.yml specifies the machine address of the queue (default localhost:5672). You can also specify the queue type with the queue param.
48
+
49
+ ## Running the workers
50
+
51
+ Assuming you ran `rake setup` in Rails, you can type:
52
+
53
+ script/sweatshop
54
+
55
+ 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.
56
+
57
+ script/sweatshop -d
58
+ script/sweatshop -d stop
59
+
60
+ 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.
61
+
62
+ # REQUIREMENTS
63
+
64
+ i_can_daemonize
65
+ memcache (for kestrel)
66
+ carrot (for rabbit)
67
+
68
+ # LICENSE
69
+
70
+ Copyright (c) 2009 Amos Elliston, Geni.com; Published under The MIT License, see License
data/Rakefile ADDED
@@ -0,0 +1,42 @@
1
+ require 'rake'
2
+ require 'rake/testtask'
3
+ require 'rake/rdoctask'
4
+ require 'rcov/rcovtask'
5
+
6
+ begin
7
+ require 'jeweler'
8
+ Jeweler::Tasks.new do |s|
9
+ s.name = "sweat_shop"
10
+ s.summary = %Q{SweatShop is a simple asynchronous worker queue build on top of rabbitmq/ampq}
11
+ s.email = "amos@geni.com"
12
+ s.homepage = "http://github.com/famoseagle/sweat-shop"
13
+ s.description = "TODO"
14
+ s.authors = ["Amos Elliston"]
15
+ s.files = FileList["[A-Z]*", "{lib,test,config}/**/*"]
16
+ s.add_dependency('famoseagle-carrot', '= 0.6.0')
17
+ end
18
+ rescue LoadError
19
+ puts "Jeweler not available. Install it with: sudo gem install technicalpickles-jeweler -s http://gems.github.com"
20
+ end
21
+
22
+ Rake::TestTask.new
23
+
24
+ Rake::RDocTask.new do |rdoc|
25
+ rdoc.rdoc_dir = 'rdoc'
26
+ rdoc.title = 'new_project'
27
+ rdoc.options << '--line-numbers' << '--inline-source'
28
+ rdoc.rdoc_files.include('README*')
29
+ rdoc.rdoc_files.include('lib/**/*.rb')
30
+ end
31
+
32
+ Rcov::RcovTask.new do |t|
33
+ t.libs << 'test'
34
+ t.test_files = FileList['test/**/*_test.rb']
35
+ t.verbose = true
36
+ end
37
+
38
+ task :default => :test
39
+
40
+ task :setup do
41
+ require File.dirname(__FILE__) + '/install'
42
+ end
data/VERSION.yml ADDED
@@ -0,0 +1,4 @@
1
+ ---
2
+ :major: 1
3
+ :minor: 0
4
+ :patch: 0
data/bin/sweatd ADDED
@@ -0,0 +1,3 @@
1
+ #!/usr/bin/env ruby
2
+
3
+ require File.dirname(__FILE__) + '/../lib/sweat_shop/sweatd'
@@ -0,0 +1,9 @@
1
+ # default options
2
+ default:
3
+ queue: rabbit
4
+ host: localhost
5
+ port: 5672
6
+ user: 'guest'
7
+ pass: 'guest'
8
+ vhost: '/'
9
+ enable: true
@@ -0,0 +1,18 @@
1
+ development:
2
+ default:
3
+ queue: rabbit
4
+ host: localhost
5
+ port: 5672
6
+ enable: true
7
+ test:
8
+ default:
9
+ queue: rabbit
10
+ host: localhost
11
+ port: 5672
12
+ enable: true
13
+ production:
14
+ default:
15
+ queue: rabbit
16
+ host: localhost
17
+ port: 5672
18
+ enable: true
@@ -0,0 +1,17 @@
1
+ module MessageQueue
2
+ class Base
3
+ attr_reader :opts
4
+ def queue_size(queue); end
5
+ def enqueue(queue, data); end
6
+ def dequeue(queue); end
7
+ def confirm(queue); end
8
+ def subscribe(queue); end
9
+ def delete(queue); end
10
+ def client; end
11
+ def stop; end
12
+
13
+ def subscribe?
14
+ false
15
+ end
16
+ end
17
+ end
@@ -0,0 +1,32 @@
1
+ module MessageQueue
2
+ class Kestrel < Base
3
+ def initialize(opts)
4
+ @servers = opts['servers']
5
+ end
6
+
7
+ def queue_size(queue)
8
+ size = 0
9
+ stats = client.stats
10
+ servers.each do |server|
11
+ size += stats[server]["queue_#{queue}_items"].to_i
12
+ end
13
+ size
14
+ end
15
+
16
+ def enqueue(queue, data)
17
+ client.set(queue, data)
18
+ end
19
+
20
+ def dequeue(queue)
21
+ client.get("#{queue}/open")
22
+ end
23
+
24
+ def confirm(queue)
25
+ client.get("#{queue}/close")
26
+ end
27
+
28
+ def client
29
+ @client ||= MemCache.new(servers)
30
+ end
31
+ end
32
+ end
@@ -0,0 +1,71 @@
1
+ require 'carrot'
2
+ module MessageQueue
3
+ class Rabbit < Base
4
+
5
+ def initialize(opts={})
6
+ @opts = opts
7
+ end
8
+
9
+ def delete(queue)
10
+ send_command do
11
+ client.queue(queue).delete
12
+ end
13
+ end
14
+
15
+ def queue_size(queue)
16
+ send_command do
17
+ client.queue(queue).message_count
18
+ end
19
+ end
20
+
21
+ def enqueue(queue, data)
22
+ send_command do
23
+ client.queue(queue, :durable => true).publish(Marshal.dump(data), :persistent => true)
24
+ end
25
+ end
26
+
27
+ def dequeue(queue)
28
+ send_command do
29
+ task = client.queue(queue).pop(:ack => true)
30
+ return unless task
31
+ Marshal.load(task)
32
+ end
33
+ end
34
+
35
+ def confirm(queue)
36
+ send_command do
37
+ client.queue(queue).ack
38
+ end
39
+ end
40
+
41
+ def send_command(&block)
42
+ retried = false
43
+ begin
44
+ block.call
45
+ rescue Carrot::AMQP::Server::ServerDown => e
46
+ if not retried
47
+ puts "Error #{e.message}. Retrying..."
48
+ @client = nil
49
+ retried = true
50
+ retry
51
+ else
52
+ raise e
53
+ end
54
+ end
55
+ end
56
+
57
+ def client
58
+ @client ||= Carrot.new(
59
+ :host => @opts['host'],
60
+ :port => @opts['port'].to_i,
61
+ :user => @opts['user'],
62
+ :pass => @opts['pass'],
63
+ :vhost => @opts['vhost']
64
+ )
65
+ end
66
+
67
+ def stop
68
+ client.stop
69
+ end
70
+ end
71
+ end
@@ -0,0 +1,405 @@
1
+ require 'optparse'
2
+ require 'timeout'
3
+
4
+ module Daemoned
5
+ class DieTime < StandardError; end
6
+ class TimeoutError < StandardError; end
7
+
8
+ def self.included(base)
9
+ base.extend ClassMethods
10
+ base.initialize_options
11
+ end
12
+
13
+ class Config
14
+ METHODS = [:script_path]
15
+ CONFIG = {}
16
+ def method_missing(name, *args)
17
+ name = name.to_s.upcase.to_sym
18
+ if name.to_s =~ /^(.*)=$/
19
+ name = $1.to_sym
20
+ CONFIG[name] = args.first
21
+ else
22
+ CONFIG[name]
23
+ end
24
+ end
25
+ end
26
+
27
+ module ClassMethods
28
+
29
+ def initialize_options
30
+ @@config = Config.new
31
+ @@config.script_path = File.expand_path(File.dirname($0))
32
+ $0 = script_name
33
+ end
34
+
35
+ def parse_options
36
+ opts = OptionParser.new do |opt|
37
+ opt.banner = "Usage: #{script_name} [options] [start|stop]"
38
+
39
+ opt.on_tail('-h', '--help', 'Show this message') do
40
+ puts opt
41
+ exit(1)
42
+ end
43
+
44
+ opt.on('--loop-every=SECONDS', 'How long to sleep between each loop') do |value|
45
+ options[:loop_every] = value
46
+ end
47
+
48
+ opt.on('-t', '--ontop', 'Stay on top (does not daemonize)') do
49
+ options[:ontop] = true
50
+ end
51
+
52
+ opt.on('--instances=NUM', 'Allow multiple instances to run simultaneously? 0 for infinite. default: 1') do |value|
53
+ self.instances = value.to_i
54
+ end
55
+
56
+ opt.on('--log-file=LOGFILE', 'Logfile to log to') do |value|
57
+ options[:log_file] = File.expand_path(value)
58
+ end
59
+
60
+ opt.on('--pid-file=PIDFILE', 'Location of pidfile') do |value|
61
+ options[:pid_file] = File.expand_path(value)
62
+ end
63
+
64
+ opt.on('--no-log-prefix', 'Do not prefix PID and date/time in log file.') do
65
+ options[:log_prefix] = false
66
+ end
67
+ end
68
+
69
+ extra_args.each do |arg|
70
+ opts.on(*arg.first) do |value|
71
+ arg.last.call(value) if arg.last
72
+ end
73
+ end
74
+
75
+ opts.parse!
76
+
77
+ if ARGV.include?('stop')
78
+ stop
79
+ elsif ARGV.include?('reload')
80
+ kill('HUP')
81
+ exit
82
+ elsif not ARGV.include?('start') and not ontop?
83
+ puts opts.help
84
+ end
85
+ end
86
+
87
+ def arg(*args, &block)
88
+ self.extra_args << [args, block]
89
+ end
90
+
91
+ def extra_args
92
+ @extra_args ||= []
93
+ end
94
+
95
+ def callbacks
96
+ @callbacks ||= {}
97
+ end
98
+
99
+ def options
100
+ @options ||= {}
101
+ end
102
+
103
+ def options=(options)
104
+ @options = options
105
+ end
106
+
107
+ def config
108
+ yield @@config
109
+ end
110
+
111
+ def before(&block)
112
+ callbacks[:before] = block
113
+ end
114
+
115
+ def after(&block)
116
+ callbacks[:after] = block
117
+ end
118
+
119
+ def sig(*signals, &block)
120
+ signals.each do |s|
121
+ callbacks["sig_#{s}".to_sym] = block
122
+ end
123
+ end
124
+
125
+ def die_if(method=nil,&block)
126
+ options[:die_if] = method || block
127
+ end
128
+
129
+ def exit_if(method=nil,&block)
130
+ options[:exit_if] = method || block
131
+ end
132
+
133
+ def callback!(callback)
134
+ callbacks[callback].call if callbacks[callback]
135
+ end
136
+
137
+ # options may include:
138
+ #
139
+ # <tt>:loop_every</tt> Fixnum (DEFAULT 0)
140
+ # How many seconds to sleep between calls to your block
141
+ #
142
+ # <tt>:timeout</tt> Fixnum (DEFAULT 0)
143
+ # Timeout in if block does not execute withing passed number of seconds
144
+ #
145
+ # <tt>:kill_timeout</tt> Fixnum (DEFAULT 120)
146
+ # Wait number of seconds before using kill -9 on daemon
147
+ #
148
+ # <tt>:die_on_timeout</tt> BOOL (DEFAULT False)
149
+ # Should the daemon continue running if a block times out, or just run the block again
150
+ #
151
+ # <tt>:ontop</tt> BOOL (DEFAULT False)
152
+ # Do not daemonize. Run in current process
153
+ #
154
+ # <tt>:before</tt> BLOCK
155
+ # Run this block after daemonizing but before begining the daemonize loop.
156
+ # You can also define the before block by putting a before do/end block in your class.
157
+ #
158
+ # <tt>:after</tt> BLOCK
159
+ # Run this block before program exists.
160
+ # You can also define the after block by putting an after do/end block in your class.
161
+ #
162
+ # <tt>:die_if</tt> BLOCK
163
+ # Run this check after each iteration of the loop. If the block returns true, throw a DieTime exception and exit
164
+ # You can also define the after block by putting an die_if do/end block in your class.
165
+ #
166
+ # <tt>:exit_if</tt> BLOCK
167
+ # Run this check after each iteration of the loop. If the block returns true, exit gracefully
168
+ # You can also define the after block by putting an exit_if do/end block in your class.
169
+ #
170
+ # <tt>:log_prefix</tt> BOOL (DEFAULT true)
171
+ # Prefix log file entries with PID and timestamp
172
+ def daemonize(opts={}, &block)
173
+ self.options = opts
174
+ parse_options
175
+ return unless ok_to_start?
176
+
177
+ puts "Starting #{script_name}..."
178
+ puts "Logging to: #{log_file}" unless ontop?
179
+
180
+ unless ontop?
181
+ safefork do
182
+ open(pid_file, 'w'){|f| f << Process.pid }
183
+ at_exit { remove_pid! }
184
+
185
+ trap('TERM') { callback!(:sig_term) }
186
+ trap('INT') { callback!(:sig_int) ; Process.kill('TERM', $$) }
187
+ trap('HUP') { callback!(:sig_hup) }
188
+
189
+ sess_id = Process.setsid
190
+ reopen_filehandes
191
+
192
+ begin
193
+ at_exit { callback!(:after) }
194
+ callback!(:before)
195
+ run_block(&block)
196
+ rescue SystemExit
197
+ rescue Exception => e
198
+ $stdout.puts "Something bad happened #{e.inspect} #{e.backtrace.join("\n")}"
199
+ end
200
+ end
201
+ else
202
+ begin
203
+ callback!(:before)
204
+ run_block(&block)
205
+ rescue SystemExit, Interrupt
206
+ callback!(:after)
207
+ end
208
+ end
209
+ end
210
+
211
+ private
212
+
213
+ def run_block(&block)
214
+ loop do
215
+ if options[:timeout]
216
+ begin
217
+ Timeout::timeout(options[:timeout].to_i) do
218
+ block.call if block
219
+ end
220
+ rescue Timeout::Error => e
221
+ if options[:die_on_timeout]
222
+ raise TimeoutError.new("#{self} timed out after #{options[:timeout]} seconds while executing block in loop")
223
+ else
224
+ $stderr.puts "#{self} timed out after #{options[:timeout]} seconds while executing block in loop #{e.backtrace.join("\n")}"
225
+ end
226
+ end
227
+ else
228
+ block.call if block
229
+ end
230
+
231
+ if options[:loop_every]
232
+ sleep options[:loop_every].to_i
233
+ elsif not block
234
+ sleep 0.1
235
+ end
236
+
237
+ break if should_exit?
238
+ raise DieTime.new('Die if conditions were met!') if should_die?
239
+ end
240
+ exit(0)
241
+ end
242
+
243
+ def should_die?
244
+ die_if = options[:die_if]
245
+ if die_if
246
+ if die_if.is_a?(Symbol) or die_if.is_a?(String)
247
+ self.send(die_if)
248
+ elsif die_if.is_a?(Proc)
249
+ die_if.call
250
+ end
251
+ else
252
+ false
253
+ end
254
+ end
255
+
256
+ def should_exit?
257
+ exit_if = options[:exit_if]
258
+ if exit_if
259
+ if exit_if.is_a?(Symbol) or exit_if.is_a?(String)
260
+ self.send(exit_if.to_sym)
261
+ elsif exit_if.is_a?(Proc)
262
+ exit_if.call
263
+ end
264
+ else
265
+ false
266
+ end
267
+ end
268
+
269
+ def ok_to_start?
270
+ return true if pid.nil?
271
+
272
+ if process_alive?
273
+ $stderr.puts "#{script_name} is already running"
274
+ return false
275
+ else
276
+ $stderr.puts "Removing stale pid: #{pid}..."
277
+ end
278
+
279
+ true
280
+ end
281
+
282
+ def stop
283
+ puts "Stopping #{script_name}..."
284
+ kill
285
+ exit
286
+ end
287
+
288
+ def kill(signal = 'TERM')
289
+ if pid.nil?
290
+ $stderr.puts "#{script_name} doesn't appear to be running"
291
+ exit(1)
292
+ end
293
+ $stdout.puts("Sending pid #{pid} signal #{signal}...")
294
+ begin
295
+ Process.kill(signal, pid)
296
+ return if signal == 'HUP'
297
+ if pid_running?(options[:kill_timeout] || 120)
298
+ $stdout.puts("Using kill -9 #{pid}")
299
+ Process.kill(9, pid)
300
+ else
301
+ $stdout.puts("Process #{pid} stopped")
302
+ end
303
+ rescue Errno::ESRCH
304
+ $stdout.puts("Couldn't #{signal} #{pid} as it wasn't running")
305
+ end
306
+ end
307
+
308
+ def pid_running?(time_to_wait = 0)
309
+ times_to_check = 1
310
+ if time_to_wait > 0.5
311
+ times_to_check = (time_to_wait / 0.5).to_i
312
+ end
313
+ times_to_check.times do
314
+ return false unless process_alive?
315
+ sleep 0.5
316
+ end
317
+ true
318
+ end
319
+
320
+ def safefork(&block)
321
+ fork_tries ||= 0
322
+ fork(&block)
323
+ rescue Errno::EWOULDBLOCK
324
+ raise if fork_tries >= 20
325
+ fork_tries += 1
326
+ sleep 5
327
+ retry
328
+ end
329
+
330
+ def process_alive?
331
+ Process.kill(0, pid)
332
+ true
333
+ rescue Errno::ESRCH => e
334
+ false
335
+ end
336
+
337
+ LOG_FORMAT = '%-6d %-19s %s'
338
+ TIME_FORMAT = '%Y/%m/%d %H:%M:%S'
339
+ def reopen_filehandes
340
+ STDIN.reopen('/dev/null')
341
+ STDOUT.reopen(log_file, 'a')
342
+ STDOUT.sync = true
343
+ STDERR.reopen(STDOUT)
344
+ if log_prefix?
345
+ def STDOUT.write(string)
346
+ if @no_prefix
347
+ @no_prefix = false if string[-1, 1] == "\n"
348
+ else
349
+ string = LOG_FORMAT % [$$, Time.now.strftime(TIME_FORMAT), string]
350
+ @no_prefix = true
351
+ end
352
+ super(string)
353
+ end
354
+ end
355
+ end
356
+
357
+ def remove_pid!
358
+ if File.file?(pid_file) and File.read(pid_file).to_i == $$
359
+ File.unlink(pid_file)
360
+ end
361
+ end
362
+
363
+ def ontop?
364
+ options[:ontop]
365
+ end
366
+
367
+ def log_prefix?
368
+ options[:log_prefix] || true
369
+ end
370
+
371
+ LOG_PATHS = ['log/', 'logs/', '../log/', '../logs/', '../../log', '../../logs', '.']
372
+ LOG_PATHS.unshift("#{RAILS_ROOT}/log") if defined?(RAILS_ROOT)
373
+ def log_dir
374
+ options[:log_dir] ||= begin
375
+ LOG_PATHS.detect do |path|
376
+ File.exists?(File.expand_path(path))
377
+ end
378
+ end
379
+ end
380
+
381
+ def log_file
382
+ options[:log_file] ||= File.expand_path("#{log_dir}/#{script_name}.log")
383
+ end
384
+
385
+ def pid_dir
386
+ options[:pid_dir] ||= log_dir
387
+ end
388
+
389
+ def pid_file
390
+ options[:pid_file] ||= File.expand_path("#{pid_dir}/#{script_name}.pid")
391
+ end
392
+
393
+ def pid
394
+ @pid ||= File.file?(pid_file) ? File.read(pid_file).to_i : nil
395
+ end
396
+
397
+ def script_name
398
+ @script_name ||= File.basename($0).gsub('.rb', '')
399
+ end
400
+
401
+ def script_name=(script_name)
402
+ @script_name = script_name
403
+ end
404
+ end
405
+ end
@@ -0,0 +1,5 @@
1
+ def metaclass; class << self; self; end; end
2
+ def meta_eval(&blk); metaclass.instance_eval(&blk); end
3
+ def meta_def(name, &blk)
4
+ meta_eval { define_method name, &blk }
5
+ end
@@ -0,0 +1,76 @@
1
+ require File.dirname(__FILE__) + '/daemoned'
2
+
3
+ module SweatShop
4
+ class Sweatd
5
+ include Daemoned
6
+ queues = []
7
+ groups = []
8
+ rails_root = nil
9
+ start_cmd = "ruby #{__FILE__} start #{ARGV.reject{|a| a == 'start'}.join(' ')}"
10
+
11
+ arg '--workers=Worker,Worker', 'Workers to service (Default is all)' do |value|
12
+ queues = value.split(',')
13
+ end
14
+
15
+ arg '--groups=GROUP,GROUP', 'Groups of queues to service' do |value|
16
+ groups = value.split(',').collect{|g| g.to_sym}
17
+ end
18
+
19
+ arg '--worker-file=WORKERFILE', 'Worker file to load' do |value|
20
+ require value
21
+ end
22
+
23
+ arg '--worker-dir=WORKERDIR', 'Directory containing workers' do |value|
24
+ Dir.glob(value + '*.rb').each{|worker| require worker}
25
+ end
26
+
27
+ arg '--rails=DIR', 'Pass in RAILS_ROOT to run this daemon in a rails environment' do |value|
28
+ rails_root = value
29
+ end
30
+
31
+ sig(:term, :int) do
32
+ puts "Shutting down sweatd..."
33
+ SweatShop.stop
34
+ end
35
+
36
+ sig(:hup) do
37
+ puts "Received HUP"
38
+ SweatShop.stop
39
+ remove_pid!
40
+ puts "Restarting sweatd with #{start_cmd}..."
41
+ `#{start_cmd}`
42
+ end
43
+
44
+ before do
45
+ if rails_root
46
+ puts "Loading Rails..."
47
+ require rails_root + '/config/environment'
48
+ end
49
+ require File.dirname(__FILE__) + '/../sweat_shop'
50
+ end
51
+
52
+ daemonize(:kill_timeout => 20) do
53
+ workers = []
54
+
55
+ if groups.any?
56
+ workers += SweatShop.workers_in_group(groups)
57
+ end
58
+
59
+ if queues.any?
60
+ workers += queues.collect{|q| Object.module_eval(q)}
61
+ end
62
+
63
+ if workers.any?
64
+ worker_str = workers.join(',')
65
+ puts "Starting #{worker_str}..."
66
+ $0 = "Sweatd: #{worker_str}"
67
+ SweatShop.do_tasks(workers)
68
+ else
69
+ puts "Starting all workers..."
70
+ $0 = 'Sweatd: all'
71
+ SweatShop.do_all_tasks
72
+ end
73
+ end
74
+
75
+ end
76
+ end
@@ -0,0 +1,142 @@
1
+ require File.dirname(__FILE__) + '/metaid'
2
+
3
+ module SweatShop
4
+ class Worker
5
+ def self.inherited(subclass)
6
+ self.workers << subclass
7
+ end
8
+
9
+ def self.method_missing(method, *args, &block)
10
+ if method.to_s =~ /^async_(.*)/
11
+ method = $1
12
+ expected_args = instance.method(method).arity
13
+ if expected_args != args.size
14
+ raise ArgumentError.new("#{method} expects #{expected_args} arguments")
15
+ end
16
+ return instance.send(method, *args) unless config['enable']
17
+
18
+ uid = ::Digest::MD5.hexdigest("#{name}:#{method}:#{args}:#{Time.now.to_f}")
19
+ task = {:args => args, :method => method, :uid => uid, :queued_at => Time.now.to_i}
20
+
21
+ log("Putting #{uid} on #{queue_name}")
22
+ enqueue(task)
23
+
24
+ uid
25
+ elsif instance.respond_to?(method)
26
+ instance.send(method, *args)
27
+ else
28
+ super
29
+ end
30
+ end
31
+
32
+ def self.instance
33
+ @instance ||= new
34
+ end
35
+
36
+ def self.config
37
+ SweatShop.config
38
+ end
39
+
40
+ def self.queue_name
41
+ @queue_name ||= self.to_s
42
+ end
43
+
44
+ def self.delete_queue
45
+ queue.delete(queue_name)
46
+ end
47
+
48
+ def self.queue_size
49
+ queue.queue_size(queue_name)
50
+ end
51
+
52
+ def self.enqueue(task)
53
+ queue.enqueue(queue_name, task)
54
+ end
55
+
56
+ def self.dequeue
57
+ queue.dequeue(queue_name)
58
+ end
59
+
60
+ def self.confirm
61
+ queue.confirm(queue_name)
62
+ end
63
+
64
+ def self.subscribe
65
+ queue.subscribe(queue_name) do |task|
66
+ do_task(task)
67
+ end
68
+ end
69
+
70
+ def self.do_tasks
71
+ while task = dequeue
72
+ do_task(task)
73
+ end
74
+ end
75
+
76
+ def self.do_task(task)
77
+ call_before_task(task)
78
+
79
+ queued_at = task[:queued_at] ? "(queued #{Time.at(task[:queued_at]).strftime('%Y/%m/%d %H:%M:%S')})" : ''
80
+ log("Dequeuing #{queue_name}::#{task[:method]} #{queued_at}")
81
+ task[:result] = instance.send(task[:method], *task[:args])
82
+
83
+ call_after_task(task)
84
+ confirm
85
+ end
86
+
87
+ def self.call_before_task(task)
88
+ superclass.call_before_task(task) if superclass.respond_to?(:call_before_task)
89
+ before_task.call(task) if before_task
90
+ end
91
+
92
+ def self.call_after_task(task)
93
+ superclass.call_after_task(task) if superclass.respond_to?(:call_after_task)
94
+ after_task.call(task) if after_task
95
+ end
96
+
97
+ def self.queue
98
+ SweatShop.queue(queue_group.to_s)
99
+ end
100
+
101
+ def self.workers
102
+ SweatShop.workers
103
+ end
104
+
105
+ def self.config
106
+ SweatShop.config
107
+ end
108
+
109
+ def self.log(msg)
110
+ SweatShop.log(msg)
111
+ end
112
+
113
+ def self.before_task(&block)
114
+ if block
115
+ @before_task = block
116
+ else
117
+ @before_task
118
+ end
119
+ end
120
+
121
+ def self.after_task(&block)
122
+ if block
123
+ @after_task = block
124
+ else
125
+ @after_task
126
+ end
127
+ end
128
+
129
+ def self.stop
130
+ instance.stop
131
+ end
132
+
133
+ # called before we exit -- subclass can implement this method
134
+ def stop; end;
135
+
136
+
137
+ def self.queue_group(group=nil)
138
+ group ? meta_def(:_queue_group){ group } : _queue_group
139
+ end
140
+ queue_group :default
141
+ end
142
+ end
data/lib/sweat_shop.rb ADDED
@@ -0,0 +1,147 @@
1
+ require 'rubygems'
2
+ require 'digest'
3
+ require 'yaml'
4
+
5
+ $:.unshift(File.dirname(__FILE__))
6
+ require 'message_queue/base'
7
+ require 'message_queue/rabbit'
8
+ require 'message_queue/kestrel'
9
+ require 'sweat_shop/worker'
10
+
11
+ module SweatShop
12
+ extend self
13
+
14
+ def workers
15
+ @workers ||= []
16
+ end
17
+
18
+ def workers=(workers)
19
+ @workers = workers
20
+ end
21
+
22
+ def workers_in_group(groups)
23
+ groups = [groups] unless groups.is_a?(Array)
24
+ if groups.include?(:all)
25
+ workers
26
+ else
27
+ workers.select do |worker|
28
+ groups.include?(worker.queue_group)
29
+ end
30
+ end
31
+ end
32
+
33
+ def do_tasks(workers)
34
+ if queue.subscribe?
35
+ EM.run do
36
+ workers.each do |worker|
37
+ worker.subscribe
38
+ end
39
+ end
40
+ else
41
+ loop do
42
+ wait = true
43
+ workers.each do |worker|
44
+ if task = worker.dequeue
45
+ worker.do_task(task)
46
+ wait = false
47
+ end
48
+ end
49
+ if stop?
50
+ workers.each do |worker|
51
+ worker.stop
52
+ end
53
+ queue.stop
54
+ exit
55
+ end
56
+ sleep 1 if wait
57
+ end
58
+ end
59
+ end
60
+
61
+ def do_all_tasks
62
+ do_tasks(
63
+ workers_in_group(:all)
64
+ )
65
+ end
66
+
67
+ def do_default_tasks
68
+ do_tasks(
69
+ workers_in_group(:default)
70
+ )
71
+ end
72
+
73
+ def config
74
+ @config ||= begin
75
+ defaults = YAML.load_file(File.dirname(__FILE__) + '/../config/defaults.yml')
76
+ if defined?(RAILS_ROOT)
77
+ file = RAILS_ROOT + '/config/sweatshop.yml'
78
+ if File.exist?(file)
79
+ YAML.load_file(file)[RAILS_ENV || 'development']
80
+ else
81
+ defaults['enable'] = false
82
+ defaults
83
+ end
84
+ else
85
+ defaults
86
+ end
87
+ end
88
+ end
89
+
90
+ def stop
91
+ @stop = true
92
+ queue.stop if queue.subscribe?
93
+ end
94
+
95
+ def stop?
96
+ @stop
97
+ end
98
+
99
+ def queue_sizes
100
+ workers.inject([]) do |all, worker|
101
+ all << [worker, worker.queue_size]
102
+ all
103
+ end
104
+ end
105
+
106
+ def pp_sizes
107
+ max_width = workers.collect{|w| w.to_s.size}.max
108
+ puts '-' * (max_width + 10)
109
+ puts queue_sizes.collect{|p| sprintf("%-#{max_width}s %2s", p.first, p.last)}.join("\n")
110
+ puts '-' * (max_width + 10)
111
+ end
112
+
113
+ def queue(type = 'default')
114
+ @queues ||= {}
115
+ @queues[type] ||= begin
116
+ qconfig = config[type] || config['default']
117
+ qtype = qconfig['queue'] || 'rabbit'
118
+ queue = constantize("MessageQueue::#{qtype.capitalize}")
119
+ queue.new(qconfig)
120
+ end
121
+ end
122
+
123
+ def queue=(queue, type = 'default')
124
+ @queues[type] = queue
125
+ end
126
+
127
+ def log(msg)
128
+ return if logger == :silent
129
+ logger ? logger.debug(msg) : puts(msg)
130
+ end
131
+
132
+ def logger
133
+ @logger
134
+ end
135
+
136
+ def logger=(logger)
137
+ @logger = logger
138
+ end
139
+
140
+ def constantize(str)
141
+ Object.module_eval("#{str}", __FILE__, __LINE__)
142
+ end
143
+ end
144
+
145
+ if defined?(RAILS_ROOT)
146
+ Dir.glob(RAILS_ROOT + '/app/workers/*.rb').each{|worker| require worker }
147
+ end
@@ -0,0 +1,13 @@
1
+ require File.dirname(__FILE__) + '/../lib/sweat_shop'
2
+ class HelloWorker < SweatShop::Worker
3
+ TEST_FILE = File.dirname(__FILE__) + '/test.txt' unless defined?(TEST_FILE)
4
+
5
+ def hello(name)
6
+ puts name
7
+ "Hi, #{name}"
8
+ end
9
+
10
+ after_task do |task|
11
+ File.open(TEST_FILE, 'w'){|f| f << task[:result]}
12
+ end
13
+ end
@@ -0,0 +1,36 @@
1
+ require File.dirname(__FILE__) + '/../lib/sweat_shop'
2
+ require File.dirname(__FILE__) + '/test_helper'
3
+ require File.dirname(__FILE__) + '/hello_worker'
4
+
5
+ class WorkerTest < Test::Unit::TestCase
6
+
7
+ def setup
8
+ File.delete(HelloWorker::TEST_FILE) if File.exist?(HelloWorker::TEST_FILE)
9
+ end
10
+
11
+ def teardown
12
+ File.delete(HelloWorker::TEST_FILE) if File.exist?(HelloWorker::TEST_FILE)
13
+ end
14
+
15
+ test "daemon" do
16
+ begin
17
+ SweatShop.queue = nil
18
+ SweatShop.logger = :silent
19
+
20
+ worker = File.expand_path(File.dirname(__FILE__) + '/hello_worker')
21
+ sweatd = "#{File.dirname(__FILE__)}/../lib/sweat_shop/sweatd.rb"
22
+ uid = HelloWorker.async_hello('Amos')
23
+
24
+ `ruby #{sweatd} --worker-file #{worker} start`
25
+ `ruby #{sweatd} stop`
26
+
27
+ File.delete('sweatd.log') if File.exist?('sweatd.log')
28
+ assert_equal 'Hi, Amos', File.read(HelloWorker::TEST_FILE)
29
+ rescue Exception => e
30
+ puts e.message
31
+ puts e.backtrace.join("\n")
32
+ fail "\n\n*** Functional test failed, is the rabbit server running on localhost? ***\n"
33
+ end
34
+ end
35
+
36
+ end
@@ -0,0 +1,14 @@
1
+ require 'test/unit'
2
+
3
+ class << Test::Unit::TestCase
4
+ def test(name, &block)
5
+ test_name = "test_#{name.gsub(/[\s\W]/,'_')}"
6
+ raise ArgumentError, "#{test_name} is already defined" if self.instance_methods.include? test_name
7
+ define_method test_name, &block
8
+ end
9
+
10
+ def xtest(name, &block)
11
+ # no-op, an empty test method is defined to prevent "no tests in testcase" errors when all tests are disabled
12
+ define_method(:test_disabled) { assert true }
13
+ end
14
+ end
@@ -0,0 +1,65 @@
1
+ require File.dirname(__FILE__) + '/test_helper'
2
+ require File.dirname(__FILE__) + '/../lib/sweat_shop'
3
+
4
+ class SweatShopTest < Test::Unit::TestCase
5
+ SweatShop.workers = []
6
+
7
+ class HelloWorker < SweatShop::Worker
8
+ def hello(name)
9
+ "Hi, #{name}"
10
+ end
11
+ end
12
+
13
+ class GroupedWorker < SweatShop::Worker
14
+ queue_group :foo
15
+ end
16
+
17
+ test "group workers" do
18
+ assert_equal [HelloWorker, GroupedWorker], SweatShop.workers_in_group(:all)
19
+ assert_equal [HelloWorker], SweatShop.workers_in_group(:default)
20
+ assert_equal [GroupedWorker], SweatShop.workers_in_group(:foo)
21
+ end
22
+
23
+ test "synch call" do
24
+ worker = HelloWorker.new
25
+ assert_equal "Hi, Amos", worker.hello('Amos')
26
+ end
27
+
28
+ test "uid" do
29
+ SweatShop.logger = :silent
30
+ uid = HelloWorker.async_hello('Amos')
31
+ assert_not_nil uid
32
+ end
33
+
34
+ test "before task" do
35
+ HelloWorker.before_task do
36
+ "hello"
37
+ end
38
+ assert_equal "hello", HelloWorker.before_task.call
39
+ end
40
+
41
+ test "after task" do
42
+ HelloWorker.after_task do
43
+ "goodbye"
44
+ end
45
+ assert_equal "goodbye", HelloWorker.after_task.call
46
+ end
47
+
48
+ test "chainable before tasks" do
49
+ MESSAGES = []
50
+ class BaseWorker < SweatShop::Worker
51
+ before_task do |task|
52
+ MESSAGES << 'base'
53
+ end
54
+ end
55
+ class SubWorker < BaseWorker
56
+ before_task do |task|
57
+ MESSAGES << 'sub'
58
+ end
59
+ end
60
+ SubWorker.call_before_task('foo')
61
+ assert_equal ['base', 'sub'], MESSAGES
62
+ SweatShop.workers.delete(BaseWorker)
63
+ SweatShop.workers.delete(SubWorker)
64
+ end
65
+ end
metadata ADDED
@@ -0,0 +1,85 @@
1
+ --- !ruby/object:Gem::Specification
2
+ name: Bira-sweat_shop
3
+ version: !ruby/object:Gem::Version
4
+ version: 1.0.0
5
+ platform: ruby
6
+ authors:
7
+ - Amos Elliston
8
+ autorequire:
9
+ bindir: bin
10
+ cert_chain: []
11
+
12
+ date: 2009-07-17 00:00:00 -07:00
13
+ default_executable: sweatd
14
+ dependencies:
15
+ - !ruby/object:Gem::Dependency
16
+ name: famoseagle-carrot
17
+ type: :runtime
18
+ version_requirement:
19
+ version_requirements: !ruby/object:Gem::Requirement
20
+ requirements:
21
+ - - "="
22
+ - !ruby/object:Gem::Version
23
+ version: 0.6.0
24
+ version:
25
+ description: TODO
26
+ email: amos@geni.com
27
+ executables:
28
+ - sweatd
29
+ extensions: []
30
+
31
+ extra_rdoc_files:
32
+ - LICENSE
33
+ - README.markdown
34
+ files:
35
+ - History.txt
36
+ - LICENSE
37
+ - README.markdown
38
+ - Rakefile
39
+ - VERSION.yml
40
+ - config/defaults.yml
41
+ - config/sweatshop.yml
42
+ - lib/message_queue/base.rb
43
+ - lib/message_queue/kestrel.rb
44
+ - lib/message_queue/rabbit.rb
45
+ - lib/sweat_shop.rb
46
+ - lib/sweat_shop/daemoned.rb
47
+ - lib/sweat_shop/metaid.rb
48
+ - lib/sweat_shop/sweatd.rb
49
+ - lib/sweat_shop/worker.rb
50
+ - test/hello_worker.rb
51
+ - test/test_functional_worker.rb
52
+ - test/test_helper.rb
53
+ - test/test_sweatshop.rb
54
+ has_rdoc: true
55
+ homepage: http://github.com/famoseagle/sweat-shop
56
+ licenses:
57
+ post_install_message:
58
+ rdoc_options:
59
+ - --charset=UTF-8
60
+ require_paths:
61
+ - lib
62
+ required_ruby_version: !ruby/object:Gem::Requirement
63
+ requirements:
64
+ - - ">="
65
+ - !ruby/object:Gem::Version
66
+ version: "0"
67
+ version:
68
+ required_rubygems_version: !ruby/object:Gem::Requirement
69
+ requirements:
70
+ - - ">="
71
+ - !ruby/object:Gem::Version
72
+ version: "0"
73
+ version:
74
+ requirements: []
75
+
76
+ rubyforge_project:
77
+ rubygems_version: 1.3.5
78
+ signing_key:
79
+ specification_version: 2
80
+ summary: SweatShop is a simple asynchronous worker queue build on top of rabbitmq/ampq
81
+ test_files:
82
+ - test/hello_worker.rb
83
+ - test/test_functional_worker.rb
84
+ - test/test_helper.rb
85
+ - test/test_sweatshop.rb