que 0.14.3 → 1.0.0.beta

data/lib/que/job_methods.rb

@@ -0,0 +1,168 @@
+# frozen_string_literal: true
+
+module Que
+  SQL[:finish_job] =
+    %{
+      UPDATE public.que_jobs
+      SET finished_at = now()
+      WHERE id = $1::bigint
+    }
+
+  SQL[:expire_job] =
+    %{
+      UPDATE public.que_jobs
+      SET error_count = error_count + 1,
+          expired_at = now()
+      WHERE id = $1::bigint
+    }
+
+  SQL[:destroy_job] =
+    %{
+      DELETE FROM public.que_jobs
+      WHERE id = $1::bigint
+    }
+
+  SQL[:set_error] =
+    %{
+      UPDATE public.que_jobs
+      SET error_count          = error_count + 1,
+          run_at               = now() + $1::float * '1 second'::interval,
+          last_error_message   = left($2::text, 500),
+          last_error_backtrace = left($3::text, 10000)
+      WHERE id = $4::bigint
+    }
+
+  module JobMethods
+    # Note that we delegate almost all methods to the result of the que_target
+    # method, which could be one of a few things, depending on the circumstance.
+
+    # Run the job with error handling and cleanup logic. Optionally support
+    # overriding the args, because it's necessary when jobs are invoked from
+    # ActiveJob.
+    def _run(args: nil, reraise_errors: false)
+      if args.nil? && que_target
+        args = que_target.que_attrs.fetch(:args)
+      end
+
+      run(*args)
+      default_resolve_action if que_target && !que_target.que_resolved
+    rescue => error
+      raise error unless que_target
+
+      que_target.que_error = error
+
+      run_error_notifier =
+        begin
+          handle_error(error)
+        rescue => error_2
+          Que.notify_error(error_2, que_target.que_attrs)
+          true
+        end
+
+      Que.notify_error(error, que_target.que_attrs) if run_error_notifier
+      retry_in_default_interval unless que_target.que_resolved
+
+      raise error if reraise_errors
+    end
+
+    private
+
+    # This method defines the object on which the various job helper methods are
+    # acting. When using Que in the default configuration this will just be
+    # self, but when using the Que adapter for ActiveJob it'll be the actual
+    # underlying job object. When running an ActiveJob::Base subclass that
+    # includes this module through a separate adapter this will be nil - hence,
+    # the defensive coding in every method that no-ops if que_target is falsy.
+    def que_target
+      raise NotImplementedError
+    end
+
+    def resolve_que_setting(*args)
+      return unless que_target
+
+      que_target.class.resolve_que_setting(*args)
+    end
+
+    def default_resolve_action
+      return unless que_target
+
+      destroy
+    end
+
+    def expire
+      return unless que_target
+
+      if id = que_target.que_attrs[:id]
+        Que.execute :expire_job, [id]
+      end
+
+      que_target.que_resolved = true
+    end
+
+    def finish
+      return unless que_target
+
+      if id = que_target.que_attrs[:id]
+        Que.execute :finish_job, [id]
+      end
+
+      que_target.que_resolved = true
+    end
+
+    def error_count
+      return 0 unless que_target
+
+      count = que_target.que_attrs.fetch(:error_count)
+      que_target.que_error ? count + 1 : count
+    end
+
+    # To be overridden in subclasses.
+    def handle_error(error)
+      return unless que_target
+
+      max = resolve_que_setting(:maximum_retry_count)
+
+      if max && error_count > max
+        expire
+      else
+        retry_in_default_interval
+      end
+    end
+
+    def retry_in_default_interval
+      return unless que_target
+
+      retry_in(resolve_que_setting(:retry_interval, error_count))
+    end
+
+    # Explicitly check for the job id in these helpers, because it won't exist
+    # if we're running synchronously.
+    def retry_in(period)
+      return unless que_target
+
+      if id = que_target.que_attrs[:id]
+        values = [period]
+
+        if e = que_target.que_error
+          values << "#{e.class}: #{e.message}".slice(0, 500) << e.backtrace.join("\n").slice(0, 10000)
+        else
+          values << nil << nil
+        end
+
+        Que.execute :set_error, values << id
+      end
+
+      que_target.que_resolved = true
+    end
+
+    def destroy
+      return unless que_target
+
+      if id = que_target.que_attrs[:id]
+        Que.execute :destroy_job, [id]
+      end
+
+      que_target.que_resolved = true
+    end
+  end
+end

data/lib/que/listener.rb

