xing-gearman-ruby 1.0.0

data/lib/gearman/client.rb

@@ -0,0 +1,139 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+
+module Gearman
+
+# = Client
+#
+# == Description
+# A client for communicating with Gearman job servers.
+class Client
+  ##
+  # Create a new client.
+  #
+  # @param job_servers  "host:port"; either a single server or an array
+  def initialize(job_servers=nil)
+    @job_servers = []  # "host:port"
+    self.job_servers = job_servers if job_servers
+    @sockets = {}  # "host:port" -> [sock1, sock2, ...]
+    @socket_to_hostport = {}  # sock -> "host:port"
+    @test_hostport = nil  # make get_job_server return a given host for testing
+    @task_create_timeout_sec = 10
+    @server_counter = -1
+    @bad_servers = []
+  end
+  attr_reader :job_servers, :bad_servers  
+  attr_accessor :test_hostport, :task_create_timeout_sec  
+
+  ##
+  # Set the job servers to be used by this client.
+  #
+  # @param servers  "host:port"; either a single server or an array
+  def job_servers=(servers)
+    @job_servers = Util.normalize_job_servers(servers)
+    self
+  end
+
+  ##
+  # Get connection info about an arbitrary (currently random, but maybe
+  # we'll do something smarter later) job server.
+  #
+  # @return  "host:port"
+  def get_job_server
+
+    raise Exception.new('No servers available') if @job_servers.empty?
+
+    @server_counter += 1
+    # Return a specific server if one's been set.
+    @test_hostport or @job_servers[@server_counter % @job_servers.size]
+  end
+  
+  def signal_bad_server(hostport)
+    @job_servers = @job_servers.reject { |s| s == hostport }
+    @bad_servers << hostport
+  end
+  ##
+  # Get a socket for a job server.
+  #
+  # @param hostport  job server "host:port"
+  # @return          a Socket
+  def get_socket(hostport, num_retries=3)
+    # If we already have an open socket to this host, return it.
+    if @sockets[hostport]
+      sock = @sockets[hostport].shift
+      @sockets.delete(hostport) if @sockets[hostport].size == 0
+      return sock
+    end
+
+    num_retries.times do |i|
+      begin
+        sock = TCPSocket.new(*hostport.split(':'))
+      rescue Exception
+      else
+
+        @socket_to_hostport[sock] = hostport
+        return sock
+      end
+    end
+
+    signal_bad_server(hostport)
+    raise RuntimeError, "Unable to connect to job server #{hostport}"
+  end
+
+  ##
+  # Relinquish a socket created by Client#get_socket.
+  #
+  # If we don't know about the socket, we just close it.
+  #
+  # @param sock  Socket
+  def return_socket(sock)
+    hostport = get_hostport_for_socket(sock)
+    if not hostport
+      inet, port, host, ip = s.addr
+      Util.err "Got socket for #{ip}:#{port}, which we don't " +
+        "know about -- closing"
+      sock.close
+      return
+    end
+    (@sockets[hostport] ||= []) << sock
+  end
+
+  def close_socket(sock)
+    sock.close
+    @socket_to_hostport.delete(sock)
+    nil
+  end
+
+  ##
+  # Given a socket from Client#get_socket, return its host and port.
+  #
+  # @param sock  Socket
+  # @return      "host:port", or nil if unregistered (which shouldn't happen)
+  def get_hostport_for_socket(sock)
+    @socket_to_hostport[sock]
+  end
+
+  ##
+  # Perform a single task.
+  #
+  # @param args  either a Task or arguments for Task.new
+  # @return      output of the task, or nil on failure
+  def do_task(*args)
+    task = Util::get_task_from_args(*args)
+
+    result = nil
+    failed = false
+    task.on_complete {|v| result = v }
+    task.on_fail { failed = true }
+
+    taskset = TaskSet.new(self)
+    taskset.add_task(task)
+    taskset.wait
+
+    failed ? nil : result
+  end
+
+end
+
+end

data/lib/gearman/server.rb

