que 0.14.3 → 1.0.0.beta

data/lib/que/active_job/extensions.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+module Que
+  module ActiveJob
+    # A module that devs can include into their ApplicationJob classes to get
+    # access to Que-like job behavior.
+    module JobExtensions
+      include JobMethods
+
+      def run(*args)
+        raise Error, "Job class #{self.class} didn't define a run() method!"
+      end
+
+      def perform(*args)
+        Que.internal_log(:active_job_perform, self) do
+          {args: args}
+        end
+
+        _run(
+          args: Que.recursively_freeze(
+            que_filter_args(
+              args.map { |a| a.is_a?(Hash) ? a.deep_symbolize_keys : a }
+            )
+          )
+        )
+      end
+
+      private
+
+      # Have helper methods like `destroy` and `retry_in` delegate to the actual
+      # job object. If the current job is being run through an ActiveJob adapter
+      # other than Que's, this will return nil, which is fine.
+      def que_target
+        Thread.current[:que_current_job]
+      end
+
+      # Filter out :_aj_symbol_keys constructs so that keywords work as
+      # expected.
+      def que_filter_args(thing)
+        case thing
+        when Array
+          thing.map { |t| que_filter_args(t) }
+        when Hash
+          thing.each_with_object({}) do |(k, v), hash|
+            hash[k] = que_filter_args(v) unless k == :_aj_symbol_keys
+          end
+        else
+          thing
+        end
+      end
+    end
+
+    # A module that we mix into ActiveJob's wrapper for Que::Job, to maintain
+    # backwards-compatibility with internal changes we make.
+    module WrapperExtensions
+      # The Rails adapter (built against a pre-1.0 version of this gem)
+      # assumes that it can access a job's id via job.attrs["job_id"]. So,
+      # oblige it.
+      def attrs
+        {"job_id" => que_attrs[:id]}
+      end
+
+      def run(args)
+        # Our ActiveJob extensions expect to be able to operate on the actual
+        # job object, but there's no way to access it through ActiveJob. So,
+        # scope it to the current thread. It's a bit messy, but it's the best
+        # option under the circumstances (doesn't require hacking ActiveJob in
+        # any more extensive way).
+
+        # There's no reason this logic should ever nest, because it wouldn't
+        # make sense to run a worker inside of a job, but even so, assert that
+        # nothing absurd is going on.
+        Que.assert NilClass, Thread.current[:que_current_job]
+
+        begin
+          Thread.current[:que_current_job] = self
+
+          # We symbolize the args hash but ActiveJob doesn't like that :/
+          super(args.deep_stringify_keys)
+        ensure
+          # Also assert that the current job state was only removed now, but
+          # unset the job first so that an assertion failure doesn't mess up
+          # the state any more than it already has.
+          current = Thread.current[:que_current_job]
+          Thread.current[:que_current_job] = nil
+          Que.assert(self, current)
+        end
+      end
+    end
+  end
+end
+
+class ActiveJob::QueueAdapters::QueAdapter
+  class JobWrapper < Que::Job
+    prepend Que::ActiveJob::WrapperExtensions
+  end
+end

data/lib/que/active_record/connection.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module Que
+  module ActiveRecord
+    module Connection
+      class << self
+        private
+
+        # Check out a PG::Connection object from ActiveRecord's pool.
+        def checkout
+          wrap_in_rails_executor do
+            ::ActiveRecord::Base.connection_pool.with_connection do |conn|
+               yield conn.raw_connection
+            end
+          end
+        end
+
+        # Use Rails' executor (if present) to make sure that the connection
+        # we're using isn't taken from us while the block runs. See
+        # https://github.com/chanks/que/issues/166#issuecomment-274218910
+        def wrap_in_rails_executor
+          if defined?(::Rails.application.executor)
+            ::Rails.application.executor.wrap { yield }
+          else
+            yield
+          end
+        end
+      end
+
+      module Middleware
+        class << self
+          def call(job)
+            yield
+
+            # ActiveRecord will check out connections to the current thread when
+            # queries are executed and not return them to the pool until
+            # explicitly requested to. I'm not wild about this API design, and
+            # it doesn't pose a problem for the typical case of workers using a
+            # single PG connection (since we ensure that connection is checked
+            # in and checked out responsibly), but since ActiveRecord supports
+            # connections to multiple databases, it's easy for people using that
+            # feature to unknowingly leak connections to other databases. So,
+            # take the additional step of telling ActiveRecord to check in all
+            # of the current thread's connections after each job is run.
+            ::ActiveRecord::Base.clear_active_connections!
+          end
+        end
+      end
+    end
+  end
+end