@@ -0,0 +1,176 @@
+# frozen_string_literal: true
+
+module Que
+  class Listener
+    MESSAGE_FORMATS = {}
+
+    attr_reader :connection, :channel
+
+    def initialize(connection:, channel: nil)
+      @connection = connection
+      @channel    = channel || "que_listener_#{connection.backend_pid}"
+
+      Que.internal_log :listener_instantiate, self do
+        {
+          backend_pid: connection.backend_pid,
+        }
+      end
+    end
+
+    def listen
+      connection.execute "LISTEN #{channel}"
+    end
+
+    def wait_for_grouped_messages(timeout)
+      messages = wait_for_messages(timeout)
+
+      output = {}
+
+      messages.each do |message|
+        message_type = message.delete(:message_type)
+
+        (output[message_type.to_sym] ||= []) << message.freeze
+      end
+
+      output
+    end
+
+    def wait_for_messages(timeout)
+      # Make sure we never pass nil to this method, so we don't hang the thread.
+      Que.assert(Numeric, timeout)
+
+      Que.internal_log :listener_waiting, self do
+        {
+          backend_pid: connection.backend_pid,
+          channel:     channel,
+          timeout:     timeout,
+        }
+      end
+
+      accumulated_messages = []
+
+      # Notifications often come in batches (especially when a transaction that
+      # inserted many jobs commits), so we want to loop and pick up all the
+      # received notifications before continuing.
+      loop do
+        notification_received =
+          connection.wait_for_notify(timeout) do |channel, pid, payload|
+            # We've received a notification, so zero out the timeout before we
+            # loop again to check for another message. This ensures that we
+            # don't wait an additional `timeout` seconds after processing the
+            # final message before this method returns.
+            timeout = 0
+
+            Que.internal_log(:listener_received_notification, self) do
+              {
+                channel:     channel,
+                backend_pid: connection.backend_pid,
+                source_pid:  pid,
+                payload:     payload,
+              }
+            end
+
+            # Be very defensive about the message we receive - it may not be
+            # valid JSON or have the structure we expect.
+            next unless message = parse_payload(payload)
+
+            case message
+            when Array then accumulated_messages.concat(message)
+            when Hash  then accumulated_messages << message
+            else raise Error, "Unexpected parse_payload output: #{message.class}"
+            end
+          end
+
+        break unless notification_received
+      end
+
+      return accumulated_messages if accumulated_messages.empty?
+
+      Que.internal_log(:listener_received_messages, self) do
+        {
+          backend_pid: connection.backend_pid,
+          channel:     channel,
+          messages:    accumulated_messages,
+        }
+      end
+
+      accumulated_messages.keep_if do |message|
+        next unless message.is_a?(Hash)
+        next unless type = message[:message_type]
+        next unless type.is_a?(String)
+        next unless format = MESSAGE_FORMATS[type.to_sym]
+
+        if message_matches_format?(message, format)
+          true
+        else
+          error_message = [
+            "Message of type '#{type}' doesn't match format!",
+            # Massage message and format a bit to make these errors more readable.
+            "Message: #{Hash[message.reject{|k,v| k == :message_type}.sort_by{|k,v| k}].inspect}",
+            "Format: #{Hash[format.sort_by{|k,v| k}].inspect}",
+          ].join("\n")
+
+          Que.notify_error_async(Error.new(error_message))
+          false
+        end
+      end
+
+      Que.internal_log(:listener_filtered_messages, self) do
+        {
+          backend_pid: connection.backend_pid,
+          channel:     channel,
+          messages:    accumulated_messages,
+        }
+      end
+
+      accumulated_messages
+    end
+
+    def unlisten
+      # Be sure to drain all notifications so that any code that uses this
+      # connection later doesn't receive any nasty surprises.
+      connection.execute "UNLISTEN *"
+      connection.drain_notifications
+
+      Que.internal_log :listener_unlisten, self do
+        {
+          backend_pid: connection.backend_pid,
+          channel:     channel,
+        }
+      end
+    end
+
+    private
+
+    def parse_payload(payload)
+      Que.deserialize_json(payload)
+    rescue JSON::ParserError => e
+      Que.notify_error_async(e)
+      nil
+    end
+
+    def message_matches_format?(message, format)
+      message_has_all_keys?(message, format) &&
+        message_has_no_excess_keys?(message, format) &&
+        message_keys_all_valid?(message, format)
+    end
+
+    def message_has_all_keys?(message, format)
+      format.all? { |k,v| message.has_key?(k) }
+    end
+
+    def message_has_no_excess_keys?(message, format)
+      message.all? { |k,v| format.has_key?(k) || k == :message_type }
+    end
+
+    def message_keys_all_valid?(message, format)
+      message.all? do |key, value|
+        if type = format[key]
+          Que.assert?(type, value)
+        else
+          true
+        end
+      end
+    end
+  end
+end

