faktory_worker_ruby 0.7.0 → 1.0.1

data/lib/faktory/cli.rb

@@ -8,6 +8,11 @@ require 'optparse'
 require 'erb'
 require 'fileutils'
 
+module Faktory
+  class CLI
+  end
+end
+
 require 'faktory'
 require 'faktory/util'
 
@@ -235,6 +240,10 @@ module Faktory
           opts[:tag] = arg
         end
 
+        o.on '-l', '--label LABEL', "Process label to use in Faktory UI" do |arg|
+          (opts[:labels] ||= []) << arg
+        end
+
         o.on "-q", "--queue QUEUE[,WEIGHT]", "Queues to process with optional weights" do |arg|
           queue, weight = arg.split(",")
           parse_queue opts, queue, weight

data/lib/faktory/client.rb

@@ -1,15 +1,29 @@
 require 'socket'
 require 'json'
 require 'uri'
+require 'digest'
 require 'securerandom'
+require 'timeout'
+require 'faktory/io'
 
 module Faktory
-  class CommandError < StandardError;end
-  class ParseError < StandardError;end
-
+  class BaseError < StandardError; end
+  class CommandError < BaseError; end
+  class ParseError < BaseError; end
+
+  # Faktory::Client provides a low-level connection to a Faktory server
+  # and APIs which map to Faktory commands.
+  #
+  # Most APIs will return `true` if the operation succeeded or raise a
+  # Faktory::BaseError if there was an unexpected error.
   class Client
+    # provides gets() and read() that respect a read timeout
+    include Faktory::ReadTimeout
+
     @@random_process_wid = ""
 
+    DEFAULT_TIMEOUT = 5.0
+
     HASHER = proc do |iter, pwd, salt|
       sha = Digest::SHA256.new
       hashing = pwd + salt
@@ -35,10 +49,13 @@ module Faktory
     # MY_FAKTORY_URL=tcp://:somepass@my-server.example.com:7419
     #
     # Note above, the URL can contain the password for secure installations.
-    def initialize(url: uri_from_env || 'tcp://localhost:7419', debug: false)
+    def initialize(url: uri_from_env || 'tcp://localhost:7419', debug: false, timeout: DEFAULT_TIMEOUT)
+      super
       @debug = debug
       @location = URI(url)
-      open
+      @timeout = timeout
+
+      open(@timeout)
     end
 
     def close
@@ -52,23 +69,95 @@ module Faktory
     def flush
       transaction do
         command "FLUSH"
-        ok!
+        ok
+      end
+    end
+
+    def create_batch(batch, &block)
+      bid = transaction do
+        command "BATCH NEW", Faktory.dump_json(batch.to_h)
+        result!
+      end
+      batch.instance_variable_set(:@bid, bid)
+
+      old = Thread.current["faktory_batch"]
+      Thread.current["faktory_batch"] = batch
+      begin
+        # any jobs pushed in this block will implicitly have
+        # their `bid` attribute set so they are associated
+        # with the current batch.
+        yield batch
+      ensure
+        Thread.current[:faktory_batch] = old
+      end
+      transaction do
+        command "BATCH COMMIT", bid
+        ok
       end
+      bid
     end
 
+    def batch_status(bid)
+      transaction do
+        command "BATCH STATUS", bid
+        Faktory.load_json result!
+      end
+    end
+
+    def reopen_batch(b)
+      transaction do
+        command "BATCH OPEN", b.bid
+        ok
+      end
+      old = Thread.current[:faktory_batch]
+      Thread.current[:faktory_batch] = b
+      begin
+        # any jobs pushed in this block will implicitly have
+        # their `bid` attribute set so they are associated
+        # with the current batch.
+        yield b
+      ensure
+        Thread.current[:faktory_batch] = old
+      end
+      transaction do
+        command "BATCH COMMIT", b.bid
+        ok
+      end
+    end
+
+    def get_track(jid)
+      transaction do
+        command "TRACK GET", jid
+        hashstr = result!
+        JSON.parse(hashstr)
+      end
+    end
+
+    # hash must include a 'jid' element
+    def set_track(hash)
+      transaction do
+        command("TRACK SET", Faktory.dump_json(hash))
+        ok
+      end
+    end
+
+    # Push a hash corresponding to a job payload to Faktory.
+    # Hash must contain "jid", "jobtype" and "args" elements at minimum.
+    # Returned value will either be the JID String if successful OR
+    # a symbol corresponding to an error.
     def push(job)
       transaction do
