gearman-ruby 2.0.0 → 3.0.1

data/lib/gearman/testlib.rb

@@ -0,0 +1,95 @@
+#!/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='', size=12)
+    head = sock.recv(size)
+    magic, type, len = head.unpack('a4NN')
+    @tester.assert("\0REQ" == magic || "\000REQ" == 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] || 0
+    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

data/lib/gearman/util.rb

@@ -2,51 +2,207 @@
 
 require 'socket'
 require 'time'
+require 'logger'
 
 module Gearman
 
   class ServerDownException < Exception; end
 
-  # = Util
-  #
-  # == Description
-  # Static helper methods and data used by other classes.
-  class Util
+# = Util
+#
+# == Description
+# Static helper methods and data used by other classes.
+class Util
+  # Map from Integer representations of commands used in the network
+  # protocol to more-convenient symbols.
+  COMMANDS = {
+    1  => :can_do,               # W->J: FUNC
+    2  => :cant_do,              # W->J: FUNC
+    3  => :reset_abilities,      # W->J: --
+    4  => :pre_sleep,            # W->J: --
+    #5 =>  (unused),             # -      -
+    6  => :noop,                 # J->W: --
+    7  => :submit_job,           # C->J: FUNC[0]UNIQ[0]ARGS
+    8  => :job_created,          # J->C: HANDLE
+    9  => :grab_job,             # W->J: --
+    10 => :no_job,               # J->W: --
+    11 => :job_assign,           # J->W: HANDLE[0]FUNC[0]ARG
+    12 => :work_status,          # W->J/C: HANDLE[0]NUMERATOR[0]DENOMINATOR
+    13 => :work_complete,        # W->J/C: HANDLE[0]RES
+    14 => :work_fail,            # W->J/C: HANDLE
+    15 => :get_status,           # C->J: HANDLE
+    16 => :echo_req,             # ?->J: TEXT
+    17 => :echo_res,             # J->?: TEXT
+    18 => :submit_job_bg,        # C->J: FUNC[0]UNIQ[0]ARGS
+    19 => :error,                # J->?: ERRCODE[0]ERR_TEXT
+    20 => :status_res,           # C->J: HANDLE[0]KNOWN[0]RUNNING[0]NUM[0]DENOM
+    21 => :submit_job_high,      # C->J: FUNC[0]UNIQ[0]ARGS
+    22 => :set_client_id,        # W->J: [RANDOM_STRING_NO_WHITESPACE]
+    23 => :can_do_timeout,       # W->J: FUNC[0]TIMEOUT
+    24 => :all_yours,            # REQ    Worker
+    25 => :work_exception,       # W->J: HANDLE[0]ARG
+    26 => :option_req,           # C->J: TEXT
+    27 => :option_res,           # J->C: TEXT
+    28 => :work_data,            # REQ    Worker
+    29 => :work_warning,         # W->J/C: HANDLE[0]MSG
+    30 => :grab_job_uniq,        # REQ    Worker
+    31 => :job_assign_uniq,      # RES    Worker
+    32 => :submit_job_high_bg,   # C->J: FUNC[0]UNIQ[0]ARGS
+    33 => :submit_job_low,       # C->J: FUNC[0]UNIQ[0]ARGS
+    34 => :submit_job_low_bg,    # C->J: FUNC[0]UNIQ[0]ARGS
+    35 => :submit_job_sched,     # REQ    Client
+    36 => :submit_job_epoch      # C->J: FUNC[0]UNIQ[0]EPOCH[0]ARGS
+  }
+
+  # Map e.g. 'can_do' => 1
+  NUMS = COMMANDS.invert
+
+  # Default job server port.
+  DEFAULT_PORT = 4730
+
+  def Util.logger=(logger)
+    @logger = logger
+  end
+
+  def Util.logger
+    @logger ||=
+      begin
+        l = Logger.new($stdout)
+        l.level = Logger::FATAL
+        l
+      end
+  end
 
-    @@debug = false
+  ##
+  # Construct a request packet.
+  #
+  # @param type_name  command type's name (see COMMANDS)
+  # @param arg        optional data to pack into the command
+  # @return           packet (as a string)
+  def Util.pack_request(type_name, arg='')
+    type_num = NUMS[type_name.to_sym]
+    raise InvalidArgsError, "Invalid type name '#{type_name}'" unless type_num
+    arg = '' if not arg
+    "\0REQ" + [type_num, arg.size].pack('NN') + arg
+  end
 
