gearman-ruby 3.0.8 → 4.0.2

data/lib/gearman/connection.rb

@@ -0,0 +1,158 @@
+require 'socket'
+
+module Gearman
+  class Connection
+    include Logging
+
+    def initialize(hostname, port)
+      @hostname = hostname
+      @port = port
+      @real_socket = nil
+    end
+
+    attr_reader :hostname, :port, :state, :socket
+
+    ##
+    # Check server health status by sending an ECHO request
+    # Return true / false
+    ##
+    def is_healthy?
+      if @real_socket == nil
+        logger.debug "Performing health check for #{self}"
+        begin
+          request = Packet.pack_request("echo_req", "ping")
+          response = send_request(request, 3)
+          logger.debug "Health check response for #{self} is #{response.inspect}"
+          raise ProtocolError unless response[0] == :echo_res and response[1] == "ping"
+          return true
+        rescue NetworkError
+          logger.debug "NetworkError -- unhealthy"
+          return false
+        rescue ProtocolError
+          logger.debug "ProtocolError -- unhealthy"
+          return false
+        end
+      end
+    end
+
+    ##
+    # @param num_retries  Number of times to retry
+    # @return             This connection's Socket
+    ##
+    def socket(num_retries=3)
+      # If we already have an open socket to this host, return it.
+      return @real_socket if @real_socket
+      num_retries.times do |i|
+        begin
+          logger.debug("Attempt ##{i} to connect to #{hostname}:#{port}")
+          @real_socket = TCPSocket.new(hostname, port)
+        rescue Exception => e
+          logger.error("Unable to connect: #{e}")
+          # Swallow error so we can retry -> num_retries times
+        else
+          return @real_socket
+        end
+      end
+
+      raise_exception("Unable to connect to job server #{hostname}:#{port}")
+    end
+
+    def close_socket
+      @real_socket.close if @real_socket
+      @real_socket = nil
+      true
+    end
+
+    def raise_exception(message)
+      close_socket
+      raise NetworkError, message
+    end
+
+    ##
+    # 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 timed_recv(sock, len, timeout=nil)
+      data = ''
+      start_time = Time.now.to_f
+      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
+        begin
+          data += sock.readpartial(len - data.size)
+        rescue
+          close_socket
+          raise NetworkError, "Unable to read data from socket."
+        end
+      end
+      if data.size < len
+        now = Time.now.to_f
+        if now > end_time
+          time_lapse = now - start_time
+          raise SocketTimeoutError, "Took too long to read data: #{time_lapse} sec. to read on a  #{timeout} sec. timeout"
+        else
+          raise_exception("Read #{data.size} byte(s) instead of #{len}")
+        end
+      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 read_response(timeout=nil)
+      end_time = Time.now.to_f + timeout if timeout
+      head = timed_recv(socket, 12, timeout)
+      magic, type, len = head.unpack('a4NN')
+      raise ProtocolError, "Invalid magic '#{magic}'" unless magic == "\0RES"
+      buf = len > 0 ?
+          timed_recv(socket, len, timeout ? end_time - Time.now.to_f : nil) : ''
+      type = Packet::COMMANDS[type]
+      raise ProtocolError, "Invalid packet type #{type}" unless type
+      [type, buf]
+    end
+
+    ##
+    # Send a request packet over a socket that needs a response.
+    #
+    # @param sock  Socket connected to a job server
+    # @param req   request packet to send
+    # @result response from server
+    def send_request(req, timeout = nil)
+      send_update(req, timeout)
+      return read_response(timeout)
+    end
+
+    def send_update(req, timeout = nil)
+      len = with_safe_socket_op{ socket.write(req) }
+      if len != req.size
+        raise_exception("Wrote #{len} instead of #{req.size}")
+      end
+    end
+
+    def with_safe_socket_op
+      begin
+        yield
+      rescue Exception => ex
+        raise_exception(ex.message)
+      end
+    end
+
+    def to_host_port
+      "#{hostname}:#{port}"
+    end
+
+    def to_s
+      "#{hostname}:#{port} (connected: #{@real_socket != nil})"
+    end
+
+  end
+end

data/lib/gearman/connection_pool.rb

