gearman-ruby 3.0.8 → 4.0.2

data/lib/gearman/version.rb

@@ -1,3 +1,3 @@
 module Gearman
-  VERSION = "3.0.8"
+  VERSION = "4.0.2"
 end

data/lib/gearman/worker.rb

@@ -1,475 +1,248 @@
 #!/usr/bin/env ruby
 
 require 'set'
-require 'socket'
-require 'thread'
+
+require 'gearman/worker/callbacks'
+require 'gearman/worker/ability'
+require 'gearman/worker/job'
 
 module Gearman
 
-# = 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
+  class Worker
+    include Logging
+    include Callbacks
+
     ##
-    # Create a new ability.
+    # Create a new worker.
     #
-    # @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
+    # @param job_servers  "host:port"; either a single server or an array
+    # @param opts         hash of additional options
+    def initialize(job_servers=nil, opts={})
+      @abilities            = {}
+      @client_id            = opts[:client_id] || generate_id
+      @connection_pool      = ConnectionPool.new(job_servers)
+      @network_timeout_sec  = opts[:network_timeout_sec] || 10
+      @reconnect_sec        = opts[:reconnect_sec] || 30
+      @status               = :preparing
+      @worker_enabled       = true
+
+      # Add callback for when connections occur -- register abilities and send client id
+      @connection_pool.on_connection do |connection|
+        connection.send_update(Packet.pack_request(:set_client_id, @client_id))
+        @abilities.each do |func_name, ability|
+          announce_ability(func_name, ability.timeout, connection)
+        end
+      end
     end
-    attr_reader :timeout
+
+    attr_accessor :client_id, :reconnect_sec, :network_timeout_sec, :worker_enabled, :status
 
     ##
-    # 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)
+    # @return A random string of 30 characters from a-z
+    def generate_id
+      chars = ('a'..'z').to_a
+      Array.new(30) { chars[rand(chars.size)] }.join
     end
-  
-  end
 
-  # = 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
-    attr_reader :handle
-    def initialize(sock, handle)
-      @socket = sock
-      @handle = handle
+    # Generate CAN_DO (or CAN_DO_TIMEOUT) packet and submit it
+    def announce_ability(func_name, timeout, connection)
+      cmd = timeout ? :can_do_timeout : :can_do
+      arg = timeout ? "#{func_name}\0#{timeout.to_s}" : func_name
+      connection.send_update(Packet.pack_request(cmd, arg))
+      logger.debug "Announced ability #{func_name}"
     end
 
     ##
-    # 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
+    # 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_name function name (without prefix)
+    # @param timeout   the server will give up on us if we don't finish
+    #                  a task in this many seconds
+    # @param block     Block to associate with the function
+    def add_ability(func_name, timeout=nil, &block)
+      @abilities[func_name] = Ability.new(func_name, block, timeout)
+      @connection_pool.with_all_connections do |connection|
+        announce_ability(func_name, timeout, connection)
+      end
     end
 
     ##
-    # Send data before job completes
-    def send_data(data)
-      req = Util.pack_request(:work_data, "#{@handle}\0#{data}")
-      Util.send_request(@socket, req)
-      self
+    # Callback for after an ability runs
+    def after_ability(func, &block)
+      abilities[func].after_complete(block)
     end
 
     ##
-    # Send a warning explicitly
-    def report_warning(warning)
-      req = Util.pack_request(:work_warning, "#{@handle}\0#{warning}")
-      Util.send_request(@socket, req)
-      self
+    # Let job servers know that we're no longer able to do something via CANT_DO
+    #
+    # @param func  function name
+    def remove_ability(func)
+      @abilities.delete(func)
+      req = Packet.pack_request(:cant_do, func)
+      @connection_pool.with_all_connections do  |connection|
+        connection.send_update(req)
+      end
     end
-  end
-
-  module Callbacks
 
-    %w(connect grab_job no_job job_assign work_complete work_fail
-       work_exception).each do |event|
+    ##
+    # Handle a job_assign packet.
+    #
+    # @param data       data in the packet
+    # @param connection Connection where the data originated
+    def handle_job_assign(data, connection)
+      handle, func, data = data.split("\0", 3)
+
+      if not func
+        logger.error "Ignoring JOB_ASSIGN with no function from #{connection}"
+        return false
+      end
 
-      define_method("on_#{event}") do |&callback|
-        instance_variable_set("@__on_#{event}", callback)
+      if not handle
+        logger.error "Ignoring JOB_ASSIGN with no job handle from #{connection}"
+        return false
       end
 