-    ##
-    # Enable or disable debugging output (off by default).
-    #
-    # @param v  print debugging output
-    def Util.debug=(v)
-      @@debug = v
+  ##
+  # Return a Task based on the passed-in arguments.
+  #
+  # @param args  either a single Task object or the arguments accepted by
+  #              Task.new
+  # @return      Task object
+  def Util.get_task_from_args(*args)
+    if (args[0].class == Task || args[0].class.superclass == Task)
+      return args[0]
+    elsif args.size <= 3
+      return Task.new(*args)
+    else
+      raise InvalidArgsError, 'Incorrect number of args to get_task_from_args'
     end
+  end
 
-    ##
-    # Log a message if debugging is enabled.
-    #
-    # @param str  message to log
-    def Util.log(str, force=false)
-      puts "#{Time.now.strftime '%Y-%m-%d %H:%M:%S'} #{str}" if force or @@debug
+  ##
+  # Read from a socket, giving up if it doesn't finish quickly enough.
+  # NetworkError is thrown if we don't read all the bytes in time.
+  #
+  # @param sock     Socket from which we read
+  # @param len      number of bytes to read
+  # @param timeout  maximum number of seconds we'll take; nil for no timeout
+  # @return         full data that was read
+  def Util.timed_recv(sock, len, timeout=nil)
+    data = ''
+    end_time = Time.now.to_f + timeout if timeout
+    while data.size < len and (not timeout or Time.now.to_f < end_time) do
+      IO::select([sock], nil, nil, timeout ? end_time - Time.now.to_f : nil) \
+        or break
+      data += sock.readpartial(len - data.size)
     end
+    if data.size < len
+      raise NetworkError, "Read #{data.size} byte(s) instead of #{len}"
+    end
+    data
+  end
+
+  ##
+  # Read a response packet from a socket.
+  #
+  # @param sock     Socket connected to a job server
+  # @param timeout  timeout in seconds, nil for no timeout
+  # @return         array consisting of integer packet type and data
+  def Util.read_response(sock, timeout=nil)
+    #debugger
+    end_time = Time.now.to_f + timeout if timeout
+    head = timed_recv(sock, 12, timeout)
+    magic, type, len = head.unpack('a4NN')
+    raise ProtocolError, "Invalid magic '#{magic}'" unless magic == "\0RES"
+    buf = len > 0 ?
+      timed_recv(sock, len, timeout ? end_time - Time.now.to_f : nil) : ''
+    type = COMMANDS[type]
+    raise ProtocolError, "Invalid packet type #{type}" unless type
+    [type, buf]
+  end
 
-    ##
-    # Log a message no matter what.
-    #
-    # @param str  message to log
-    def Util.err(str)
-      log(str, true)
+  ##
+  # Send a request packet over a socket.
+  #
+  # @param sock  Socket connected to a job server
+  # @param req   request packet to send
+  def Util.send_request(sock, req)
+    len = with_safe_socket_op{ sock.write(req) }
+    if len != req.size
+      raise NetworkError, "Wrote #{len} instead of #{req.size}"
     end
+  end
 
-    def Util.ability_name_with_prefix(prefix,name)
-      "#{prefix}\t#{name}"
+  ##
+  # Add default ports to a job server or list of servers.
+  #
+  # @param servers  a server hostname or "host:port" or array of servers
+  # @return         an array of "host:port" strings
+  def Util.normalize_job_servers(servers)
+    if servers.class == String or servers.class == Symbol
+      servers = [ servers.to_s ]
     end
+    servers.map {|s| s =~ /:/ ? s : "#{s}:#{DEFAULT_PORT}" }
+  end
+
+  ##
+  # Convert job server info and a handle into a string.
+  #
+  # @param hostport  "host:port" of job server
+  # @param handle    job server-returned handle for a task
+  # @return          "host:port//handle"
+  def Util.handle_to_str(hostport, handle)
+    "#{hostport}//#{handle}"
+  end
 