@@ -0,0 +1,131 @@
+require 'thread'
+
+module Gearman
+  class ConnectionPool
+    include Logging
+
+    DEFAULT_PORT = 4730
+
+    def initialize(servers = [])
+      @bad_servers          = []
+      @coalesce_connections = {}
+      @connection_handler   = nil
+      @job_servers          = []
+      @reconnect_seconds    = 10
+      @server_counter       = 0   # Round-robin distribution of requests
+      @servers_mutex        = Mutex.new
+
+      add_servers(servers)
+      start_reconnect_thread
+    end
+
+    def add_connection(connection)
+      @servers_mutex.synchronize do
+        if connection.is_healthy?
+          activate_connection(connection)
+        else
+          deactivate_connection(connection)
+        end
+      end
+    end
+
+    def add_host_port(host_port)
+      host, port = host_port.split(":")
+      connection = Connection.new(host, port.to_i)
+      add_connection(connection)
+    end
+
+    def add_servers(servers)
+      if servers.class == String or servers.class == Symbol
+        servers = [ servers.to_s ]
+      end
+
+      servers = servers.map {|s| s =~ /:/ ? s : "#{s}:#{DEFAULT_PORT}" }
+
+      servers.each do |host_port|
+        add_host_port(host_port)
+      end
+    end
+
+    def get_connection(coalesce_key = nil)
+      @servers_mutex.synchronize do
+
+        logger.debug "Available job servers: #{@job_servers.inspect}"
+        raise NoJobServersError if @job_servers.empty?
+        @server_counter += 1
+        @job_servers[@server_counter % @job_servers.size]
+      end
+    end
+
+    def on_connection(&block)
+      @connection_handler = block
+    end
+
+    def poll_connections(timeout = nil)
+      available_sockets = []
+      @servers_mutex.synchronize do
+        available_sockets.concat @job_servers.collect { |conn| conn.socket }
+      end
+      if available_sockets.size > 0
+      	logger.debug "Polling on #{available_sockets.size} available server(s) with a #{timeout} second timeout"
+      	IO::select(available_sockets, nil, nil, timeout)
+      end
+    end
+
+    def with_all_connections(&block)
+      @servers_mutex.synchronize do
+        @job_servers.each do |connection|
+          begin
+            block.call(connection)
+          rescue NetworkError => ex
+            logger.debug "Error with #{connection}, marking as bad"
+            deactivate_connection(connection)
+          end
+        end
+      end
+    end
+
+
+    private
+
+      def deactivate_connection(connection)
+        @job_servers.reject! { |c| c == connection }
+        @bad_servers << connection
+      end
+
+      def activate_connection(connection)
+        @bad_servers.reject! { |c| c == connection }
+        @job_servers << connection
+        @connection_handler.call(connection) if @connection_handler
+      end
+
+      def start_reconnect_thread
+        Thread.new do
+          loop do
+            @servers_mutex.synchronize do
+              # If there are any failed servers, try to reconnect to them.
+              update_job_servers unless @bad_servers.empty?
+            end
+            sleep @reconnect_seconds
+          end
+        end.run
+      end
+
+      def update_job_servers
+        logger.debug "Found #{@bad_servers.size} zombie connections, checking pulse."
+        @bad_servers.each do |connection|
+          begin
+            message = "Testing server #{connection}..."
+            if connection.is_healthy?
+              logger.debug "#{message} Connection is healthy, putting back into service"
+              activate_connection(connection)
+            else
+              logger.debug "#{message} Still down."
+            end
+          end
+        end
+      end
+
+
+  end
+end

data/lib/gearman/exceptions.rb

@@ -0,0 +1,24 @@
+module Gearman
+
+  class InvalidArgsError < Exception
+  end
+
+  class ProtocolError < Exception
+  end
+
+  class NetworkError < Exception
+  end
+
+  class NoJobServersError < Exception
+  end
+
+  class JobQueueError < Exception
+  end
+
+  class SocketTimeoutError < Exception
+  end
+
+  class ServerDownException < Exception
+  end
+
+end

data/lib/gearman/logging.rb

@@ -0,0 +1,19 @@
+module Gearman
+  module Logging
+
+    def self.included(target)
+      target.extend ClassMethods
+    end
+
+    def logger
+      self.class.logger
+    end
+
+    module ClassMethods
+      def logger
+        Gearman.logger
+      end
+    end
+
+  end
+end

data/lib/gearman/packet.rb

@@ -0,0 +1,61 @@
+module Gearman
+  class Packet
+    # 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
+
+    ##
+    # 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 Packet.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
+
+  end
+end

data/lib/gearman/task.rb

@@ -213,7 +213,7 @@ module Gearman
       end
 
       mode = modes.join('_')
-      Util::pack_request(mode, args.join("\0"))
+      Packet::pack_request(mode, args.join("\0"))
     end
   end
 

data/lib/gearman/task_set.rb

@@ -0,0 +1,67 @@
+require 'time'
+
+module Gearman
+  class TaskSet
+    include Logging
+
+    def initialize(client)
+      @client = client
+      @tasks_in_progress = []
+      @finished_tasks = []
+    end
+
+    ##
+    # Add a new task to this TaskSet.
+    #
+    # @param args  A Task object
+    # @return      true if the task was created successfully, false otherwise
+    def add_task(task)
+      @tasks_in_progress << task
+    end
+
+    ##
+    # Wait for all tasks in the set to finish.
+    #
+    # @param timeout  maximum amount of time to wait, in seconds - if this is nil, waits forever
+    def wait(timeout = 1)
+      end_time = if timeout
+        Time.now.to_f + timeout
+      else
+        nil
+      end
+
+      while not @tasks_in_progress.empty?
+        remaining = if end_time
+          (t = end_time - Time.now.to_f) > 0 ? t : 0
+        else
+          nil
+        end
+        begin
+          task = @tasks_in_progress.pop
+          if
+            @client.submit_job(task, true, remaining)
+            @finished_tasks << task
+          end
+        rescue SocketTimeoutError
+          return false
+        end
+
+      end
+
+      @finished_tasks.each do |t|
+        if ( (t.background.nil? || t.background == false) && !t.successful)
+          logger.warn "GearmanRuby: TaskSet failed"
+          return false
+        end
+      end
+      true
+    end
+
+    # Wait for all tasks in set to finish, with no timeout
+    def wait_forever
+	wait(nil)
+    end
+
+  end
+
+end