gearman-ruby 2.0.0 → 3.0.1

data/lib/gearman/client.rb

@@ -1,80 +1,152 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+
 module Gearman
-  class Client
-    attr_accessor :uniq, :jobs
 
-    def initialize(job_servers, opts = {})
-      @reactors = []
-      @jobs = {}
+# = 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 options
+  #
+  # @options options to pass to the servers  "exeptions"
+  def option_request(opts)
+    Util.logger.debug "GearmanRuby: Send options request with #{opts}"
+    request = Util.pack_request("option_req", opts)
+    sock= self.get_socket(self.get_job_server)
+    Util.send_request(sock, request)
+    response = Util.read_response(sock, 20)
+    raise ProtocolError, response[1] if response[0]==:error
+  end
+
+  ##
+  # 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
 
-      @job_servers = Array[*job_servers]
+    raise Exception.new('No servers available') if @job_servers.empty?
 
-      @uniq = opts.delete(:uniq)
-      @opts = opts
+    @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
 
-    # Run a Task or Taskset
-    def run(taskset, timeout = nil, async = false)
-      timeout ||= 0
-      use_em_stop = EM.reactor_running?
-      EM.run do
-        @taskset ||= Taskset.new
-        @taskset += Taskset.create(taskset)
-
-        job_servers = @job_servers.dup
-        @reactors.each do |reactor|
-          unless reactor.connected?
-            reactor.reconnect true
-            reactor.keep_connected = async
-            reactor.callback { create_job(@taskset.shift, reactor) }
-            job_servers.delete reactor.to_s
-          else
-            create_job(@taskset.shift, reactor)
-          end
-        end
-        job_servers.each do |hostport|
-          host, port = hostport.split(":")
-          reactor = Gearman::Evented::ClientReactor.connect(host, port, @opts)
-          reactor.keep_connected = async
-          reactor.callback { create_job(@taskset.shift, reactor) }
-          @reactors << reactor
-        end
-
-        if timeout > 0
-          if use_em_stop
-            EM.add_timer(timeout) { EM.stop }
-          else
-            sleep timeout
-          end
-        elsif !async
-          Thread.new do
-            loop do
-              sleep 0.1
-              live = 0
-              @reactors.each {|reactor| live += 1 if reactor.connected? }
-              break if live == 0
-            end
-          end.join
-        end
+    num_retries.times do |i|
+      begin
+        sock = TCPSocket.new(*hostport.split(':'))
+      rescue Exception
+      else
+
+        @socket_to_hostport[sock] = hostport
+        return sock
       end
     end
 
-    private
+    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.logger.error "GearmanRuby: 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
 
-      def create_job(task, reactor = nil)
-        return unless task
-        reactor ||= @reactors[rand(@reactors.size)]
-        unless reactor.connected?
-          log "create_job: server #{reactor} not connected"
-          EM.next_tick { create_job(task) }
-          return
-        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)
 
-        reactor.submit_job(task) {|handle| create_job(@taskset.shift) }
-      end
+    result = nil
+    failed = false
+    task.on_complete {|v| result = v }
+    task.on_fail { failed = true }
 
-      def log(msg)
-        Gearman::Util.log(msg)
-      end
+    taskset = TaskSet.new(self)
+    taskset.add_task(task)
+    taskset.wait
 
+    failed ? nil : result
   end
+
+end
+
 end

data/lib/gearman/server.rb

@@ -45,14 +45,14 @@ class Server
   def send_command(name)
     response = ''
     socket.puts(name)
-    while true do
+    while true do 
       if buf = socket.recv_nonblock(65536) rescue nil
-        response << buf
+        response << buf 
         return response if response =~ /\n.\n$/
       end
     end
   end
-
+  
   ##
   # Returns results of a 'status' command.
   #
@@ -68,7 +68,7 @@ class Server
     end
     status
   end
-
+  
   ##
   # Returns results of a 'workers' command.
   #

data/lib/gearman/task.rb

@@ -1,20 +1,58 @@
+#!/usr/bin/env ruby
+
+require 'digest/sha1'
+
 module Gearman
+
+  # = Task
+  #
+  # == Description
+  # A task submitted to a Gearman job server.
   class Task
 