data/lib/que/active_record/model.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Que
+  module ActiveRecord
+    class Model < ::ActiveRecord::Base
+      self.table_name = :que_jobs
+
+      t = arel_table
+
+      scope :errored,     -> { where(t[:error_count].gt(0)) }
+      scope :not_errored, -> { where(t[:error_count].eq(0)) }
+
+      scope :expired,     -> { where(t[:expired_at].not_eq(nil)) }
+      scope :not_expired, -> { where(t[:expired_at].eq(nil)) }
+
+      scope :finished,     -> { where(t[:finished_at].not_eq(nil)) }
+      scope :not_finished, -> { where(t[:finished_at].eq(nil)) }
+
+      scope :scheduled,     -> { where(t[:run_at].gt("now()")) }
+      scope :not_scheduled, -> { where(t[:run_at].lteq("now()")) }
+
+      scope :ready,     -> { not_errored.not_expired.not_finished.not_scheduled }
+      scope :not_ready, -> { where(t[:error_count].gt(0).or(t[:expired_at].not_eq(nil)).or(t[:finished_at].not_eq(nil)).or(t[:run_at].gt("now()"))) }
+
+      class << self
+        def by_job_class(job_class)
+          job_class = job_class.name if job_class.is_a?(Class)
+          where(
+            "que_jobs.job_class = ? OR (que_jobs.job_class = 'ActiveJob::QueueAdapters::QueAdapter::JobWrapper' AND que_jobs.args->0->>'job_class' = ?)",
+            job_class, job_class,
+          )
+        end
+
+        def by_queue(queue)
+          where(arel_table[:queue].eq(queue))
+        end
+
+        def by_tag(tag)
+          where("que_jobs.data @> ?", JSON.dump(tags: [tag]))
+        end
+
+        def by_args(*args)
+          where("que_jobs.args @> ?", JSON.dump(args))
+        end
+      end
+    end
+  end
+end

data/lib/que/connection.rb

@@ -0,0 +1,179 @@
+# frozen_string_literal: true
+
+# A simple wrapper class around connections that basically just improves the
+# query API a bit. Currently, our connection pool wrapper discards these
+# connection wrappers once the connection is returned to the source connection
+# pool, so this class isn't currently suitable for storing data about the
+# connection long-term (like what statements it has prepared, for example).
+
+# If we wanted to do that, we'd probably need to sneak a reference to the
+# wrapper into the PG::Connection object itself, by just setting a instance
+# variable that's something namespaced and hopefully safe, like
+# `@que_connection_wrapper`. It's a bit ugly, but it should ensure that we don't
+# cause any memory leaks in esoteric setups where one-off connections are being
+# established and then garbage-collected.
+
+require 'time' # For Time.parse
+require 'set'
+
+module Que
+  class Connection
+    extend Forwardable
+
+    attr_reader :wrapped_connection
+
+    def_delegators :wrapped_connection, :backend_pid, :wait_for_notify
+
+    class << self
+      def wrap(conn)
+        case conn
+        when self
+          conn
+        when PG::Connection
+          conn.instance_variable_get(:@que_wrapper) ||
+          conn.instance_variable_set(:@que_wrapper, new(conn))
+        else
+          raise Error, "Unsupported input for Connection.wrap: #{conn.class}"
+        end
+      end
+    end
+
+    def initialize(connection)
+      @wrapped_connection = connection
+      @prepared_statements = Set.new
+    end
+
+    def execute(command, params = nil)
+      sql =
+        case command
+        when Symbol then SQL[command]
+        when String then command
+        else raise Error, "Bad command! #{command.inspect}"
+        end
+
+      params = convert_params(params) if params
+      result = execute_sql(sql, params)
+
+      Que.internal_log :connection_execute, self do
+        {
+          backend_pid: backend_pid,
+          command:     command,
+          params:      params,
+          ntuples:     result.ntuples,
+        }
+      end
+
+      convert_result(result)
+    end
+
+    def execute_prepared(command, params = nil)
+      Que.assert(Symbol, command)
+
+      if !Que.use_prepared_statements || in_transaction?
+        return execute(command, params)
+      end
+
+      name = "que_#{command}"
+
+      begin
+        unless @prepared_statements.include?(command)
+          wrapped_connection.prepare(name, SQL[command])
+          @prepared_statements.add(command)
+          prepared_just_now = true
+        end
+
+        convert_result(
+          wrapped_connection.exec_prepared(name, params)
+        )
+      rescue ::PG::InvalidSqlStatementName => error
+        # Reconnections on ActiveRecord can cause the same connection
+        # objects to refer to new backends, so recover as well as we can.
+
+        unless prepared_just_now
+          Que.log level: :warn, event: :reprepare_statement, command: command
+          @prepared_statements.delete(command)
+          retry
+        end
+
+        raise error
+      end
+    end
+
+    def next_notification
+      wrapped_connection.notifies
+    end
+
+    def drain_notifications
+      loop { break if next_notification.nil? }
+    end
+
+    def in_transaction?
+      wrapped_connection.transaction_status != ::PG::PQTRANS_IDLE
+    end
+
+    private
+
+    def convert_params(params)
+      params.map do |param|
+        case param
+        when Time
+          # The pg gem unfortunately doesn't convert fractions of time
+          # instances, so cast them to a string.
+          param.strftime('%Y-%m-%d %H:%M:%S.%6N %z')
+        when Array, Hash
+          # Handle JSON.
+          Que.serialize_json(param)
+        else
+          param
+        end
+      end
+    end
+
+    def execute_sql(sql, params)
+      # Some versions of the PG gem dislike an empty/nil params argument.
+      if params && !params.empty?
+        wrapped_connection.async_exec(sql, params)
+      else
+        wrapped_connection.async_exec(sql)
+      end
+    end
+
+    # Procs used to convert strings from Postgres into Ruby types.
+    CAST_PROCS = {
+      # Boolean
+      16   => 't'.method(:==),
+      # Timestamp with time zone
+      1184 => Time.method(:parse),
+    }
+
+    # JSON, JSONB
+    CAST_PROCS[114] = CAST_PROCS[3802] = -> (j) { Que.deserialize_json(j) }
+
+    # Integer, bigint, smallint
+    CAST_PROCS[23] = CAST_PROCS[20] = CAST_PROCS[21] = proc(&:to_i)
+
+    CAST_PROCS.freeze
+
+    def convert_result(result)
+      output = result.to_a
+
+      result.fields.each_with_index do |field, index|
+        symbol = field.to_sym
+
+        if converter = CAST_PROCS[result.ftype(index)]
+          output.each do |hash|
+            value = hash.delete(field)
+            value = converter.call(value) if value
+            hash[symbol] = value
+          end
+        else
+          output.each do |hash|
+            hash[symbol] = hash.delete(field)
+          end
+        end
+      end
+
+      output
+    end
+  end
+end