@@ -0,0 +1,94 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+require 'gearman'
+
+module Gearman
+
+# = Server
+#
+# == Description
+# A client for managing Gearman job servers.
+class Server
+  ##
+  # Create a new client.
+  #
+  # @param job_servers  "host:port"; either a single server or an array
+  # @param prefix       function name prefix (namespace)
+  def initialize(hostport)
+    @hostport = hostport  # "host:port"
+  end
+  attr_reader :hostport
+
+  ##
+  # Get a socket for a job server.
+  #
+  # @param hostport  job server "host:port"
+  # @return          a Socket
+  def socket(num_retries=3)
+    return @socket if @socket
+    num_retries.times do
+      begin
+        sock = TCPSocket.new(*hostport.split(':'))
+      rescue Exception
+      else
+        return @socket = sock
+      end
+    end
+    raise RuntimeError, "Unable to connect to job server #{hostport}"
+  end
+
+  ##
+  # Sends a command to the server.
+  #
+  # @return a response string
+  def send_command(name)
+    response = ''
+    socket.puts(name)
+    while true do 
+      if buf = socket.recv_nonblock(65536) rescue nil
+        response << buf 
+        return response if response =~ /\n.\n$/
+      end
+    end
+  end
+  
+  ##
+  # Returns results of a 'status' command.
+  #
+  # @return a hash of abilities with queued, active and workers keys.
+  def status
+    status = {}
+    if response = send_command('status')
+      response.split("\n").each do |line|
+        if line.match /^([A-Za-z_]+)\t([A-Za-z_]+)\t(\d+)\t(\d+)\t(\d+)$/
+          (status[$1] ||= {})[$2] = { :queue => $3, :active => $4, :workers => $5 }
+        end
+      end
+    end
+    status
+  end
+  
+  ##
+  # Returns results of a 'workers' command.
+  #
+  # @return an array of worker hashes, containing host, status and functions keys.
+  def workers
+    workers = []
+    if response = send_command('workers')
+      response.split("\n").each do |line|
+        if line.match /^(\d+)\s([a-z0-9\:]+)\s([A-Z\-])\s:\s([a-z_\s\t]+)$/
+          func_parts = $4.split(' ')
+          functions = []
+          while !func_parts.empty?
+            functions << func_parts.shift + '.' + func_parts.shift
+          end
+          workers << { :host => $2, :status => $3, :functions => functions }
+        end
+      end
+    end
+    workers
+  end
+end
+
+end

data/lib/gearman/task.rb