-    attr_reader :name, :payload, :retries_done
-    attr_accessor :retries, :priority, :background, :poll_status_interval
+    ##
+    # 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?
+      @uniq = nil       # Initialize to nil
+      %w{on_complete on_fail on_retry on_exception on_status on_warning on_data
+         uniq retry_count priority hash background}.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
+    
+    end
 
-    def initialize(name, payload = nil, opts = {})
-      @name       = name.to_s
-      @payload    = payload || ''
-      @priority   = opts.delete(:priority).to_sym rescue nil
-      @background = opts.delete(:background) ? true : false
+    attr_accessor :uniq, :retry_count, :priority, :background, :epoch
+    attr_reader :successful, :func, :arg, :retries_done
 
-      @retries_done = 0
-      @retries      = opts.delete(:retries) || 0
+    ##
+    # Schedule this job to run at a certain time (like `cron`)
+    # XXX: But there is no wildcard??
+    #
+    # @param time     Ruby Time object that represents when to run the thing
+    def schedule(time)
+      @scheduled_at = time
+    end
 
-      @poll_status_interval = opts.delete(:poll_status_interval)
-      @uniq = opts.has_key?(:uuid) ? opts.delete(:uuid) : `uuidgen`.strip
+    ##
+    # 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
 
     ##
@@ -73,27 +111,109 @@ module Gearman
       @on_data = 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 should_retry?
-      return false if @retries_done >= @retries
+    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
 
-    def background?
-      background
+    ##
+    # Record an exception.
+    #
+    def handle_exception(exception)
+      @on_exception.call(exception) if @on_exception
+      self
     end
 