data/lib/que/connection_pool.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+# A wrapper around whatever connection pool we're using. Mainly just asserts
+# that the source connection pool is reentrant and thread-safe.
+
+module Que
+  class ConnectionPool
+    def initialize(&block)
+      @connection_proc = block
+      @checked_out     = Set.new
+      @mutex           = Mutex.new
+      @thread_key      = "que_connection_pool_#{object_id}".to_sym
+    end
+
+    def checkout
+      # Do some asserting to ensure that the connection pool we're using is
+      # behaving properly.
+      @connection_proc.call do |conn|
+        # Did this pool already have a connection for this thread?
+        preexisting = wrapped = current_connection
+
+        begin
+          if preexisting
+            # If so, check that the connection we just got is the one we expect.
+            if preexisting.wrapped_connection.backend_pid != conn.backend_pid
+              raise Error, "Connection pool is not reentrant! previous: #{preexisting.wrapped_connection.inspect} now: #{conn.inspect}"
+            end
+          else
+            # If not, make sure that it wasn't promised to any other threads.
+            sync do
+              Que.assert(@checked_out.add?(conn.backend_pid)) do
+                "Connection pool didn't synchronize access properly! (entrance: #{conn.backend_pid})"
+              end
+            end
+
+            self.current_connection = wrapped = Connection.wrap(conn)
+          end
+
+          yield(wrapped)
+        ensure
+          if preexisting.nil?
+            # We're at the top level (about to return this connection to the
+            # pool we got it from), so mark it as no longer ours.
+            self.current_connection = nil
+
+            sync do
+              Que.assert(@checked_out.delete?(conn.backend_pid)) do
+                "Connection pool didn't synchronize access properly! (exit: #{conn.backend_pid})"
+              end
+            end
+          end
+        end
+      end
+    end
+
+    def execute(*args)
+      checkout { |conn| conn.execute(*args) }
+    end
+
+    def in_transaction?
+      checkout { |conn| conn.in_transaction? }
+    end
+
+    private
+
+    def sync
+      @mutex.synchronize { yield }
+    end
+
+    def current_connection
+      Thread.current[@thread_key]
+    end
+
+    def current_connection=(c)
+      Thread.current[@thread_key] = c
+    end
+  end
+end