data/lib/que/locker.rb

@@ -0,0 +1,466 @@
+# frozen_string_literal: true
+
+# The Locker class encapsulates a thread that is listening/polling for new
+# jobs in the DB, locking them, passing their primary keys to workers, then
+# cleaning up by unlocking them once the workers are done.
+
+require 'set'
+
+module Que
+  Listener::MESSAGE_FORMATS[:job_available] =
+    {
+      queue:    String,
+      id:       Integer,
+      run_at:   TIME_REGEX,
+      priority: Integer,
+    }
+
+  SQL[:clean_lockers] =
+    %{
+      DELETE FROM public.que_lockers
+      WHERE pid = pg_backend_pid()
+      OR pid NOT IN (SELECT pid FROM pg_stat_activity)
+    }
+
+  SQL[:register_locker] =
+    %{
+      INSERT INTO public.que_lockers
+      (
+        pid,
+        worker_count,
+        worker_priorities,
+        ruby_pid,
+        ruby_hostname,
+        listening,
+        queues
+      )
+      VALUES
+      (
+        pg_backend_pid(),
+        $1::integer,
+        $2::integer[],
+        $3::integer,
+        $4::text,
+        $5::boolean,
+        $6::text[]
+      )
+    }
+
+  class Locker
+    attr_reader :thread, :workers, :job_cache, :locks
+
+    MESSAGE_RESOLVERS = {}
+    RESULT_RESOLVERS  = {}
+
+    MESSAGE_RESOLVERS[:job_available] =
+      -> (messages) {
+        metajobs = messages.map { |key| Metajob.new(key) }
+        push_jobs(lock_jobs(job_cache.accept?(metajobs)))
+      }
+
+    RESULT_RESOLVERS[:job_finished] =
+      -> (messages) { finish_jobs(messages.map{|m| m.fetch(:metajob)}) }
+
+    DEFAULT_POLL_INTERVAL      = 5.0
+    DEFAULT_WAIT_PERIOD        = 50
+    DEFAULT_MINIMUM_QUEUE_SIZE = 2
+    DEFAULT_MAXIMUM_QUEUE_SIZE = 8
+    DEFAULT_WORKER_COUNT       = 6
+    DEFAULT_WORKER_PRIORITIES  = [10, 30, 50].freeze
+
+    def initialize(
+      queues:             [Que.default_queue],
+      connection:         nil,
+      listen:             true,
+      poll:               true,
+      poll_interval:      DEFAULT_POLL_INTERVAL,
+      wait_period:        DEFAULT_WAIT_PERIOD,
+      maximum_queue_size: DEFAULT_MAXIMUM_QUEUE_SIZE,
+      minimum_queue_size: DEFAULT_MINIMUM_QUEUE_SIZE,
+      worker_count:       DEFAULT_WORKER_COUNT,
+      worker_priorities:  DEFAULT_WORKER_PRIORITIES,
+      on_worker_start:    nil
+    )
+
+      # Sanity-check all our arguments, since some users may instantiate Locker
+      # directly.
+      Que.assert [TrueClass, FalseClass], listen
+      Que.assert [TrueClass, FalseClass], poll
+
+      Que.assert Numeric, poll_interval
+      Que.assert Numeric, wait_period
+      Que.assert Integer, worker_count
+
+      Que.assert Array, worker_priorities
+      worker_priorities.each { |p| Que.assert(Integer, p) }
+
+      all_worker_priorities = worker_priorities.values_at(0...worker_count)
+
+      # We use a JobCache to track jobs and pass them to workers, and a
+      # ResultQueue to receive messages from workers.
+      @job_cache = JobCache.new(
+        maximum_size: maximum_queue_size,
+        minimum_size: minimum_queue_size,
+        priorities:   all_worker_priorities.uniq,
+      )
+
+      @result_queue = ResultQueue.new
+
+      Que.internal_log :locker_instantiate, self do
+        {
+          queues:             queues,
+          listen:             listen,
+          poll:               poll,
+          poll_interval:      poll_interval,
+          wait_period:        wait_period,
+          maximum_queue_size: maximum_queue_size,
+          minimum_queue_size: minimum_queue_size,
+          worker_count:       worker_count,
+          worker_priorities:  worker_priorities,
+        }
+      end
+
+      # Local cache of which advisory locks are held by this connection.
+      @locks = Set.new
+
+      @queue_names = queues.is_a?(Hash) ? queues.keys : queues
+      @wait_period = wait_period.to_f / 1000 # Milliseconds to seconds.
+
+      # If the worker_count exceeds the array of priorities it'll result in
+      # extra workers that will work jobs of any priority. For example, the
+      # default worker_count of 6 and the default worker priorities of [10, 30,
+      # 50] will result in three workers that only work jobs that meet those
+      # priorities, and three workers that will work any job.
+      @workers =
+        all_worker_priorities.map do |priority|
+          Worker.new(
+            priority:       priority,
+            job_cache:      @job_cache,
+            result_queue:   @result_queue,
+            start_callback: on_worker_start,
+          )
+        end
+
+      # To prevent race conditions, let every worker get into a ready state
+      # before starting up the locker thread.
+      loop do
+        break if job_cache.waiting_count == workers.count
+        sleep 0.001
+      end
+
+      pool =
+        if connection
+          # Wrap the given connection in a dummy connection pool.
+          ConnectionPool.new { |&block| block.call(connection) }
+        else
+          Que.pool
+        end
+
+      @thread =
+        Thread.new do
+          # An error causing this thread to exit is a bug in Que, which we want
+          # to know about ASAP, so propagate the error if it happens.
+          Thread.current.abort_on_exception = true
+
+          # Give this thread priority, so it can promptly respond to NOTIFYs.
+          Thread.current.priority = 1
+
+          pool.checkout do |connection|
+            original_application_name =
+              connection.
+              execute("SHOW application_name").
+              first.
+              fetch(:application_name)
+
+            begin
+              @connection = connection
+
+              connection.execute(
+                "SELECT set_config('application_name', $1, false)",
+                ["Que Locker: #{connection.backend_pid}"]
+              )
+
+              Poller.setup(connection)
+
+              if listen
+                @listener = Listener.new(connection: connection)
+              end
+
+              if poll
+                @pollers =
+                  queues.map do |queue, interval|
+                    Poller.new(
+                      connection:    connection,
+                      queue:         queue,
+                      poll_interval: interval || poll_interval,
+                    )
+                  end
+              end
+
+              work_loop
+            ensure
+              connection.execute(
+                "SELECT set_config('application_name', $1, false)",
+                [original_application_name]
+              )
+
+              Poller.cleanup(connection)
+            end
+          end
+        end
+    end
+
+    def stop!
+      stop; wait_for_stop
+    end
+
+    def stop
+      @job_cache.stop
+      @stop = true
+    end
+
+    def stopping?
+      @stop
+    end
+
+    def wait_for_stop
+      @thread.join
+    end
+
+    private
+
+    attr_reader :connection, :pollers
+
+    def work_loop
+      Que.log(
+        level: :debug,
+        event: :locker_start,
+        queues: @queue_names,
+      )
+
+      Que.internal_log :locker_start, self do
+        {
+          backend_pid: connection.backend_pid,
+          worker_priorities: workers.map(&:priority),
+          pollers: pollers && pollers.map { |p| [p.queue, p.poll_interval] }
+        }
+      end
+
+      begin
+        @listener.listen if @listener
+
+        # A previous locker that didn't exit cleanly may have left behind
+        # a bad locker record, so clean up before registering.
+        connection.execute :clean_lockers
+        connection.execute :register_locker, [
+          @workers.count,
+          "{#{@workers.map(&:priority).map{|p| p || 'NULL'}.join(',')}}",
+          Process.pid,
+          CURRENT_HOSTNAME,
+          !!@listener,
+          "{\"#{@queue_names.join('","')}\"}",
+        ]
+
+        {} while cycle
+
+        Que.log(
+          level: :debug,
+          event: :locker_stop,
+        )
+
+        unlock_jobs(@job_cache.clear)
+
+        @workers.each(&:wait_until_stopped)
+
+        handle_results
+      ensure
+        connection.execute :clean_lockers
+
+        @listener.unlisten if @listener
+      end
+    end
+
+    def cycle
+      # Poll at the start of a cycle, so that when the worker starts up we can
+      # load up the queue with jobs immediately.
+      poll
+
+      # If we got the stop call while we were polling, break before going to
+      # sleep.
+      return if @stop
+
+      # The main sleeping part of the cycle. If this is a listening locker, this
+      # is where we wait for notifications.
+      wait
+
+      # Manage any job output we got while we were sleeping.
+      handle_results
+
+      # If we haven't gotten the stop signal, cycle again.
+      !@stop
+    end
+
+    def poll
+      # Only poll when there are pollers to use (that is, when polling is
+      # enabled) and when the local queue has dropped below the configured
+      # minimum size.
+      return unless pollers && job_cache.jobs_needed?
+
+      pollers.each do |poller|
+        priorities = job_cache.available_priorities
+        break if priorities.empty?
+
+        Que.internal_log(:locker_polling, self) { {priorities: priorities, held_locks: @locks.to_a, queue: poller.queue} }
+
+        if metajobs = poller.poll(priorities: priorities, held_locks: @locks)
+          metajobs.each do |metajob|
+            mark_id_as_locked(metajob.id)
+          end
+
+          push_jobs(metajobs)
+        end
+      end
+    end
+
+    def wait
+      if @listener
+        @listener.wait_for_grouped_messages(@wait_period).each do |type, messages|
+          if resolver = MESSAGE_RESOLVERS[type]
+            instance_exec messages, &resolver
+          else
+            raise Error, "Unexpected message type: #{type.inspect}"
+          end
+        end
+      else
+        sleep(@wait_period)
+      end
+    end
+
+    def handle_results
+      messages_by_type =
+        @result_queue.clear.group_by{|r| r.fetch(:message_type)}
+
+      messages_by_type.each do |type, messages|
+        if resolver = RESULT_RESOLVERS[type]
+          instance_exec messages, &resolver
+        else
+          raise Error, "Unexpected result type: #{type.inspect}"
+        end
+      end
+    end
+
+    def lock_jobs(metajobs)
+      metajobs.reject! { |m| @locks.include?(m.id) }
+      return metajobs if metajobs.empty?
+
+      ids = metajobs.map{|m| m.id.to_i}
+
+      Que.internal_log :locker_locking, self do
+        {
+          backend_pid: connection.backend_pid,
+          ids:         ids,
+        }
+      end
+
+      jobs =
+        connection.execute \
+          <<-SQL
+            WITH jobs AS (
+              SELECT * FROM que_jobs WHERE id IN (#{ids.join(', ')})
+            )
+            SELECT * FROM jobs WHERE pg_try_advisory_lock(id)
+          SQL
+
+      jobs_by_id = {}
+
+      jobs.each do |job|
+        id = job.fetch(:id)
+        mark_id_as_locked(id)
+        jobs_by_id[id] = job
+      end
+
+      metajobs.keep_if do |metajob|
+        if job = jobs_by_id[metajob.id]
+          metajob.set_job(job)
+          true
+        else
+          false
+        end
+      end
+    end
+
+    def push_jobs(metajobs)
+      return if metajobs.empty?
+
+      # First check that the jobs are all still visible/available in the DB.
+      ids = metajobs.map(&:id)
+
+      verified_ids =
+        connection.execute(
+          <<-SQL
+            SELECT id
+            FROM public.que_jobs
+            WHERE finished_at IS NULL
+              AND expired_at IS NULL
+              AND id IN (#{ids.join(', ')})
+          SQL
+        ).map{|h| h[:id]}.to_set
+
+      good, bad = metajobs.partition{|mj| verified_ids.include?(mj.id)}
+
+      displaced = @job_cache.push(*good) || []
+
+      # Unlock any low-importance jobs the new ones may displace.
+      if bad.any? || displaced.any?
+        unlock_jobs(bad + displaced)
+      end
+    end
+
+    def finish_jobs(metajobs)
+      unlock_jobs(metajobs)
+    end
+
+    def unlock_jobs(metajobs)
+      return if metajobs.empty?
+
+      # Unclear how untrusted input would get passed to this method, but since
+      # we need string interpolation here, make sure we only have integers.
+      ids = metajobs.map { |job| job.id.to_i }
+
+      Que.internal_log :locker_unlocking, self do
+        {
+          backend_pid: connection.backend_pid,
+          ids:         ids,
+        }
+      end
+
+      values = ids.join('), (')
+
+      results =
+        connection.execute \
+          "SELECT pg_advisory_unlock(v.i) FROM (VALUES (#{values})) v (i)"
+
+      results.each do |result|
+        Que.assert(result.fetch(:pg_advisory_unlock)) do
+          [
+            "Tried to unlock a job we hadn't locked!",
+            results.inspect,
+            ids.inspect,
+          ].join(' ')
+        end
+      end
+
+      ids.each do |id|
+        Que.assert(@locks.delete?(id)) do
+          "Tried to remove a local lock that didn't exist!: #{id}"
+        end
+      end
+    end
+
+    def mark_id_as_locked(id)
+      Que.assert(@locks.add?(id)) do
+        "Tried to lock a job that was already locked: #{id}"
+      end
+    end
+  end
+end