-      define_method("run_#{event}_callback") do
-        callback = instance_variable_get("@__on_#{event}")
-        return unless callback
+      logger.info "Got JOB_ASSIGN with handle #{handle} and #{data.size} byte(s) from #{connection}"
 
-        begin
-          callback.call
-        rescue Exception => e
-          Util.logger.error "GearmanRuby: #{event} failed: #{e.inspect}"
-        end
+      ability = @abilities[func]
+
+      if ability == nil
+        logger.error "Ignoring JOB_ASSIGN for unsupported function #{func} with handle #{handle} from #{connection}"
+        connection.send_update(Packet.pack_request(:work_fail, handle))
+        return false
       end
-    end
 
-  end
+      exception = nil
+      begin
+        ret = ability.run(data, Job.new(connection, handle))
+      rescue Exception => e
+        exception = e
+        logger.debug "Exception: #{e}\n#{e.backtrace.join("\n")}\n"
+      end
 
-  # Provides callbacks for internal worker use:
-  #
-  # def named_metric(metric)
-  #   "HardWorker.#{Process.pid}.#{metric}"
-  # end
-  #
-  # worker = Gearman::Worker.new
-  # worker.on_grab_job { StatsD.increment(named_metric('grab_job')) }
-  # worker.on_job_assign { StatsD.increment(named_metric('job_assign')) }
-  # worker.on_no_job { StatsD.increment(named_metric('no_job')) }
-  # worker.on_work_complete { StatsD.increment(named_metric('work_complete')) }
-  include Callbacks
-
-  ##
-  # 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
-    @after_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
+      packets = if ret && exception.nil?
+              logger.debug "Sending WORK_COMPLETE for #{handle} with #{ret.to_s.size} byte(s) to #{connection}"
+              run_work_complete_callback
+              [Packet.pack_request(:work_complete, "#{handle}\0#{ret.to_s}")]
+            elsif exception.nil?
+              logger.debug "Sending WORK_FAIL for #{handle} to #{connection}"
+              run_work_fail_callback
+              [Packet.pack_request(:work_fail, handle)]
+            elsif exception
+              logger.debug "Sending WORK_EXCEPTION for #{handle} to #{connection}"
+              run_work_exception_callback
+              [Packet.pack_request(:work_exception, "#{handle}\0#{exception.message}")]
+            end
 
-  # 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
+      packets.each do |packet|
+        connection.send_update(packet)
       end
-    end.run
-  end
 
-  def job_servers
-    servers = nil
-    @servers_mutex.synchronize do
-      servers = @sockets.keys + @bad_servers
+      true
     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.info "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.info "GearmanRuby: Connecting to server #{server}"
-          run_connect_callback
-          @sockets[server] = connect(server)
-        rescue NetworkError
-          @bad_servers << server
-          Util.logger.info "GearmanRuby: Unable to connect to #{server}"
-        end
+    ##
+    # Handle a message for the worker
+    #
+    # @param type       Packet type (NO_JOB, JOB_ASSIGN, NO_OP)
+    # @param data       Opaque data being passed with the message
+    # @param connection The Connection object where the message originates
+    # @return
+    def handle_work_message(type, data, connection)
+      case type
+        when :no_job
+          logger.info "Got NO_JOB from #{connection}"
+          run_no_job_callback
+        when :job_assign
+          @status = :working
+          run_job_assign_callback
+          return worker_enabled if handle_job_assign(data, connection)
+        when :noop
+          # We'll have to read again
+          logger.debug "Received NOOP while polling. Ignoring NOOP"
+        else
+          logger.error "Got unexpected #{type.to_s} from #{connection}"
       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
 
+    ##
+    # Do a single job and return.
+    def work
+      grab_job_req = Packet.pack_request(:grab_job)
+      type, data = nil
 
-  ##
-  # Add an after-ability hook
-  #
-  # The passed-in block of code will be executed after the work block for 
-  # jobs with the same function name. It takes two arguments, the result of
-  # the work and the original job data. This way, if you need to hook into 
-  # *after* the job_complete packet is sent to the server, you can do so. 
-  #
-  # N.B The after-ability hook ONLY runs if the ability was successful and no 
-  # exceptions were raised.
-  #
-  # @param func     function name (without prefix)
-  #
-  def after_ability(func, &f)
-    @after_abilities[func] = Ability.new(f)
-  end
+      loop do
+        @status = :preparing
+  
+        @connection_pool.with_all_connections do |connection|
+          begin
+            logger.debug "Sending GRAB_JOB to #{connection}"
+            run_grab_job_callback
+            type, data = connection.send_request(grab_job_req, @network_timeout_sec)
+            handle_work_message(type, data, connection)
+          end while type == :job_assign
+        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
+        logger.info "Sending PRE_SLEEP and going to sleep for #{@reconnect_sec} second(s)"
+        @connection_pool.with_all_connections do |connection|
+            connection.send_update(Packet.pack_request(:pre_sleep))
+        end
 
