xing-gearman-ruby 1.0.0

data/lib/gearman/util.rb

@@ -0,0 +1,212 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+require 'time'
+
+module Gearman
+
+  class ServerDownException < Exception; end
+
+# = 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
+    23 => :can_do_timeout,  # W->J: FUNC[0]TIMEOUT
+    2  => :cant_do,         # W->J: FUNC
+    3  => :reset_abilities, # W->J: --
+    22 => :set_client_id,   # W->J: [RANDOM_STRING_NO_WHITESPACE]
+    4  => :pre_sleep,       # W->J: --
+
+    6  => :noop,            # J->W: --
+    7  => :submit_job,      # C->J: FUNC[0]UNIQ[0]ARGS
+    21 => :submit_job_high, # C->J: FUNC[0]UNIQ[0]ARGS
+    18 => :submit_job_bg,   # 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
+    20 => :status_res,      # C->J: HANDLE[0]KNOWN[0]RUNNING[0]NUM[0]DENOM
+
+    16 => :echo_req,        # ?->J: TEXT
+    17 => :echo_res,        # J->?: TEXT
+
+    19 => :error,           # J->?: ERRCODE[0]ERR_TEXT
+  }
+
+  # Map e.g. 'can_do' => 1
+  NUMS = COMMANDS.invert
+
+  # Default job server port.
+  DEFAULT_PORT = 7003
+
+  @@debug = false
+
+  ##
+  # Enable or disable debugging output (off by default).
+  #
+  # @param v  print debugging output
+  def Util.debug=(v)
+    @@debug = v
+  end
+
+  ##
+  # 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
+
+  ##
+  # 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
+      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
+
+  ##
+  # 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
+
+  ##
+  # 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
+
+  ##
+  # 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
+
+  ##
+  # 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
+
+  ##
+  # 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
+  end
+
+  ##
+  # Log a message no matter what.
+  #
+  # @param str  message to log
+  def Util.err(str)
+    log(str, true)
+  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

@@ -0,0 +1,341 @@
+#!/usr/bin/env ruby
+
+require 'set'
+require 'socket'
+require 'thread'
+
+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
+    ##
+    # 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
+
+    ##
+    # 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
+  #
+  # == 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
+
+    ##
+    # 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
+  end
+
+  ##
+  # 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
+    self.job_servers = job_servers if job_servers
+    start_reconnect_thread
+  end
+  attr_accessor :client_id, :reconnect_sec, :network_timeout_sec, :bad_servers
+
+  # 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.log "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.log "Connecting to server #{server}"
+          @sockets[server] = connect(server)
+        rescue NetworkError, Errno::ECONNRESET
+          @bad_servers << server
+          Util.log "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 Errno::ECONNREFUSED
+      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) }
+    @sockets[hostport] = 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.err "Ignoring job_assign with no function from #{hostport}"
+      return false
+    end
+
+    Util.log "Got job_assign with handle #{handle} and #{data.size} byte(s) " +
+      "from #{hostport}"
+
+    ability = @abilities[func]
+    if not ability
+      Util.err "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
+
+    ret = ability.run(data, Job.new(sock, handle))
+
+    cmd = nil
+    if ret
+      ret = ret.to_s
+      Util.log "Sending work_complete for #{handle} with #{ret.size} byte(s) " +
+        "to #{hostport}"
+      cmd = Util.pack_request(:work_complete, "#{handle}\0#{ret}")
+    else
+      Util.log "Sending work_fail for #{handle} to #{hostport}"
+      cmd = Util.pack_request(:work_fail, handle)
+    end
+
+    Util.send_request(sock, cmd)
+    true
+  end
+
+  ##
+  # Do a single job and return.
+  def work
+    req = Util.pack_request(:grab_job)
+    loop do
+      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.log "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.log "Got no_job from #{hostport}"
+              break
+            when :job_assign
+              return if handle_job_assign(data, sock, hostport)
+              break
+            else
+              Util.log "Got #{type.to_s} from #{hostport}"
+            end
+          rescue Gearman::ServerDownException, NetworkError, Errno::ECONNRESET
+            Util.log "Server #{hostport} timed out or lost connection; 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.log "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
+
+      # 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

data/lib/gearman.rb

@@ -0,0 +1,76 @@
+#!/usr/bin/env ruby
+#
+# = Name
+# Gearman
+#
+# == Description
+# This file provides a Ruby interface for communicating with the Gearman
+# distributed job system.
+#
+# "Gearman is a system to farm out work to other machines, dispatching
+# function calls to machines that are better suited to do work, to do work
+# in parallel, to load balance lots of function calls, or to call functions
+# between languages."  -- http://www.danga.com/gearman/
+#
+# == Version
+# 0.0.1
+#
+# == Author
+# Daniel Erat <dan-ruby@erat.org>
+#
+# == License
+# This program is free software; you can redistribute it and/or modify it
+# under the terms of either:
+#
+# a) the GNU General Public License as published by the Free Software
+#    Foundation; either version 1, or (at your option) any later version,
+#    or
+#
+# b) the "Artistic License" which comes with Perl.
+
+# = Gearman
+#
+# == Usage
+#  require 'gearman'
+#
+#  # Create a new client and tell it about two job servers.
+#  c = Gearman::Client.new
+#  c.job_servers = ['127.0.0.1:7003', '127.0.0.1:7004']
+#
+#  # Create two tasks, using an "add" function to sum two numbers.
+#  t1 = Gearman::Task.new('add', '5 + 2')
+#  t2 = Gearman::Task.new('add', '1 + 3')
+#
+#  # Make the tasks print the data they get back from the server.
+#  t1.on_complete {|d| puts "t1 got #{d}" }
+#  t2.on_complete {|d| puts "t2 got #{d}" }
+#
+#  # Create a taskset, add the two tasks to it, and wait until they finish.
+#  ts = Gearman::TaskSet.new(c)
+#  ts.add_task(t1)
+#  ts.add_task(t2)
+#  ts.wait
+#
+# Or, a more simple example:
+#
+#  c = Gearman::Client.new('127.0.0.1')
+#  puts c.do_task('add', '2 + 2')
+#
+module Gearman
+
+require File.dirname(__FILE__) + '/gearman/client'
+require File.dirname(__FILE__) + '/gearman/task'
+require File.dirname(__FILE__) + '/gearman/taskset'
+require File.dirname(__FILE__) + '/gearman/util'
+require File.dirname(__FILE__) + '/gearman/worker'
+
+class InvalidArgsError < Exception
+end
+
+class ProtocolError < Exception
+end
+
+class NetworkError < Exception
+end
+
+end