@@ -0,0 +1,128 @@
+#!/usr/bin/env ruby
+
+module Gearman
+
+# = Task
+#
+# == Description
+# A task submitted to a Gearman job server.
+class Task
+  ##
+  # Create a new Task object.
+  #
+  # @param func  function name
+  # @param arg   argument to the function
+  # @param opts  hash of additional options
+  def initialize(func, arg='', opts={})
+    @func = func.to_s
+    @arg = arg or ''  # TODO: use something more ref-like?
+    %w{uniq on_complete on_fail on_retry on_status retry_count
+       high_priority}.map {|s| s.to_sym }.each do |k|
+      instance_variable_set "@#{k}", opts[k]
+      opts.delete k
+    end
+    if opts.size > 0
+      raise InvalidArgsError, 'Invalid task args: ' + opts.keys.sort.join(', ')
+    end
+    @retry_count ||= 0
+    @successful = false
+    @retries_done = 0
+    @hash = nil
+  end
+  attr_accessor :uniq, :retry_count, :high_priority
+  attr_reader :successful, :func, :arg
+
+  ##
+  # Internal method to reset this task's state so it can be run again.
+  # Called by TaskSet#add_task.
+  def reset_state
+    @retries_done = 0
+    @successful = false
+    self
+  end
+
+  ##
+  # Set a block of code to be executed when this task completes
+  # successfully.  The returned data will be passed to the block.
+  def on_complete(&f)
+    @on_complete = f
+  end
+
+  ##
+  # Set a block of code to be executed when this task fails.
+  def on_fail(&f)
+    @on_fail = f
+  end
+
+  ##
+  # Set a block of code to be executed when this task is retried after
+  # failing.  The number of retries that have been attempted (including the
+  # current one) will be passed to the block.
+  def on_retry(&f)
+    @on_retry = f
+  end
+
+  ##
+  # Set a block of code to be executed when we receive a status update for
+  # this task.  The block will receive two arguments, a numerator and
+  # denominator describing the task's status.
+  def on_status(&f)
+    @on_status = f
+  end
+
+  ##
+  # Handle completion of the task.
+  #
+  # @param data  data returned from the server (doesn't include handle)
+  def handle_completion(data)
+    @successful = true
+    @on_complete.call(data) if @on_complete
+    self
+  end
+
+  ##
+  # Record a failure and check whether we should be retried.
+  #
+  # @return  true if we should be resubmitted; false otherwise
+  def handle_failure
+    if @retries_done >= @retry_count
+      @on_fail.call if @on_fail
+      return false
+    end
+    @retries_done += 1
+    @on_retry.call(@retries_done) if @on_retry
+    true
+  end
+
+  ##
+  # Handle a status update for the task.
+  def handle_status(numerator, denominator)
+    @on_status.call(numerator, denominator) if @on_status
+    self
+  end
+
+  ##
+  # Return a hash that we can use to execute identical tasks on the same
+  # job server.
+  #
+  # @return  hashed value, based on @arg if @uniq is '-', on @uniq if it's
+  #          set to something else, and just nil if @uniq is nil
+  def get_uniq_hash
+    return @hash if @hash
+    merge_on = (@uniq and @uniq == '-') ? @arg : @uniq
+    @hash = merge_on ? merge_on.hash.to_s : ''
+  end
+
+  ##
+  # Construct a packet to submit this task to a job server.
+  #
+  # @param background  ??
+  # @return            String representation of packet
+  def get_submit_packet(background=false)
+    mode = 'submit_job' +
+      (background ? '_bg' : @high_priority ? '_high' : '')
+    Util::pack_request(mode, [func, get_uniq_hash, arg].join("\0"))
+  end
+end
+
+end

data/lib/gearman/taskset.rb

