workhorse 0.0.1 → 0.0.2

Sign up to get free protection for your applications and to get access to all the features.
@@ -0,0 +1,51 @@
1
+ module Workhorse
2
+ # Abstraction layer of a simple thread pool implementation used by the worker.
3
+ class Pool
4
+ attr_reader :mutex
5
+
6
+ def initialize(size)
7
+ @size = size
8
+ @executor = Concurrent::ThreadPoolExecutor.new(
9
+ min_threads: 0,
10
+ max_threads: @size,
11
+ max_queue: 0,
12
+ fallback_policy: :abort,
13
+ auto_terminate: false
14
+ )
15
+ @mutex = Mutex.new
16
+ @active_threads = Concurrent::AtomicFixnum.new(0)
17
+ end
18
+
19
+ # Posts a new work unit to the pool.
20
+ def post
21
+ mutex.synchronize do
22
+ if @active_threads.value >= @size
23
+ fail 'All threads are busy.'
24
+ end
25
+
26
+ active_threads = @active_threads
27
+
28
+ active_threads.increment
29
+
30
+ @executor.post do
31
+ begin
32
+ yield
33
+ ensure
34
+ active_threads.decrement
35
+ end
36
+ end
37
+ end
38
+ end
39
+
40
+ # Returns the number of idle threads.
41
+ def idle
42
+ @size - @active_threads.value
43
+ end
44
+
45
+ # Shuts down the pool
46
+ def shutdown
47
+ @executor.shutdown
48
+ @executor.wait_for_termination
49
+ end
50
+ end
51
+ end
@@ -0,0 +1,144 @@
1
+ module Workhorse
2
+ class Worker
3
+ LOG_LEVELS = %i[fatal error warn info debug].freeze
4
+ SHUTDOWN_SIGNALS = %w[TERM INT].freeze
5
+
6
+ attr_reader :queues
7
+ attr_reader :state
8
+ attr_reader :pool_size
9
+ attr_reader :polling_interval
10
+ attr_reader :mutex
11
+ attr_reader :logger
12
+
13
+ # Instantiates and starts a new worker with the given arguments and then
14
+ # waits for its completion (i.e. an interrupt).
15
+ def self.start_and_wait(*args)
16
+ worker = new(*args)
17
+ worker.start
18
+ worker.wait
19
+ end
20
+
21
+ # Instantiates a new worker. The worker is not automatically started.
22
+ #
23
+ # @param queues [Array] The queues you want this worker to process. If an
24
+ # empty array is given, any queues will be processed. Queues need to be
25
+ # specified as a symbol. To also process jobs without a queue, supply
26
+ # `nil` within the array.
27
+ # @param pool_size [Integer] The number of jobs that will be processed
28
+ # simultaneously. If this parameter is not given, it will be set to the
29
+ # number of given queues + 1.
30
+ # @param polling_interval [Integer] Interval in seconds the database will
31
+ # be polled for new jobs. Set this as high as possible to avoid
32
+ # unnecessary database load Set this as high as possible to avoid
33
+ # unnecessary database load.
34
+ # @param auto_terminate [Boolean] Whether to automatically shut down the
35
+ # worker properly on INT and TERM signals.
36
+ # @param quiet [Boolean] If this is set to `false`, the worker will also log
37
+ # to STDOUT.
38
+ # @param logger [Logger] An optional logger the worker will append to. This
39
+ # can be any instance of ruby's `Logger` but is commonly set to
40
+ # `Rails.logger`.
41
+ def initialize(queues: [], pool_size: nil, polling_interval: 5, auto_terminate: true, quiet: true, logger: nil)
42
+ @queues = queues
43
+ @pool_size = pool_size || queues.size + 1
44
+ @polling_interval = polling_interval
45
+ @auto_terminate = auto_terminate
46
+ @state = :initialized
47
+ @quiet = quiet
48
+
49
+ @mutex = Mutex.new
50
+ @pool = Pool.new(@pool_size)
51
+ @poller = Workhorse::Poller.new(self)
52
+ @logger = logger
53
+
54
+ fail 'Polling interval must be an integer.' unless @polling_interval.is_a?(Integer)
55
+ end
56
+
57
+ def log(text, level = :info)
58
+ text = "[Job worker #{id}] #{text}"
59
+ puts text unless @quiet
60
+ return unless logger
61
+ fail "Log level #{level} is not available. Available are #{LOG_LEVELS.inspect}." unless LOG_LEVELS.include?(level)
62
+ logger.send(level, text.strip)
63
+ end
64
+
65
+ def id
66
+ @id ||= "#{Socket.gethostname}.#{Process.pid}.#{SecureRandom.hex(3)}"
67
+ end
68
+
69
+ # Starts the worker. This call is not blocking - call {wait} for this
70
+ # purpose.
71
+ def start
72
+ mutex.synchronize do
73
+ assert_state! :initialized
74
+ log 'Starting up'
75
+ @state = :running
76
+ @poller.start
77
+ log 'Started up'
78
+
79
+ trap_termination if @auto_terminate
80
+ end
81
+ end
82
+
83
+ def assert_state!(state)
84
+ fail "Expected worker to be in state #{state} but current state is #{self.state}." unless self.state == state
85
+ end
86
+
87
+ # Shuts down worker and DB poller. Jobs currently beeing processed are
88
+ # properly finished before this method returns. Subsequent calls to this
89
+ # method are ignored.
90
+ def shutdown
91
+ mutex.synchronize do
92
+ return if @state == :shutdown
93
+ assert_state! :running
94
+ log 'Shutting down'
95
+ @state = :shutdown
96
+
97
+ @poller.shutdown
98
+ @pool.shutdown
99
+ log 'Shut down'
100
+ end
101
+ end
102
+
103
+ # Waits until the worker is shut down. This only happens if shutdown gets
104
+ # called - either by another thread or by enabling `auto_terminate` and
105
+ # receiving a respective signal. Use this method to let worker run
106
+ # undefinitely.
107
+ def wait
108
+ assert_state! :running
109
+ @poller.wait
110
+ end
111
+
112
+ def idle
113
+ @pool.idle
114
+ end
115
+
116
+ def perform(db_job)
117
+ mutex.synchronize do
118
+ assert_state! :running
119
+ log "Posting job #{db_job.id} to thread pool"
120
+
121
+ @pool.post do
122
+ begin
123
+ Workhorse::Performer.new(db_job, self).perform
124
+ rescue => e
125
+ log %(#{e.message}\n#{e.backtrace.join("\n")}), :error
126
+ end
127
+ end
128
+ end
129
+ end
130
+
131
+ private
132
+
133
+ def trap_termination
134
+ SHUTDOWN_SIGNALS.each do |signal|
135
+ Signal.trap(signal) do
136
+ Thread.new do
137
+ log "\nCaught #{signal}, shutting worker down..."
138
+ shutdown
139
+ end
140
+ end
141
+ end
142
+ end
143
+ end
144
+ end
@@ -0,0 +1,20 @@
1
+ ActiveRecord::Schema.define do
2
+ self.verbose = false
3
+
4
+ create_table :jobs, force: true do |t|
5
+ t.string :state, null: false, default: 'waiting'
6
+ t.string :queue, null: true
7
+ t.text :handler, null: false
8
+
9
+ t.string :locked_by
10
+ t.datetime :locked_at
11
+
12
+ t.datetime :started_at
13
+
14
+ t.datetime :succeeded_at
15
+ t.datetime :failed_at
16
+ t.text :last_error
17
+
18
+ t.timestamps null: false
19
+ end
20
+ end
data/test/lib/jobs.rb ADDED
@@ -0,0 +1,46 @@
1
+ class BasicJob
2
+ class_attribute :results
3
+ self.results = Concurrent::Array.new
4
+
5
+ def initialize(some_param: nil, sleep_time: 1)
6
+ @some_param = some_param
7
+ @sleep_time = sleep_time
8
+ end
9
+
10
+ def perform
11
+ results << @some_param
12
+ sleep @sleep_time
13
+ end
14
+ end
15
+
16
+ class DbConnectionTestJob
17
+ class_attribute :db_connections
18
+ self.db_connections = Concurrent::Array.new
19
+
20
+ def perform
21
+ db_connections << ActiveRecord::Base.connection.object_id
22
+ end
23
+ end
24
+
25
+ class DummyRailsOpsOp
26
+ class_attribute :results
27
+ self.results = Concurrent::Array.new
28
+
29
+ def self.run!(params = {})
30
+ new(params).run!
31
+ end
32
+
33
+ def initialize(params = {})
34
+ @params = params
35
+ end
36
+
37
+ def run!
38
+ perform
39
+ end
40
+
41
+ private
42
+
43
+ def perform
44
+ results << @params
45
+ end
46
+ end
@@ -0,0 +1,29 @@
1
+ require 'minitest/autorun'
2
+ require 'active_record'
3
+ require 'mysql2'
4
+ require 'benchmark'
5
+ require 'jobs'
6
+
7
+ class WorkhorseTest < ActiveSupport::TestCase
8
+ def setup
9
+ Workhorse::DbJob.delete_all
10
+ end
11
+
12
+ protected
13
+
14
+ def work(time = 2, options = {})
15
+ options[:pool_size] ||= 5
16
+ options[:polling_interval] ||= 1
17
+
18
+ w = Workhorse::Worker.new(options)
19
+ w.start
20
+ sleep time
21
+ w.shutdown
22
+ end
23
+ end
24
+
25
+ ActiveRecord::Base.logger = Logger.new('debug.log')
26
+ ActiveRecord::Base.establish_connection adapter: 'mysql2', database: 'workhorse', username: 'travis', password: '', pool: 10, host: :localhost
27
+
28
+ require 'db_schema'
29
+ require 'workhorse'
@@ -0,0 +1,42 @@
1
+ require 'test_helper'
2
+
3
+ class Workhorse::EnqueuerTest < WorkhorseTest
4
+ def test_basic
5
+ assert_equal 0, Workhorse::DbJob.all.count
6
+ Workhorse::Enqueuer.enqueue BasicJob.new
7
+ assert_equal 1, Workhorse::DbJob.all.count
8
+
9
+ db_job = Workhorse::DbJob.first
10
+ assert_equal 'waiting', db_job.state
11
+ assert_equal Marshal.dump(BasicJob.new), db_job.handler
12
+ assert_nil db_job.locked_by
13
+ assert_nil db_job.queue
14
+ assert_nil db_job.locked_at
15
+ assert_nil db_job.started_at
16
+ assert_nil db_job.last_error
17
+ assert_not_nil db_job.created_at
18
+ assert_not_nil db_job.updated_at
19
+ end
20
+
21
+ def test_with_queue
22
+ assert_equal 0, Workhorse::DbJob.all.count
23
+ Workhorse::Enqueuer.enqueue BasicJob.new, queue: :q1
24
+ assert_equal 1, Workhorse::DbJob.all.count
25
+
26
+ db_job = Workhorse::DbJob.first
27
+ assert_equal 'q1', db_job.queue
28
+ end
29
+
30
+ def test_op
31
+ Workhorse::Enqueuer.enqueue_op DummyRailsOpsOp, { foo: :bar }, queue: :q1
32
+
33
+ w = Workhorse::Worker.new(queues: [:q1])
34
+ w.start
35
+ sleep 1.5
36
+ w.shutdown
37
+
38
+ assert_equal 'succeeded', Workhorse::DbJob.first.state
39
+
40
+ assert_equal [{ foo: :bar }], DummyRailsOpsOp.results
41
+ end
42
+ end
@@ -0,0 +1,18 @@
1
+ require 'test_helper'
2
+
3
+ class Workhorse::WorkerTest < WorkhorseTest
4
+ # This test makes sure that concurrent jobs always work in different database
5
+ # connections.
6
+ def test_db_connections
7
+ w = Workhorse::Worker.new polling_interval: 1, pool_size: 5
8
+ 5.times do
9
+ Workhorse::Enqueuer.enqueue DbConnectionTestJob.new
10
+ end
11
+ w.start
12
+ sleep 2
13
+ w.shutdown
14
+
15
+ assert_equal 5, DbConnectionTestJob.db_connections.count
16
+ assert_equal 5, DbConnectionTestJob.db_connections.uniq.count
17
+ end
18
+ end
@@ -0,0 +1,13 @@
1
+ require 'test_helper'
2
+
3
+ class Workhorse::PollerTest < WorkhorseTest
4
+ def test_interruptable_sleep
5
+ w = Workhorse::Worker.new(polling_interval: 60)
6
+ w.start
7
+ sleep 2
8
+
9
+ Timeout.timeout(1.5) do
10
+ w.shutdown
11
+ end
12
+ end
13
+ end
@@ -0,0 +1,72 @@
1
+ require 'test_helper'
2
+
3
+ class Workhorse::PoolTest < WorkhorseTest
4
+ def test_idle
5
+ with_pool 5 do |p|
6
+ assert_equal 5, p.idle
7
+
8
+ 4.times do |_i|
9
+ p.post do
10
+ sleep 1
11
+ end
12
+ end
13
+
14
+ sleep 0.5
15
+ assert_equal 1, p.idle
16
+
17
+ sleep 1
18
+ assert_equal 5, p.idle
19
+ end
20
+ end
21
+
22
+ def test_overflow
23
+ with_pool 5 do |p|
24
+ 5.times { p.post { sleep 1 } }
25
+
26
+ exception = assert_raises do
27
+ p.post { sleep 1 }
28
+ end
29
+
30
+ assert_equal 'All threads are busy.', exception.message
31
+ end
32
+ end
33
+
34
+ def test_work
35
+ with_pool 5 do |p|
36
+ counter = Concurrent::AtomicFixnum.new(0)
37
+
38
+ 5.times do
39
+ p.post do
40
+ sleep 1
41
+ counter.increment
42
+ end
43
+ end
44
+
45
+ sleep 1.2
46
+
47
+ assert_equal 5, counter.value
48
+
49
+ 2.times do
50
+ p.post do
51
+ sleep 1
52
+ counter.increment
53
+ end
54
+ end
55
+
56
+ sleep 1.2
57
+
58
+ assert_equal 7, counter.value
59
+ end
60
+ end
61
+
62
+ private
63
+
64
+ def with_pool(size)
65
+ p = Workhorse::Pool.new(size)
66
+ begin
67
+ yield(p)
68
+ ensure
69
+ p.shutdown
70
+ end
71
+ end
72
+ end