-    class << self
-      alias :ability_name_for_perl :ability_name_with_prefix
+  ##
+  # Reverse Util.handle_to_str.
+  #
+  # @param str  "host:port//handle"
+  # @return     [hostport, handle]
+  def Util.str_to_handle(str)
+    str =~ %r{^([^:]+:\d+)//(.+)}
+    return [$1, $3]
+  end
+
+  def self.with_safe_socket_op
+    begin
+      yield
+    rescue Exception => ex
+      raise ServerDownException.new(ex.message)
     end
+  end
 
+  def Util.ability_name_with_prefix(prefix,name)
+    "#{prefix}\t#{name}"
   end
 
+  class << self
+    alias :ability_name_for_perl :ability_name_with_prefix
+  end
+
+end
+
 end

data/lib/gearman/worker.rb

@@ -1,39 +1,375 @@
+#!/usr/bin/env ruby
+
+require 'set'
+require 'socket'
+require 'thread'
+
 module Gearman
-  class Worker
 
-    attr_reader :abilities
+# = Worker
+#
+# == Description
+# A worker that can connect to a Gearman server and perform tasks.
+#
+# == Usage
+#  require 'gearman'
+#
+#  w = Gearman::Worker.new('127.0.0.1')
+#
+#  # Add a handler for a "sleep" function that takes a single argument, the
+#  # number of seconds to sleep before reporting success.
+#  w.add_ability('sleep') do |data,job|
+#    seconds = data
+#    (1..seconds.to_i).each do |i|
+#      sleep 1
+#      # Report our progress to the job server every second.
+#      job.report_status(i, seconds)
+#    end
+#    # Report success.
+#    true
+#  end
+#  loop { w.work }
+class Worker
+  # = Ability
+  #
+  # == Description
+  # Information about an ability that we possess.
+  class Ability
+    ##
+    # Create a new ability.
+    #
+    # @param block    code to run
+    # @param timeout  server gives up on us after this many seconds
+    def initialize(block, timeout=nil)
+      @block = block
+      @timeout = timeout
+    end
+    attr_reader :timeout
 
-    def initialize(job_servers, opts = {})
-      @reactors = []
-      @abilities = {}
+    ##
+    # Run the block of code.
+    #
+    # @param data  data passed to us by a client
+    # @param job   interface to report job information to the server
+    def run(data, job)
+      @block.call(data, job)
+    end
+  end
 
-      @job_servers = Array[*job_servers]
+  # = Job
+  #
+  # == Description
+  # Interface to allow a worker to report information to a job server.
+  class Job
+    ##
+    # Create a new Job.
+    #
+    # @param sock    Socket connected to job server
+    # @param handle  job server-supplied job handle
+    def initialize(sock, handle)
+      @socket = sock
+      @handle = handle
+    end
 
-      @opts = opts
+    ##
+    # Report our status to the job server.
+    def report_status(numerator, denominator)
+      req = Util.pack_request(
+        :work_status, "#{@handle}\0#{numerator}\0#{denominator}")
+      Util.send_request(@socket, req)
+      self
     end
 
-    def add_ability(name, timeout = nil, &f)
-      remove_ability(name) if @abilities.has_key?(name)
-      @abilities[name] = { :callback => f, :timeout => timeout }
+    ##
+    # Send data before job completes
+    def send_data(data)
+      req = Util.pack_request(:work_data, "#{@handle}\0#{data}")
+      Util.send_request(@socket, req)
+      self
     end
 
-    def remove_ability(name)
-      @abilities.delete(name)
+    ##
+    # Send a warning explicitly
+    def report_warning(warning)
+      req = Util.pack_request(:work_warning, "#{@handle}\0#{warning}")
+      Util.send_request(@socket, req)
+      self
     end
+  end
 
-    def has_ability?(name)
-      @abilities.has_key?(name)
+  ##
+  # Create a new worker.
+  #
+  # @param job_servers  "host:port"; either a single server or an array
+  # @param opts         hash of additional options
+  def initialize(job_servers=nil, opts={})
+    chars = ('a'..'z').to_a
+    @client_id = Array.new(30) { chars[rand(chars.size)] }.join
+    @sockets = {}  # "host:port" -> Socket
+    @abilities = {}  # "funcname" -> Ability
+    @bad_servers = []  # "host:port"
+    @servers_mutex = Mutex.new
+    %w{client_id reconnect_sec
+       network_timeout_sec}.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 worker args: ' + opts.keys.sort.join(', ')
     end
+    @reconnect_sec = 30 if not @reconnect_sec
+    @network_timeout_sec = 5 if not @network_timeout_sec
+    @worker_enabled = true
+    @status = :preparing
+    self.job_servers = job_servers if job_servers
+    start_reconnect_thread
+  end
+  attr_accessor :client_id, :reconnect_sec, :network_timeout_sec, :bad_servers, :worker_enabled, :status
 
-    def work
-      EM.run do
-        @job_servers.each do |hostport|
-          host, port = hostport.split(":")
-          opts = { :abilities => @abilities }.merge(@opts)
-          Gearman::Evented::WorkerReactor.connect(host, port, opts)
+  # Start a thread to repeatedly attempt to connect to down job servers.
+  def start_reconnect_thread
+    Thread.new do
+      loop do
+        @servers_mutex.synchronize do
+          # If there are any failed servers, try to reconnect to them.
+          if not @bad_servers.empty?
+            update_job_servers(@sockets.keys + @bad_servers)
+          end
         end
+        sleep @reconnect_sec
       end
+    end.run
+  end
+
+  def job_servers
+    servers = nil
+    @servers_mutex.synchronize do
+      servers = @sockets.keys + @bad_servers
+    end
+    servers
+  end
+
+  ##
+  # Connect to job servers to be used by this worker.
+  #
+  # @param servers  "host:port"; either a single server or an array
+  def job_servers=(servers)
+    @servers_mutex.synchronize do
+      update_job_servers(servers)
+    end
+  end
+
+  # Internal function to actually connect to servers.
+  # Caller must acquire @servers_mutex before calling us.
+  #
+  # @param servers  "host:port"; either a single server or an array
+  def update_job_servers(servers)
+    @bad_servers = []
+    servers = Set.new(Util.normalize_job_servers(servers))
+    # Disconnect from servers that we no longer care about.
+    @sockets.each do |server,sock|
+      if not servers.include? server
+        Util.logger.debug "GearmanRuby: Disconnecting from old server #{server}"
+        sock.close
+        @sockets.delete(server)
+      end
+    end
+    # Connect to new servers.
+    servers.each do |server|
+      if not @sockets[server]
+        begin
+          Util.logger.debug "GearmanRuby: Connecting to server #{server}"
+          @sockets[server] = connect(server)
+        rescue NetworkError
+          @bad_servers << server
+          Util.logger.debug "GearmanRuby: Unable to connect to #{server}"
+        end
+      end
+    end
+  end
+  private :update_job_servers
+
+  ##
+  # Connect to a job server.
+  #
+  # @param hostport  "hostname:port"
+  def connect(hostport)
+    begin
+      # FIXME: handle timeouts
+      sock = TCPSocket.new(*hostport.split(':'))
+    rescue SocketError, Errno::ECONNREFUSED, Errno::ECONNRESET, Errno::EHOSTUNREACH
+      raise NetworkError
+    rescue Exception => e
+      Util.logger.debug "GearmanRuby: Unhandled exception while connecting to #{hostport} : #{e} (raising NetworkError exception)"
+      raise NetworkError
+    end
+    # FIXME: catch exceptions; do something smart
+    Util.send_request(sock, Util.pack_request(:set_client_id, @client_id))
+    @abilities.each {|f,a| announce_ability(sock, f, a.timeout) }
+    sock
+  end
+  private :connect
+
+  ##
+  # Announce an ability over a particular socket.
+  #
+  # @param sock     Socket connect to a job server
+  # @param func     function name (including prefix)
+  # @param timeout  the server will give up on us if we don't finish
+  #                 a task in this many seconds
+  def announce_ability(sock, func, timeout=nil)
+    begin
+      cmd = timeout ? :can_do_timeout : :can_do
+      arg = timeout ? "#{func}\0#{timeout.to_s}" : func
+      Util.send_request(sock, Util.pack_request(cmd, arg))
+    rescue Exception => ex
+      bad_servers << @sockets.keys.detect{|hp| @sockets[hp] == sock}
+    end
+  end
+  private :announce_ability
+
+  ##
+  # Add a new ability, announcing it to job servers.
+  #
+  # The passed-in block of code will be executed for jobs of this function
+  # type.  It'll receive two arguments, the data supplied by the client and
+  # a Job object.  If it returns nil or false, the server will be informed
+  # that the job has failed; otherwise the return value of the block will
+  # be passed back to the client in String form.
+  #
+  # @param func     function name (without prefix)
+  # @param timeout  the server will give up on us if we don't finish
+  #                 a task in this many seconds
+  def add_ability(func, timeout=nil, &f)
+    @abilities[func] = Ability.new(f, timeout)
+    @sockets.values.each {|s| announce_ability(s, func, timeout) }
+  end
+
+  ##
+  # Let job servers know that we're no longer able to do something.
+  #
+  # @param func  function name
+  def remove_ability(func)
+    @abilities.delete(func)
+    req = Util.pack_request(:cant_do, func)
+    @sockets.values.each {|s| Util.send_request(s, req) }
+  end
+
+  ##
+  # Handle a job_assign packet.
+  #
+  # @param data      data in the packet
+  # @param sock      Socket on which the packet arrived
+  # @param hostport  "host:port"
+  def handle_job_assign(data, sock, hostport)
+    handle, func, data = data.split("\0", 3)
+    if not func
+      Util.logger.error "GearmanRuby: Ignoring job_assign with no function from #{hostport}"
+      return false
+    end
+
+    Util.logger.error "GearmanRuby: Got job_assign with handle #{handle} and #{data.size} byte(s) " +
+      "from #{hostport}"
+
+    ability = @abilities[func]
+    if not ability
+      Util.logger.error "Ignoring job_assign for unsupported func #{func} " +
+        "with handle #{handle} from #{hostport}"
+      Util.send_request(sock, Util.pack_request(:work_fail, handle))
+      return false
     end
 
+    exception = nil
+    begin
+      ret = ability.run(data, Job.new(sock, handle))
+    rescue Exception => e
+      exception = e
+    end
+
+
+    cmd = if ret && exception.nil?
+      ret = ret.to_s
+      Util.logger.debug "GearmanRuby: Sending work_complete for #{handle} with #{ret.size} byte(s) " +
+        "to #{hostport}"
+      [ Util.pack_request(:work_complete, "#{handle}\0#{ret}") ]
+    elsif exception.nil?
+      Util.logger.debug "GearmanRuby: Sending work_fail for #{handle} to #{hostport}"
+      [ Util.pack_request(:work_fail, handle) ]
+    elsif exception
+      Util.logger.debug "GearmanRuby: Sending work_exception for #{handle} to #{hostport}"
+      [ Util.pack_request(:work_exception, "#{handle}\0#{exception.message}") ]
+    end
+
+    cmd.each {|p| Util.send_request(sock, p) }
+    true
   end
+
+  ##
+  # Do a single job and return.
+  def work
+    req = Util.pack_request(:grab_job)
+    loop do
+      @status = :preparing
+      bad_servers = []
+      # We iterate through the servers in sorted order to make testing
+      # easier.
+      servers = nil
+      @servers_mutex.synchronize { servers = @sockets.keys.sort }
+      servers.each do |hostport|
+        Util.logger.debug "GearmanRuby: Sending grab_job to #{hostport}"
+        sock = @sockets[hostport]
+        Util.send_request(sock, req)
+
+        # Now that we've sent grab_job, we need to keep reading packets
+        # until we see a no_job or job_assign response (there may be a noop
+        # waiting for us in response to a previous pre_sleep).
+        loop do
+          begin
+            type, data = Util.read_response(sock, @network_timeout_sec)
+            case type
+            when :no_job
+              Util.logger.debug "GearmanRuby: Got no_job from #{hostport}"
+              break
+            when :job_assign
+              @status = :working
+              return worker_enabled if handle_job_assign(data, sock, hostport)
+              break
+            else
+              Util.logger.debug "GearmanRuby: Got #{type.to_s} from #{hostport}"
+            end
+          rescue Exception
+            Util.logger.debug "GearmanRuby: Server #{hostport} timed out or lost connection (#{$!.inspect}); marking bad"
+            bad_servers << hostport
+            break
+          end
+        end
+      end
+
+      @servers_mutex.synchronize do
+        bad_servers.each do |hostport|
+          @sockets[hostport].close if @sockets[hostport]
+          @bad_servers << hostport if @sockets[hostport]
+          @sockets.delete(hostport)
+        end
+      end
+
+      Util.logger.debug "GearmanRuby: Sending pre_sleep and going to sleep for #{@reconnect_sec} sec"
+      @servers_mutex.synchronize do
+        @sockets.values.each do |sock|
+          Util.send_request(sock, Util.pack_request(:pre_sleep))
+        end
+      end
+
+      return false unless worker_enabled
+      @status = :waiting
+
+      # FIXME: We could optimize things the next time through the 'each' by
+      # sending the first grab_job to one of the servers that had a socket
+      # with data in it.  Not bothering with it for now.
+      IO::select(@sockets.values, nil, nil, @reconnect_sec)
+    end
+  end
+end
+
 end