@@ -0,0 +1,241 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+require 'time'
+
+module Gearman
+
+# = TaskSet
+#
+# == Description
+# A set of tasks submitted to a Gearman job server.
+class TaskSet
+  def initialize(client)
+    @client = client
+    @task_waiting_for_handle = nil
+    @tasks_in_progress = {}  # "host:port//handle" -> [job1, job2, ...]
+    @finished_tasks = []  # tasks that have completed or failed
+    @sockets = {}  # "host:port" -> Socket
+    @merge_hash_to_hostport = {}  # Fixnum -> "host:port"
+  end
+
+  ##
+  # Add a new task to this TaskSet.
+  #
+  # @param args  either a Task or arguments for Task.new
+  # @return      true if the task was created successfully, false otherwise
+  def add_task(*args)
+    task = Util::get_task_from_args(*args)
+    add_task_internal(task, true)
+  end
+
+  ##
+  # Internal function to add a task.
+  #
+  # @param task         Task to add
+  # @param reset_state  should we reset task state?  true if we're adding a
+  #                     new task; false if we're rescheduling one that's
+  #                     failed
+  # @return             true if the task was created successfully, false
+  #                     otherwise
+  def add_task_internal(task, reset_state=true)
+    task.reset_state if reset_state
+    req = task.get_submit_packet()
+
+    @task_waiting_for_handle = task
+    # FIXME: We need to loop here in case we get a bad job server, or the
+    # job creation fails (see how the server reports this to us), or ...
+
+    merge_hash = task.get_uniq_hash
+
+    looking_for_socket = true
+
+    should_try_rehash = true
+    while(looking_for_socket)
+      begin
+        hostport = if should_try_rehash
+          (@merge_hash_to_hostport[merge_hash] or @client.get_job_server)
+        else
+          @client.get_job_server
+        end
+
+        @merge_hash_to_hostport[merge_hash] = hostport if merge_hash
+        sock = (@sockets[hostport] or @client.get_socket(hostport))
+        looking_for_socket = false
+      rescue RuntimeError
+        should_try_rehash = false
+      end
+    end
+    Util.log "Using socket #{sock.inspect} for #{hostport}"
+    Util.send_request(sock, req)
+    while @task_waiting_for_handle
+      begin
+        read_packet(sock, @client.task_create_timeout_sec)
+      rescue NetworkError
+        Util.log "Got timeout on read from #{hostport}"
+        @task_waiting_for_handle = nil
+        @client.close_socket(sock)
+        return false
+      end
+    end
+
+    @sockets[hostport] ||= sock
+    true
+  end
+  private :add_task_internal
+
+  ##
+  # Handle a 'job_created' response from a job server.
+  #
+  # @param hostport  "host:port" of job server
+  # @param data      data returned in packet from server
+  def handle_job_created(hostport, data)
+    Util.log "Got job_created with handle #{data} from #{hostport}"
+    if not @task_waiting_for_handle
+      raise ProtocolError, "Got unexpected job_created notification " +
+        "with handle #{data} from #{hostport}"
+    end
+    js_handle = Util.handle_to_str(hostport, data)
+    task = @task_waiting_for_handle
+    @task_waiting_for_handle = nil
+    (@tasks_in_progress[js_handle] ||= []) << task
+    nil
+  end
+  private :handle_job_created
+
+  ##
+  # Handle a 'work_complete' response from a job server.
+  #
+  # @param hostport  "host:port" of job server
+  # @param data      data returned in packet from server
+  def handle_work_complete(hostport, data)
+    handle, data = data.split("\0", 2)
+    Util.log "Got work_complete with handle #{handle} and " +
+      "#{data ? data.size : '0'} byte(s) of data from #{hostport}"
+    js_handle = Util.handle_to_str(hostport, handle)
+    tasks = @tasks_in_progress.delete(js_handle)
+    if not tasks
+      raise ProtocolError, "Got unexpected work_complete with handle " +
+        "#{handle} from #{hostport} (no task by that name)"
+    end
+    tasks.each do |t|
+      t.handle_completion(data)
+      @finished_tasks << t
+    end
+    nil
+  end
+
+  private :handle_work_complete
+
+  ##
+  # Handle a 'work_fail' response from a job server.
+  #
+  # @param hostport  "host:port" of job server
+  # @param data      data returned in packet from server
+  def handle_work_fail(hostport, data)
+    Util.log "Got work_fail with handle #{data} from #{hostport}"
+    js_handle = Util.handle_to_str(hostport, data)
+    tasks = @tasks_in_progress.delete(js_handle)
+    if not tasks
+      raise ProtocolError, "Got unexpected work_fail with handle " +
+        "#{data} from #{hostport} (no task by that name)"
+    end
+    tasks.each do |t|
+      if t.handle_failure
+        add_task_internal(t, false)
+      else
+        @finished_tasks << t
+      end
+    end
+  end
+  private :handle_work_fail
+
+  ##
+  # Handle a 'work_status' response from a job server.
+  #
+  # @param hostport  "host:port" of job server
+  # @param data      data returned in packet from server
+  def handle_work_status(hostport, data)
+    handle, num, den = data.split("\0", 3)
+    Util.log "Got work_status with handle #{handle} from #{hostport}: " +
+      "#{num}/#{den}"
+    js_handle = Util.handle_to_str(hostport, handle)
+    tasks = @tasks_in_progress[js_handle]
+    if not tasks
+      raise ProtocolError, "Got unexpected work_status with handle " +
+        "#{handle} from #{hostport} (no task by that name)"
+    end
+    tasks.each {|t| t.handle_status(num, den) }
+  end
+  private :handle_work_status
+
+  ##
+  # Read and process a packet from a socket.
+  #
+  # @param sock  socket connected to a job server
+  def read_packet(sock, timeout=nil)
+    hostport = @client.get_hostport_for_socket(sock)
+    if not hostport
+      raise RuntimeError, "Client doesn't know host/port for socket " +
+        sock.inspect
+    end
+    type, data = Util.read_response(sock, timeout)
+    case type
+    when :job_created
+      handle_job_created(hostport, data)
+    when :work_complete
+      handle_work_complete(hostport, data)
+    when :work_fail
+      handle_work_fail(hostport, data)
+    when :work_status
+      handle_work_status(hostport, data)
+    else
+      Util.log "Got #{type.to_s} from #{hostport}"
+    end
+    nil
+  end
+  private :read_packet
+
+  ##
+  # Wait for all tasks in the set to finish.
+  #
+  # @param timeout  maximum amount of time to wait, in seconds
+  def wait(timeout=1)
+    end_time = Time.now.to_f + timeout
+    while not @tasks_in_progress.empty?
+      remaining = end_time - Time.now.to_f
+      ready_socks = IO::select(
+        @sockets.values, nil, nil, remaining > 0 ? remaining : 0)
+      if not ready_socks or not ready_socks[0]
+        Util.log "Timed out while waiting for tasks to finish"
+        # not sure what state the connections are in, so just be lame and
+        # close them for now
+        @sockets.values.each {|s| @client.close_socket(s) }
+        @sockets = {}
+        return false
+      end
+      ready_socks[0].each do |sock|
+        begin
+          read_packet(sock, end_time - Time.now.to_f)
+        rescue ProtocolError
+          hostport = @client.get_hostport_for_socket(sock)
+          Util.log "Ignoring bad packet from #{hostport}"
+        rescue NetworkError
+          hostport = @client.get_hostport_for_socket(sock)
+          Util.log "Got timeout on read from #{hostport}"
+        end
+      end
+    end
+    @sockets.values.each {|s| @client.return_socket(s) }
+    @sockets = {}
+    @finished_tasks.each do |t|
+      if not t.successful
+        Util.log "Taskset failed"
+        return false
+      end
+    end
+    true
+  end
+end
+
+end

