xing-gearman-ruby 1.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,139 @@
1
+ #!/usr/bin/env ruby
2
+
3
+ require 'socket'
4
+
5
+ module Gearman
6
+
7
+ # = Client
8
+ #
9
+ # == Description
10
+ # A client for communicating with Gearman job servers.
11
+ class Client
12
+ ##
13
+ # Create a new client.
14
+ #
15
+ # @param job_servers "host:port"; either a single server or an array
16
+ def initialize(job_servers=nil)
17
+ @job_servers = [] # "host:port"
18
+ self.job_servers = job_servers if job_servers
19
+ @sockets = {} # "host:port" -> [sock1, sock2, ...]
20
+ @socket_to_hostport = {} # sock -> "host:port"
21
+ @test_hostport = nil # make get_job_server return a given host for testing
22
+ @task_create_timeout_sec = 10
23
+ @server_counter = -1
24
+ @bad_servers = []
25
+ end
26
+ attr_reader :job_servers, :bad_servers
27
+ attr_accessor :test_hostport, :task_create_timeout_sec
28
+
29
+ ##
30
+ # Set the job servers to be used by this client.
31
+ #
32
+ # @param servers "host:port"; either a single server or an array
33
+ def job_servers=(servers)
34
+ @job_servers = Util.normalize_job_servers(servers)
35
+ self
36
+ end
37
+
38
+ ##
39
+ # Get connection info about an arbitrary (currently random, but maybe
40
+ # we'll do something smarter later) job server.
41
+ #
42
+ # @return "host:port"
43
+ def get_job_server
44
+
45
+ raise Exception.new('No servers available') if @job_servers.empty?
46
+
47
+ @server_counter += 1
48
+ # Return a specific server if one's been set.
49
+ @test_hostport or @job_servers[@server_counter % @job_servers.size]
50
+ end
51
+
52
+ def signal_bad_server(hostport)
53
+ @job_servers = @job_servers.reject { |s| s == hostport }
54
+ @bad_servers << hostport
55
+ end
56
+ ##
57
+ # Get a socket for a job server.
58
+ #
59
+ # @param hostport job server "host:port"
60
+ # @return a Socket
61
+ def get_socket(hostport, num_retries=3)
62
+ # If we already have an open socket to this host, return it.
63
+ if @sockets[hostport]
64
+ sock = @sockets[hostport].shift
65
+ @sockets.delete(hostport) if @sockets[hostport].size == 0
66
+ return sock
67
+ end
68
+
69
+ num_retries.times do |i|
70
+ begin
71
+ sock = TCPSocket.new(*hostport.split(':'))
72
+ rescue Exception
73
+ else
74
+
75
+ @socket_to_hostport[sock] = hostport
76
+ return sock
77
+ end
78
+ end
79
+
80
+ signal_bad_server(hostport)
81
+ raise RuntimeError, "Unable to connect to job server #{hostport}"
82
+ end
83
+
84
+ ##
85
+ # Relinquish a socket created by Client#get_socket.
86
+ #
87
+ # If we don't know about the socket, we just close it.
88
+ #
89
+ # @param sock Socket
90
+ def return_socket(sock)
91
+ hostport = get_hostport_for_socket(sock)
92
+ if not hostport
93
+ inet, port, host, ip = s.addr
94
+ Util.err "Got socket for #{ip}:#{port}, which we don't " +
95
+ "know about -- closing"
96
+ sock.close
97
+ return
98
+ end
99
+ (@sockets[hostport] ||= []) << sock
100
+ end
101
+
102
+ def close_socket(sock)
103
+ sock.close
104
+ @socket_to_hostport.delete(sock)
105
+ nil
106
+ end
107
+
108
+ ##
109
+ # Given a socket from Client#get_socket, return its host and port.
110
+ #
111
+ # @param sock Socket
112
+ # @return "host:port", or nil if unregistered (which shouldn't happen)
113
+ def get_hostport_for_socket(sock)
114
+ @socket_to_hostport[sock]
115
+ end
116
+
117
+ ##
118
+ # Perform a single task.
119
+ #
120
+ # @param args either a Task or arguments for Task.new
121
+ # @return output of the task, or nil on failure
122
+ def do_task(*args)
123
+ task = Util::get_task_from_args(*args)
124
+
125
+ result = nil
126
+ failed = false
127
+ task.on_complete {|v| result = v }
128
+ task.on_fail { failed = true }
129
+
130
+ taskset = TaskSet.new(self)
131
+ taskset.add_task(task)
132
+ taskset.wait
133
+
134
+ failed ? nil : result
135
+ end
136
+
137
+ end
138
+
139
+ end
@@ -0,0 +1,94 @@
1
+ #!/usr/bin/env ruby
2
+
3
+ require 'socket'
4
+ require 'gearman'
5
+
6
+ module Gearman
7
+
8
+ # = Server
9
+ #
10
+ # == Description
11
+ # A client for managing Gearman job servers.
12
+ class Server
13
+ ##
14
+ # Create a new client.
15
+ #
16
+ # @param job_servers "host:port"; either a single server or an array
17
+ # @param prefix function name prefix (namespace)
18
+ def initialize(hostport)
19
+ @hostport = hostport # "host:port"
20
+ end
21
+ attr_reader :hostport
22
+
23
+ ##
24
+ # Get a socket for a job server.
25
+ #
26
+ # @param hostport job server "host:port"
27
+ # @return a Socket
28
+ def socket(num_retries=3)
29
+ return @socket if @socket
30
+ num_retries.times do
31
+ begin
32
+ sock = TCPSocket.new(*hostport.split(':'))
33
+ rescue Exception
34
+ else
35
+ return @socket = sock
36
+ end
37
+ end
38
+ raise RuntimeError, "Unable to connect to job server #{hostport}"
39
+ end
40
+
41
+ ##
42
+ # Sends a command to the server.
43
+ #
44
+ # @return a response string
45
+ def send_command(name)
46
+ response = ''
47
+ socket.puts(name)
48
+ while true do
49
+ if buf = socket.recv_nonblock(65536) rescue nil
50
+ response << buf
51
+ return response if response =~ /\n.\n$/
52
+ end
53
+ end
54
+ end
55
+
56
+ ##
57
+ # Returns results of a 'status' command.
58
+ #
59
+ # @return a hash of abilities with queued, active and workers keys.
60
+ def status
61
+ status = {}
62
+ if response = send_command('status')
63
+ response.split("\n").each do |line|
64
+ if line.match /^([A-Za-z_]+)\t([A-Za-z_]+)\t(\d+)\t(\d+)\t(\d+)$/
65
+ (status[$1] ||= {})[$2] = { :queue => $3, :active => $4, :workers => $5 }
66
+ end
67
+ end
68
+ end
69
+ status
70
+ end
71
+
72
+ ##
73
+ # Returns results of a 'workers' command.
74
+ #
75
+ # @return an array of worker hashes, containing host, status and functions keys.
76
+ def workers
77
+ workers = []
78
+ if response = send_command('workers')
79
+ response.split("\n").each do |line|
80
+ if line.match /^(\d+)\s([a-z0-9\:]+)\s([A-Z\-])\s:\s([a-z_\s\t]+)$/
81
+ func_parts = $4.split(' ')
82
+ functions = []
83
+ while !func_parts.empty?
84
+ functions << func_parts.shift + '.' + func_parts.shift
85
+ end
86
+ workers << { :host => $2, :status => $3, :functions => functions }
87
+ end
88
+ end
89
+ end
90
+ workers
91
+ end
92
+ end
93
+
94
+ end
@@ -0,0 +1,128 @@
1
+ #!/usr/bin/env ruby
2
+
3
+ module Gearman
4
+
5
+ # = Task
6
+ #
7
+ # == Description
8
+ # A task submitted to a Gearman job server.
9
+ class Task
10
+ ##
11
+ # Create a new Task object.
12
+ #
13
+ # @param func function name
14
+ # @param arg argument to the function
15
+ # @param opts hash of additional options
16
+ def initialize(func, arg='', opts={})
17
+ @func = func.to_s
18
+ @arg = arg or '' # TODO: use something more ref-like?
19
+ %w{uniq on_complete on_fail on_retry on_status retry_count
20
+ high_priority}.map {|s| s.to_sym }.each do |k|
21
+ instance_variable_set "@#{k}", opts[k]
22
+ opts.delete k
23
+ end
24
+ if opts.size > 0
25
+ raise InvalidArgsError, 'Invalid task args: ' + opts.keys.sort.join(', ')
26
+ end
27
+ @retry_count ||= 0
28
+ @successful = false
29
+ @retries_done = 0
30
+ @hash = nil
31
+ end
32
+ attr_accessor :uniq, :retry_count, :high_priority
33
+ attr_reader :successful, :func, :arg
34
+
35
+ ##
36
+ # Internal method to reset this task's state so it can be run again.
37
+ # Called by TaskSet#add_task.
38
+ def reset_state
39
+ @retries_done = 0
40
+ @successful = false
41
+ self
42
+ end
43
+
44
+ ##
45
+ # Set a block of code to be executed when this task completes
46
+ # successfully. The returned data will be passed to the block.
47
+ def on_complete(&f)
48
+ @on_complete = f
49
+ end
50
+
51
+ ##
52
+ # Set a block of code to be executed when this task fails.
53
+ def on_fail(&f)
54
+ @on_fail = f
55
+ end
56
+
57
+ ##
58
+ # Set a block of code to be executed when this task is retried after
59
+ # failing. The number of retries that have been attempted (including the
60
+ # current one) will be passed to the block.
61
+ def on_retry(&f)
62
+ @on_retry = f
63
+ end
64
+
65
+ ##
66
+ # Set a block of code to be executed when we receive a status update for
67
+ # this task. The block will receive two arguments, a numerator and
68
+ # denominator describing the task's status.
69
+ def on_status(&f)
70
+ @on_status = f
71
+ end
72
+
73
+ ##
74
+ # Handle completion of the task.
75
+ #
76
+ # @param data data returned from the server (doesn't include handle)
77
+ def handle_completion(data)
78
+ @successful = true
79
+ @on_complete.call(data) if @on_complete
80
+ self
81
+ end
82
+
83
+ ##
84
+ # Record a failure and check whether we should be retried.
85
+ #
86
+ # @return true if we should be resubmitted; false otherwise
87
+ def handle_failure
88
+ if @retries_done >= @retry_count
89
+ @on_fail.call if @on_fail
90
+ return false
91
+ end
92
+ @retries_done += 1
93
+ @on_retry.call(@retries_done) if @on_retry
94
+ true
95
+ end
96
+
97
+ ##
98
+ # Handle a status update for the task.
99
+ def handle_status(numerator, denominator)
100
+ @on_status.call(numerator, denominator) if @on_status
101
+ self
102
+ end
103
+
104
+ ##
105
+ # Return a hash that we can use to execute identical tasks on the same
106
+ # job server.
107
+ #
108
+ # @return hashed value, based on @arg if @uniq is '-', on @uniq if it's
109
+ # set to something else, and just nil if @uniq is nil
110
+ def get_uniq_hash
111
+ return @hash if @hash
112
+ merge_on = (@uniq and @uniq == '-') ? @arg : @uniq
113
+ @hash = merge_on ? merge_on.hash.to_s : ''
114
+ end
115
+
116
+ ##
117
+ # Construct a packet to submit this task to a job server.
118
+ #
119
+ # @param background ??
120
+ # @return String representation of packet
121
+ def get_submit_packet(background=false)
122
+ mode = 'submit_job' +
123
+ (background ? '_bg' : @high_priority ? '_high' : '')
124
+ Util::pack_request(mode, [func, get_uniq_hash, arg].join("\0"))
125
+ end
126
+ end
127
+
128
+ end
@@ -0,0 +1,241 @@
1
+ #!/usr/bin/env ruby
2
+
3
+ require 'socket'
4
+ require 'time'
5
+
6
+ module Gearman
7
+
8
+ # = TaskSet
9
+ #
10
+ # == Description
11
+ # A set of tasks submitted to a Gearman job server.
12
+ class TaskSet
13
+ def initialize(client)
14
+ @client = client
15
+ @task_waiting_for_handle = nil
16
+ @tasks_in_progress = {} # "host:port//handle" -> [job1, job2, ...]
17
+ @finished_tasks = [] # tasks that have completed or failed
18
+ @sockets = {} # "host:port" -> Socket
19
+ @merge_hash_to_hostport = {} # Fixnum -> "host:port"
20
+ end
21
+
22
+ ##
23
+ # Add a new task to this TaskSet.
24
+ #
25
+ # @param args either a Task or arguments for Task.new
26
+ # @return true if the task was created successfully, false otherwise
27
+ def add_task(*args)
28
+ task = Util::get_task_from_args(*args)
29
+ add_task_internal(task, true)
30
+ end
31
+
32
+ ##
33
+ # Internal function to add a task.
34
+ #
35
+ # @param task Task to add
36
+ # @param reset_state should we reset task state? true if we're adding a
37
+ # new task; false if we're rescheduling one that's
38
+ # failed
39
+ # @return true if the task was created successfully, false
40
+ # otherwise
41
+ def add_task_internal(task, reset_state=true)
42
+ task.reset_state if reset_state
43
+ req = task.get_submit_packet()
44
+
45
+ @task_waiting_for_handle = task
46
+ # FIXME: We need to loop here in case we get a bad job server, or the
47
+ # job creation fails (see how the server reports this to us), or ...
48
+
49
+ merge_hash = task.get_uniq_hash
50
+
51
+ looking_for_socket = true
52
+
53
+ should_try_rehash = true
54
+ while(looking_for_socket)
55
+ begin
56
+ hostport = if should_try_rehash
57
+ (@merge_hash_to_hostport[merge_hash] or @client.get_job_server)
58
+ else
59
+ @client.get_job_server
60
+ end
61
+
62
+ @merge_hash_to_hostport[merge_hash] = hostport if merge_hash
63
+ sock = (@sockets[hostport] or @client.get_socket(hostport))
64
+ looking_for_socket = false
65
+ rescue RuntimeError
66
+ should_try_rehash = false
67
+ end
68
+ end
69
+ Util.log "Using socket #{sock.inspect} for #{hostport}"
70
+ Util.send_request(sock, req)
71
+ while @task_waiting_for_handle
72
+ begin
73
+ read_packet(sock, @client.task_create_timeout_sec)
74
+ rescue NetworkError
75
+ Util.log "Got timeout on read from #{hostport}"
76
+ @task_waiting_for_handle = nil
77
+ @client.close_socket(sock)
78
+ return false
79
+ end
80
+ end
81
+
82
+ @sockets[hostport] ||= sock
83
+ true
84
+ end
85
+ private :add_task_internal
86
+
87
+ ##
88
+ # Handle a 'job_created' response from a job server.
89
+ #
90
+ # @param hostport "host:port" of job server
91
+ # @param data data returned in packet from server
92
+ def handle_job_created(hostport, data)
93
+ Util.log "Got job_created with handle #{data} from #{hostport}"
94
+ if not @task_waiting_for_handle
95
+ raise ProtocolError, "Got unexpected job_created notification " +
96
+ "with handle #{data} from #{hostport}"
97
+ end
98
+ js_handle = Util.handle_to_str(hostport, data)
99
+ task = @task_waiting_for_handle
100
+ @task_waiting_for_handle = nil
101
+ (@tasks_in_progress[js_handle] ||= []) << task
102
+ nil
103
+ end
104
+ private :handle_job_created
105
+
106
+ ##
107
+ # Handle a 'work_complete' response from a job server.
108
+ #
109
+ # @param hostport "host:port" of job server
110
+ # @param data data returned in packet from server
111
+ def handle_work_complete(hostport, data)
112
+ handle, data = data.split("\0", 2)
113
+ Util.log "Got work_complete with handle #{handle} and " +
114
+ "#{data ? data.size : '0'} byte(s) of data from #{hostport}"
115
+ js_handle = Util.handle_to_str(hostport, handle)
116
+ tasks = @tasks_in_progress.delete(js_handle)
117
+ if not tasks
118
+ raise ProtocolError, "Got unexpected work_complete with handle " +
119
+ "#{handle} from #{hostport} (no task by that name)"
120
+ end
121
+ tasks.each do |t|
122
+ t.handle_completion(data)
123
+ @finished_tasks << t
124
+ end
125
+ nil
126
+ end
127
+
128
+ private :handle_work_complete
129
+
130
+ ##
131
+ # Handle a 'work_fail' response from a job server.
132
+ #
133
+ # @param hostport "host:port" of job server
134
+ # @param data data returned in packet from server
135
+ def handle_work_fail(hostport, data)
136
+ Util.log "Got work_fail with handle #{data} from #{hostport}"
137
+ js_handle = Util.handle_to_str(hostport, data)
138
+ tasks = @tasks_in_progress.delete(js_handle)
139
+ if not tasks
140
+ raise ProtocolError, "Got unexpected work_fail with handle " +
141
+ "#{data} from #{hostport} (no task by that name)"
142
+ end
143
+ tasks.each do |t|
144
+ if t.handle_failure
145
+ add_task_internal(t, false)
146
+ else
147
+ @finished_tasks << t
148
+ end
149
+ end
150
+ end
151
+ private :handle_work_fail
152
+
153
+ ##
154
+ # Handle a 'work_status' response from a job server.
155
+ #
156
+ # @param hostport "host:port" of job server
157
+ # @param data data returned in packet from server
158
+ def handle_work_status(hostport, data)
159
+ handle, num, den = data.split("\0", 3)
160
+ Util.log "Got work_status with handle #{handle} from #{hostport}: " +
161
+ "#{num}/#{den}"
162
+ js_handle = Util.handle_to_str(hostport, handle)
163
+ tasks = @tasks_in_progress[js_handle]
164
+ if not tasks
165
+ raise ProtocolError, "Got unexpected work_status with handle " +
166
+ "#{handle} from #{hostport} (no task by that name)"
167
+ end
168
+ tasks.each {|t| t.handle_status(num, den) }
169
+ end
170
+ private :handle_work_status
171
+
172
+ ##
173
+ # Read and process a packet from a socket.
174
+ #
175
+ # @param sock socket connected to a job server
176
+ def read_packet(sock, timeout=nil)
177
+ hostport = @client.get_hostport_for_socket(sock)
178
+ if not hostport
179
+ raise RuntimeError, "Client doesn't know host/port for socket " +
180
+ sock.inspect
181
+ end
182
+ type, data = Util.read_response(sock, timeout)
183
+ case type
184
+ when :job_created
185
+ handle_job_created(hostport, data)
186
+ when :work_complete
187
+ handle_work_complete(hostport, data)
188
+ when :work_fail
189
+ handle_work_fail(hostport, data)
190
+ when :work_status
191
+ handle_work_status(hostport, data)
192
+ else
193
+ Util.log "Got #{type.to_s} from #{hostport}"
194
+ end
195
+ nil
196
+ end
197
+ private :read_packet
198
+
199
+ ##
200
+ # Wait for all tasks in the set to finish.
201
+ #
202
+ # @param timeout maximum amount of time to wait, in seconds
203
+ def wait(timeout=1)
204
+ end_time = Time.now.to_f + timeout
205
+ while not @tasks_in_progress.empty?
206
+ remaining = end_time - Time.now.to_f
207
+ ready_socks = IO::select(
208
+ @sockets.values, nil, nil, remaining > 0 ? remaining : 0)
209
+ if not ready_socks or not ready_socks[0]
210
+ Util.log "Timed out while waiting for tasks to finish"
211
+ # not sure what state the connections are in, so just be lame and
212
+ # close them for now
213
+ @sockets.values.each {|s| @client.close_socket(s) }
214
+ @sockets = {}
215
+ return false
216
+ end
217
+ ready_socks[0].each do |sock|
218
+ begin
219
+ read_packet(sock, end_time - Time.now.to_f)
220
+ rescue ProtocolError
221
+ hostport = @client.get_hostport_for_socket(sock)
222
+ Util.log "Ignoring bad packet from #{hostport}"
223
+ rescue NetworkError
224
+ hostport = @client.get_hostport_for_socket(sock)
225
+ Util.log "Got timeout on read from #{hostport}"
226
+ end
227
+ end
228
+ end
229
+ @sockets.values.each {|s| @client.return_socket(s) }
230
+ @sockets = {}
231
+ @finished_tasks.each do |t|
232
+ if not t.successful
233
+ Util.log "Taskset failed"
234
+ return false
235
+ end
236
+ end
237
+ true
238
+ end
239
+ end
240
+
241
+ end
@@ -0,0 +1,96 @@
1
+ #!/usr/bin/env ruby
2
+
3
+ require 'socket'
4
+ require 'thread'
5
+
6
+ class FakeJobServer
7
+ def initialize(tester,port=nil)
8
+ @tester = tester
9
+ @serv = TCPserver.open(0) if port.nil?
10
+ @serv = TCPserver.open('localhost',port) unless port.nil?
11
+ @port = @serv.addr[1]
12
+ end
13
+ attr_reader :port
14
+
15
+ def server_socket
16
+ @serv
17
+ end
18
+
19
+ def stop
20
+ @serv.close
21
+ end
22
+
23
+ def start
24
+ @serv = TCPserver.open(@port)
25
+ end
26
+
27
+ def expect_connection
28
+ sock = @serv.accept
29
+ return sock
30
+ end
31
+
32
+ def expect_closed(sock)
33
+ @tester.assert_true(sock.closed?)
34
+ end
35
+
36
+ def expect_request(sock, exp_type, exp_data='')
37
+ head = sock.recv(12)
38
+ magic, type, len = head.unpack('a4NN')
39
+ @tester.assert_equal("\0REQ", magic)
40
+ @tester.assert_equal(Gearman::Util::NUMS[exp_type.to_sym], type)
41
+ data = len > 0 ? sock.recv(len) : ''
42
+ @tester.assert_equal(exp_data, data)
43
+ end
44
+
45
+ def expect_any_request(sock)
46
+ head = sock.recv(12)
47
+ end
48
+
49
+ def expect_anything_and_close_socket(sock)
50
+ head = sock.recv(12)
51
+ sock.close
52
+ end
53
+
54
+ def send_response(sock, type, data='', bogus_size=nil)
55
+ type_num = Gearman::Util::NUMS[type.to_sym]
56
+ raise RuntimeError, "Invalid type #{type}" if not type_num
57
+ response = "\0RES" + [type_num, (bogus_size or data.size)].pack('NN') + data
58
+ sock.write(response)
59
+ end
60
+ end
61
+
62
+ class TestScript
63
+ def initialize
64
+ @mutex = Mutex.new
65
+ @cv = ConditionVariable.new
66
+ @blocks = []
67
+ end
68
+
69
+ def loop_forever
70
+ loop do
71
+ f = nil
72
+ @mutex.synchronize do
73
+ @cv.wait(@mutex) if @blocks.empty?
74
+ f = @blocks[0] if not @blocks.empty?
75
+ end
76
+ f.call if f
77
+ @mutex.synchronize do
78
+ @blocks.shift
79
+ @cv.signal if @blocks.empty?
80
+ end
81
+ end
82
+ end
83
+
84
+ def exec(&f)
85
+ @mutex.synchronize do
86
+ @blocks << f
87
+ @cv.signal
88
+ end
89
+ end
90
+
91
+ def wait
92
+ @mutex.synchronize do
93
+ @cv.wait(@mutex) if not @blocks.empty?
94
+ end
95
+ end
96
+ end