-    def dispatch(event, *args)
-      callback = instance_variable_get("@#{event}".to_sym)
-      callback.call(*args) if callback
+    ##
+    # Handle a status update for the task.
+    def handle_status(numerator, denominator)
+      @on_status.call(numerator, denominator) if @on_status
+      self
+    end
+
+    ##
+    # Handle a warning.
+    #
+    def handle_warning(message)
+      @on_warning.call(message) if @on_warning
+      self
+    end
+
+    ##
+    # Handle (partial) data
+    def handle_data(data)
+      @on_data.call(data) if @on_data
+      self
+    end
+
+    ##
+    # Return a hash that we can use to execute identical tasks on the same
+    # job server.
+    #
+    # @return  hashed value, based on @func + @arg (sorted letters and then SHA1'd) if @uniq is nil, 
+    #          or the SHA1 of @uniq if it's set to something
+    #
+    def get_uniq_hash
+      return @hash if @hash
+    
+      if @uniq.nil? 
+        string = (@func+@arg).split(//).sort.join
+      else 
+        string = @uniq
+      end
+    
+      @hash = Digest::SHA1.hexdigest(string)
     end
 
-    def hash
-      @uniq
+    ##
+    # Construct a packet to submit this task to a job server.
+    #
+    # @return            String representation of packet
+    def get_submit_packet()
+      modes = ['submit_job']
+    
+      if @scheduled_at
+        modes << 'epoch'
+        args = [func, get_uniq_hash, @scheduled_at.to_i, arg]
+      else
+        if @priority
+          modes << 'high' if @priority == :high
+          modes << 'low'  if @priority == :low
+        else
+          modes << 'bg' if @background
+        end
+        args = [func, get_uniq_hash, arg]
+      end
+    
+      mode = modes.join('_')
+      Util::pack_request(mode, args.join("\0"))
+    end
+  end
+  
+  class BackgroundTask < Task
+    def initialize(*args)
+      super
+      @background = true
     end
   end
+
 end

data/lib/gearman/taskset.rb

@@ -1,11 +1,286 @@
+#!/usr/bin/env ruby
+
+require 'socket'
+require 'time'
+
 module Gearman
-  class Taskset < ::Array
 
-    def self.create(task_or_taskset)
-      [*task_or_taskset]
+# = 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.logger.debug "GearmanRuby: 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.logger.debug "GearmanRuby: Got timeout on read from #{hostport}"
+        @task_waiting_for_handle = nil
+        @client.close_socket(sock)
+        return false
+      end
     end
 
-    alias :add :<<
-    alias :add_task :add
+    @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.logger.debug "GearmanRuby: 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
+    if(task.background)
+      @finished_tasks << task
+    else
+      (@tasks_in_progress[js_handle] ||= []) << task
+    end
+    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.logger.debug "GearmanRuby: Got work_complete with handle #{handle} and #{data ? data.size : '0'} byte(s) of data from #{hostport}"
+    tasks_in_progress(hostport, handle, true).each do |t|
+      t.handle_completion(data)
+      @finished_tasks << t
+    end
+    nil
+  end
+  private :handle_work_complete
+
+  ##
+  # Handle a 'work_exception' response from a job server.
+  #
+  # @param hostport  "host:port" of job server
+  # @param data      data returned in packet from server
+  def handle_work_exception(hostport, data)
+    handle, exception = data.split("\0", 2)
+    Util.logger.debug "GearmanRuby: Got work_exception with handle #{handle} from #{hostport}: '#{exception}'"
+    tasks_in_progress(hostport, handle).each {|t| t.handle_exception(exception) }
+  end
+  private :handle_work_exception
+
+  ##
+  # 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.logger.debug "GearmanRuby: Got work_fail with handle #{data} from #{hostport}"
+    tasks_in_progress(hostport, data, true).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.logger.debug "GearmanRuby: Got work_status with handle #{handle} from #{hostport}: #{num}/#{den}"
+    tasks_in_progress(hostport, handle).each {|t| t.handle_status(num, den) }
+  end
+  private :handle_work_status
+
+  ##
+  # Handle a 'work_warning' response from a job server.
+  #
+  # @param hostport "host:port" of job server
+  # @param data     data returned in packet from server
+  def handle_work_warning(hostport, data)
+    handle, message = data.split("\0", 2)
+    Util.logger.debug "GearmanRuby: Got work_warning with handle #{handle} from #{hostport}: '#{message}'"
+    tasks_in_progress(hostport, handle).each {|t| t.handle_warning(message) }
+  end
+  private :handle_work_warning
+
+  ##
+  # Handle a 'work_data' response from a job server
+  #
+  # @param hostport  "host:port" of a job server
+  # @param data       data returned in packet from server
+  def handle_work_data(hostport, data)
+    handle, data = data.split("\0", 2)
+    Util.logger.debug "GearmanRuby: Got work_data 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[js_handle]
+    if not tasks
+      raise ProtocolError, "Got unexpected work_data with handle #{handle} from #{hostport} (no task by that name)"
+    end
+    tasks.each {|t| t.handle_data(data) }
+  end
+  private :handle_work_data
+
+  ##
+  # 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)
+    known_types = [ :job_created,
+                    :work_complete,
+                    :work_fail,
+                    :work_status,
+                    :work_exception,
+                    :work_warning,
+                    :work_data ]
+
+    if known_types.include?(type)
+      send("handle_#{type}".to_sym, hostport, data)
+    else
+      Util.logger.debug "GearmanRuby: 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 = 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
+
+      ready_socks = remaining == 0 ? nil : IO::select(@sockets.values, nil, nil, remaining)
+      if not ready_socks or not ready_socks[0]
+        Util.logger.debug "GearmanRuby: 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 ? end_time - Time.now.to_f : nil))
+        rescue ProtocolError
+          hostport = @client.get_hostport_for_socket(sock)
+          Util.logger.debug "GearmanRuby: Ignoring bad packet from #{hostport}"
+        rescue NetworkError
+          hostport = @client.get_hostport_for_socket(sock)
+          Util.logger.debug "GearmanRuby: Got timeout on read from #{hostport}"
+        end
+      end
+    end
+
+    @sockets.values.each {|s| @client.return_socket(s) }
+    @sockets = {}
+    @finished_tasks.each do |t|
+      if ( (t.background.nil? || t.background == false) && !t.successful)
+        Util.logger.debug "GearmanRuby: Taskset failed"
+        return false
+      end
+    end
+    true
+  end
+
+  private
+    def tasks_in_progress(hostport, handle, remove_task = false)
+      js_handle = Util.handle_to_str(hostport, handle)
+      tasks = remove_task ? @tasks_in_progress.delete(js_handle) : @tasks_in_progress[js_handle]
+      if not tasks
+        raise ProtocolError, "Got unexpected work_data with handle #{handle} from #{hostport} (no task by that name)"
+      end
+      tasks
+    end
+end
+
 end