-    exception = nil
-    begin
-      ret = ability.run(data, Job.new(sock, handle))
-    rescue Exception => e
-      exception = e
-      Util.logger.debug "GearmanRuby: Exception: #{e}\n#{e.backtrace.join("\n")}\n"
-    end
+        return false unless worker_enabled
+        @status = :waiting
 
-    cmd = if ret && exception.nil?
-      Util.logger.debug "GearmanRuby: Sending work_complete for #{handle} with #{ret.to_s.size} byte(s) " +
-        "to #{hostport}"
-      run_work_complete_callback
-      [ Util.pack_request(:work_complete, "#{handle}\0#{ret.to_s}") ]
-    elsif exception.nil?
-      Util.logger.debug "GearmanRuby: Sending work_fail for #{handle} to #{hostport}"
-      run_work_fail_callback
-      [ Util.pack_request(:work_fail, handle) ]
-    elsif exception
-      Util.logger.debug "GearmanRuby: Sending work_exception for #{handle} to #{hostport}"
-      run_work_exception_callback
-      [ Util.pack_request(:work_exception, "#{handle}\0#{exception.message}") ]
-    end
+        time_asleep = Time.now
 
-    cmd.each {|p| Util.send_request(sock, p) }
-    
-    # There are cases where we might want to run something after the worker
-    # successfully completes the ability in question and sends its results
-    if ret && exception.nil?
-      after_ability = @after_abilities[func]
-      if after_ability
-        Util.logger.debug "Running after ability for #{func}..."
-        begin
-          after_ability.run(ret, data)
-        rescue Exception => e
-          Util.logger.debug "GearmanRuby: Exception: #{e}\n#{e.backtrace.join("\n")}\n"
-          nil
+        while (@status == :waiting)
+          sleep(time_asleep)
         end
-      end
-    end 
-    
-    true
-  end
 
-  ##
-  # Do a single job and return.
-  def work
-    req = Util.pack_request(:grab_job)
-    type = nil
-    data = nil
-    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}"
-        run_grab_job_callback
-        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}"
-              run_no_job_callback
-              break
-            when :job_assign
-              @status = :working
-              run_job_assign_callback
-              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.info "GearmanRuby: Server #{hostport} timed out or lost connection (#{$!.inspect}); marking bad"
-            bad_servers << hostport
-            break
-          end
-        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
+    ##
+    # Sleep and poll until timeout occurs or a NO_OP packet is received
+    # @param time_fell_asleep The time that we fell asleep (Time object)
+    def sleep(time_fell_asleep)
+      max_timeout = 30 - (Time.now - time_fell_asleep).to_i
+
+      if max_timeout > 0
+        # Use IO::select to wait for available connection data
+        @connection_pool.poll_connections(max_timeout)
       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
+      # If 30 seconds have passed, then wakeup
+      time_asleep = (Time.now - time_fell_asleep).to_f
+      @status = :wakeup if time_asleep >= 30
+  
+      # We didn't sleep for >= 30s, so we need to check for a NOOP
+      if (@status == :waiting)
+        @connection_pool.with_all_connections do |connection|
+          begin
+            type, data = connection.read_response(@network_timeout_sec)
 
-      return false unless worker_enabled
-      @status = :waiting
-
-      sleepTime = Time.now
-      while(@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)
-        
-        # If 30 seconds have passed, then wakeup
-        @status = :wakeup if Time.now - sleepTime > 30 
-
-        if(@status == :waiting)
-          @sockets.values.each do |sock|
-            type, data = Util.read_response(sock, @network_timeout_sec)
- 
-            # there shouldn't be anything else here, if there is, we should be able to ignore it...
-            if(type == :noop)
-              Util.logger.debug "Received NoOp while sleeping... waking up!"
+            # Wake up if we receive a NOOP packet
+            if (type == :noop)
+              logger.debug "Received NOOP while sleeping... waking up!"
               @status = :wakeup
+            else
+              logger.warn "Received something other than a NOOP packet while sleeping: #{type.to_s}"
             end
+          rescue SocketTimeoutError
+            # This is okay here.
           end
         end
       end
-      
     end
   end
-end
 
 end