data/lib/gearman/testlib.rb

@@ -0,0 +1,96 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+require 'thread'
+
+class FakeJobServer
+  def initialize(tester,port=nil)
+    @tester = tester
+    @serv = TCPserver.open(0) if port.nil?
+    @serv = TCPserver.open('localhost',port) unless port.nil?
+    @port = @serv.addr[1]
+  end
+  attr_reader :port
+
+  def server_socket
+    @serv
+  end
+
+  def stop
+    @serv.close
+  end
+
+  def start
+    @serv = TCPserver.open(@port)
+  end
+
+  def expect_connection
+    sock = @serv.accept
+    return sock
+  end
+
+  def expect_closed(sock)
+    @tester.assert_true(sock.closed?)
+  end
+
+  def expect_request(sock, exp_type, exp_data='')
+    head = sock.recv(12)
+    magic, type, len = head.unpack('a4NN')
+    @tester.assert_equal("\0REQ", magic)
+    @tester.assert_equal(Gearman::Util::NUMS[exp_type.to_sym], type)
+    data = len > 0 ? sock.recv(len) : ''
+    @tester.assert_equal(exp_data, data)
+  end
+
+  def expect_any_request(sock)
+    head = sock.recv(12)
+  end
+
+  def expect_anything_and_close_socket(sock)
+    head = sock.recv(12)
+    sock.close
+  end
+
+  def send_response(sock, type, data='', bogus_size=nil)
+    type_num = Gearman::Util::NUMS[type.to_sym]
+    raise RuntimeError, "Invalid type #{type}" if not type_num
+    response = "\0RES" + [type_num, (bogus_size or data.size)].pack('NN') + data
+    sock.write(response)
+  end
+end
+
+class TestScript
+  def initialize
+    @mutex = Mutex.new
+    @cv = ConditionVariable.new
+    @blocks = []
+  end
+
+  def loop_forever
+    loop do
+      f = nil
+      @mutex.synchronize do
+        @cv.wait(@mutex) if @blocks.empty?
+        f = @blocks[0] if not @blocks.empty?
+      end
+      f.call if f
+      @mutex.synchronize do
+        @blocks.shift
+        @cv.signal if @blocks.empty?
+      end
+    end
+  end
+
+  def exec(&f)
+    @mutex.synchronize do
+      @blocks << f
+      @cv.signal
+    end
+  end
+
+  def wait
+    @mutex.synchronize do
+      @cv.wait(@mutex) if not @blocks.empty?
+    end
+  end
+end