-        command "PUSH", JSON.generate(job)
-        ok!
-        job["jid"]
+        command "PUSH", Faktory.dump_json(job)
+        ok(job["jid"])
       end
     end
 
+    # Returns either a job hash or falsy.
     def fetch(*queues)
       job = nil
       transaction do
         command("FETCH", *queues)
-        job = result
+        job = result!
       end
       JSON.parse(job) if job
     end
@@ -76,34 +165,42 @@ module Faktory
     def ack(jid)
       transaction do
         command("ACK", %Q[{"jid":"#{jid}"}])
-        ok!
+        ok
       end
     end
 
     def fail(jid, ex)
       transaction do
-        command("FAIL", JSON.dump({ message: ex.message[0...1000],
+        command("FAIL", Faktory.dump_json({ message: ex.message[0...1000],
                           errtype: ex.class.name,
                           jid: jid,
                           backtrace: ex.backtrace}))
-        ok!
+        ok
       end
     end
 
     # Sends a heartbeat to the server, in order to prove this
     # worker process is still alive.
     #
+    # You can pass in the current_state of the process, for example during shutdown
+    # quiet and/or terminate can be supplied.
+    #
     # Return a string signal to process, legal values are "quiet" or "terminate".
     # The quiet signal is informative: the server won't allow this process to FETCH
     # any more jobs anyways.
-    def beat
+    def beat(current_state = nil)
       transaction do
-        command("BEAT", %Q[{"wid":"#{@@random_process_wid}"}])
-        str = result
+        if current_state.nil?
+          command("BEAT", %Q[{"wid":"#{@@random_process_wid}"}])
+        else
+          command("BEAT", %Q[{"wid":"#{@@random_process_wid}", "current_state":"#{current_state}"}])
+        end
+
+        str = result!
         if str == "OK"
           str
         else
-          hash = JSON.parse(str)
+          hash = Faktory.load_json(str)
           hash["state"]
         end
       end
@@ -112,8 +209,8 @@ module Faktory
     def info
       transaction do
         command("INFO")
-        str = result
-        JSON.parse(str) if str
+        str = result!
+        Faktory.load_json(str) if str
       end
     end
 
@@ -128,12 +225,14 @@ module Faktory
       @location.scheme =~ /tls/
     end
 
-    def open
+    def open(timeout = DEFAULT_TIMEOUT)
       if tls?
         sock = TCPSocket.new(@location.hostname, @location.port)
+        sock.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, true)
+
         ctx = OpenSSL::SSL::SSLContext.new
         ctx.set_params(verify_mode: OpenSSL::SSL::VERIFY_PEER)
-        ctx.ssl_version = :TLSv1_2
+        ctx.min_version = OpenSSL::SSL::TLS1_2_VERSION
 
         @sock = OpenSSL::SSL::SSLSocket.new(sock, ctx).tap do |socket|
           socket.sync_close = true
@@ -148,7 +247,7 @@ module Faktory
         "wid": @@random_process_wid,
         "hostname": Socket.gethostname,
         "pid": $$,
-        "labels": ["ruby-#{RUBY_VERSION}"],
+        "labels": Faktory.options[:labels] || ["ruby-#{RUBY_VERSION}"],
         "v": 2,
       }
 
@@ -175,11 +274,10 @@ module Faktory
         end
       end
 
-      command("HELLO", JSON.dump(payload))
-      ok!
+      command("HELLO", Faktory.dump_json(payload))
+      ok
     end
 
-
     def command(*args)
       cmd = args.join(" ")
       @sock.puts(cmd)
@@ -188,12 +286,22 @@ module Faktory
 
     def transaction
       retryable = true
+
+      # When using Faktory::Testing, you can get a client which does not actually
+      # have an underlying socket.  Now if you disable testing and try to use that
+      # client, it will crash without a socket.  This open() handles that case to
+      # transparently open a socket.
+      open(@timeout) if !@sock
+
       begin
         yield
-      rescue Errno::EPIPE, Errno::ECONNRESET
+      rescue SystemCallError, SocketError, TimeoutError
         if retryable
           retryable = false
-          open
+
+          @sock.close rescue nil
+          @sock = nil
+          open(@timeout)
           retry
         else
           raise
@@ -204,7 +312,7 @@ module Faktory
     # I love pragmatic, simple protocols.  Thanks antirez!
     # https://redis.io/topics/protocol
     def result
-      line = @sock.gets
+      line = gets
       debug "< #{line}" if @debug
       raise Errno::ECONNRESET, "No response" unless line
       chr = line[0]
@@ -213,11 +321,20 @@ module Faktory
       elsif chr == '$'
         count = line[1..-1].strip.to_i
         return nil if count == -1
-        data = @sock.read(count) if count > 0
-        line = @sock.gets # read extra linefeeds
+        data = read(count) if count > 0
+        line = gets # read extra linefeeds
         data
       elsif chr == '-'
-        raise CommandError, line[1..-1]
+        # Server can respond with:
+        #
+        # -ERR Something unexpected
+        # We raise a CommandError
+        #
+        # -NOTUNIQUE Job not unique
+        # We return ["NOTUNIQUE", "Job not unique"]
+        err = line[1..-1].split(" ", 2)
+        raise CommandError, err[1] if err[0] == "ERR"
+        err
       else
         # this is bad, indicates we need to reset the socket
         # and start fresh
@@ -225,10 +342,17 @@ module Faktory
       end
     end
 
-    def ok!
+    def ok(retval=true)
       resp = result
-      raise CommandError, resp if resp != "OK"
-      true
+      return retval if resp == "OK"
+      return resp[0].to_sym
+    end
+
+    def result!
+      resp = result
+      return nil if resp == nil
+      raise CommandError, resp[0] if !resp.is_a?(String)
+      resp
     end
 
     # FAKTORY_PROVIDER=MY_FAKTORY_URL
@@ -252,4 +376,3 @@ module Faktory
 
   end
 end
-

data/lib/faktory/connection.rb

@@ -7,7 +7,7 @@ module Faktory
       def create(options={})
         size = Faktory.worker? ? (Faktory.options[:concurrency] + 2) : 5
         ConnectionPool.new(:timeout => options[:pool_timeout] || 1, :size => size) do
-          Faktory::Client.new
+          Faktory::Client.new(**options)
         end
       end
     end

data/lib/faktory/io.rb

@@ -0,0 +1,51 @@
+require "io/wait"
+
+# this is the necessary magic to get a line-oriented protocol to
+# respect a read timeout. unfortunately Ruby sockets do not provide any
+# timeout support directly, delegating that to the IO reactor.
+module Faktory
+  class TimeoutError < Timeout::Error; end
+
+  module ReadTimeout
+    CRLF = "\r\n"
+    BUFSIZE = 16_384
+
+    # Ruby's TCP sockets do not implement timeouts.
+    # We have to implement them ourselves by using
+    # nonblocking IO and IO.select.
+    def initialize(**opts)
+      @buf = "".dup
+      @timeout = opts[:timeout] || 5
+    end
+
+    def gets
+      while (crlf = @buf.index(CRLF)).nil?
+        @buf << read_timeout(BUFSIZE)
+      end
+
+      @buf.slice!(0, crlf + 2)
+    end
+
+    def read(nbytes)
+      result = @buf.slice!(0, nbytes)
+      result << read_timeout(nbytes - result.bytesize) while result.bytesize < nbytes
+      result
+    end
+
+    private
+    def read_timeout(nbytes)
+      loop do
+        result = @sock.read_nonblock(nbytes, exception: false)
+        if result == :wait_readable
+          raise Faktory::TimeoutError unless @sock.wait_readable(@timeout)
+        elsif result == :wait_writable
+          raise Faktory::TimeoutError unless @sock.wait_writeable(@timeout)
+        elsif result == nil
+          raise Errno::ECONNRESET
+        else
+          return result
+        end
+      end
+    end
+  end
+end

data/lib/faktory/job.rb

@@ -1,4 +1,6 @@
 # frozen_string_literal: true
+require 'faktory/tracking'
+
 module Faktory
 
   ##
@@ -20,6 +22,9 @@ module Faktory
   # Note that perform_async is a class method, perform is an instance method.
   module Job
     attr_accessor :jid
+    attr_accessor :bid
+
+    include Faktory::Trackable
 
     def self.included(base)
       raise ArgumentError, "You cannot include Faktory::Job in an ActiveJob: #{base.name}" if base.ancestors.any? {|c| c.name == 'ActiveJob::Base' }
@@ -28,6 +33,16 @@ module Faktory
       base.faktory_class_attribute :faktory_options_hash
     end
 
+    def self.set(options)
+      Setter.new(options)
+    end
+
+    def batch
+      if bid
+        @batch ||= Faktory::Batch.new(bid)
+      end
+    end
+
     def logger
       Faktory.logger
     end
@@ -42,7 +57,7 @@ module Faktory
       end
 
       def perform_async(*args)
-        @opts['jobtype'.freeze].client_push(@opts.merge!('args'.freeze => args))
+        client_push(@opts.merge('args'.freeze => args))
       end
 
       # +interval+ must be a timestamp, numeric or something that acts
@@ -53,12 +68,32 @@ module Faktory
         ts = (int < 1_000_000_000 ? now + int : int)
         at = Time.at(ts).utc.to_datetime.rfc3339(9)
 
-        @opts.merge! 'args'.freeze => args, 'at'.freeze => at
+        item = @opts.merge('args'.freeze => args, 'at'.freeze => at)
+
         # Optimization to enqueue something now that is scheduled to go out now or in the past
-        @opts.delete('at'.freeze) if ts <= now
-        @opts['jobtype'.freeze].client_push(@opts)
+        item.delete('at'.freeze) if ts <= now
+
+        client_push(item)
       end
       alias_method :perform_at, :perform_in
+
+      def client_push(item) # :nodoc:
+        # stringify
+        item.keys.each do |key|
+          item[key.to_s] = item.delete(key)
+        end
+        item["jid"] ||= SecureRandom.hex(12)
+        item["queue"] ||= "default"
+
+        pool = Thread.current[:faktory_via_pool] || item["pool"] || Faktory.server_pool
+        item.delete("pool")
+
+        Faktory.client_middleware.invoke(item, pool) do
+          pool.with do |c|
+            c.push(item)
+          end
+        end
+      end
     end
 
     module ClassMethods
@@ -68,19 +103,13 @@ module Faktory
       end
 
       def perform_async(*args)
-        client_push('jobtype'.freeze => self, 'args'.freeze => args)
+        set(get_faktory_options).perform_async(*args)
       end
 
       # +interval+ must be a timestamp, numeric or something that acts
       #   numeric (like an activesupport time interval).
       def perform_in(interval, *args)
-        int = interval.to_f
-        now = Time.now.to_f
-        ts = (int < 1_000_000_000 ? now + int : int)
-        item = { 'jobtype'.freeze => self, 'args'.freeze => args }
-
-        item['at'] = Time.at(ts).utc.to_datetime.rfc3339(9) if ts > now
-        client_push(item)
+        set(get_faktory_options).perform_in(interval, *args)
       end
       alias_method :perform_at, :perform_in
 
@@ -102,23 +131,6 @@ module Faktory
         self.faktory_options_hash ||= Faktory.default_job_options
       end
 
-      def client_push(item) # :nodoc:
-        pool = Thread.current[:faktory_via_pool] || get_faktory_options['pool'.freeze] || Faktory.server_pool
-        item = get_faktory_options.merge(item)
-        # stringify
-        item.keys.each do |key|
-          item[key.to_s] = item.delete(key)
-        end
-        item["jid"] ||= SecureRandom.hex(12)
-        item["queue"] ||= "default"
-
-        Faktory.client_middleware.invoke(item, pool) do
-          pool.with do |c|
-            c.push(item)
-          end
-        end
-      end
-
       def faktory_class_attribute(*attrs)
         instance_reader = true
         instance_writer = true