que 1.0.0.beta → 1.0.0.beta2

data/lib/que/connection.rb

@@ -30,8 +30,11 @@ module Que
         when self
           conn
         when PG::Connection
-          conn.instance_variable_get(:@que_wrapper) ||
-          conn.instance_variable_set(:@que_wrapper, new(conn))
+          if conn.instance_variable_defined?(:@que_wrapper)
+            conn.instance_variable_get(:@que_wrapper)
+          else
+            conn.instance_variable_set(:@que_wrapper, new(conn))
+          end
         else
           raise Error, "Unsupported input for Connection.wrap: #{conn.class}"
         end
@@ -43,7 +46,7 @@ module Que
       @prepared_statements = Set.new
     end
 
-    def execute(command, params = nil)
+    def execute(command, params = [])
       sql =
         case command
         when Symbol then SQL[command]
@@ -51,8 +54,17 @@ module Que
         else raise Error, "Bad command! #{command.inspect}"
         end
 
-      params = convert_params(params) if params
-      result = execute_sql(sql, params)
+      params = convert_params(params)
+
+      result =
+        Que.run_sql_middleware(sql, params) do
+          # Some versions of the PG gem dislike an empty/nil params argument.
+          if params.empty?
+            wrapped_connection.async_exec(sql)
+          else
+            wrapped_connection.async_exec(sql, params)
+          end
+        end
 
       Que.internal_log :connection_execute, self do
         {
@@ -129,19 +141,16 @@ module Que
       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(:==),
+      16 => -> (value) {
+        case value
+        when String then value == 't'.freeze
+        else !!value
+        end
+      },
+
       # Timestamp with time zone
       1184 => Time.method(:parse),
     }

data/lib/que/connection_pool.rb

@@ -63,8 +63,8 @@ module Que
 
     private
 
-    def sync
-      @mutex.synchronize { yield }
+    def sync(&block)
+      @mutex.synchronize(&block)
     end
 
     def current_connection

data/lib/que/job.rb

@@ -134,7 +134,7 @@ module Que
         Que.recursively_freeze(attrs)
 
         new(attrs).tap do |job|
-          Que.run_middleware(job) do
+          Que.run_job_middleware(job) do
             job._run(reraise_errors: true)
           end
         end

data/lib/que/{job_cache.rb → job_buffer.rb}

@@ -5,63 +5,70 @@
 # minimum priority, and stopping gracefully.
 
 module Que
-  class JobCache
+  class JobBuffer
     attr_reader :maximum_size, :minimum_size, :priority_queues
 
+    # Since we use a mutex, which is not reentrant, we have to be a little
+    # careful to not call a method that locks the mutex when we've already
+    # locked it. So, as a general rule, public methods handle locking the mutex
+    # when necessary, while private methods handle the actual underlying data
+    # changes. This lets us reuse those private methods without running into
+    # locking issues.
+
     def initialize(
       maximum_size:,
       minimum_size:,
       priorities:
     )
       @maximum_size = Que.assert(Integer, maximum_size)
-      Que.assert(maximum_size >= 0) { "maximum_size for a JobCache must be at least zero!" }
+      Que.assert(maximum_size >= 0) { "maximum_size for a JobBuffer must be at least zero!" }
 
       @minimum_size = Que.assert(Integer, minimum_size)
-      Que.assert(minimum_size >= 0) { "minimum_size for a JobCache must be at least zero!" }
+      Que.assert(minimum_size >= 0) { "minimum_size for a JobBuffer must be at least zero!" }
 
       Que.assert(minimum_size <= maximum_size) do
-        "minimum queue size (#{minimum_size}) is " \
-          "greater than the maximum queue size (#{maximum_size})!"
+        "minimum buffer size (#{minimum_size}) is " \
+          "greater than the maximum buffer size (#{maximum_size})!"
       end
 
-      @stop    = false
-      @array   = []
-      @monitor = Monitor.new # TODO: Make this a mutex?
+      @stop  = false
+      @array = []
+      @mutex = Mutex.new
 
-      # Make sure that priority = nil sorts highest.
       @priority_queues = Hash[
-        priorities.sort_by{|p| p || Float::INFINITY}.map do |p|
-          [p, PriorityQueue.new(priority: p, job_cache: self)]
+        # Make sure that priority = nil sorts highest.
+        priorities.sort_by{|p| p || MAXIMUM_PRIORITY}.map do |p|
+          [p, PriorityQueue.new(priority: p, job_buffer: self)]
         end
       ].freeze
     end
 
     def push(*metajobs)
-      Que.internal_log(:job_cache_push, self) do
+      Que.internal_log(:job_buffer_push, self) do
         {
           maximum_size:  maximum_size,
           ids:           metajobs.map(&:id),
-          current_queue: @array,
+          current_queue: to_a,
         }
       end
 
       sync do
-        return metajobs if stopping?
+        return metajobs if _stopping?
 
-        @array.push(*metajobs).sort!
+        @array.concat(metajobs).sort!
 
         # Relying on the hash's contents being sorted, here.
         priority_queues.reverse_each do |_, pq|
           pq.waiting_count.times do
-            job = shift_job(pq.priority)
-            break if job.nil?
+            job = _shift_job(pq.priority)
+            break if job.nil? # False would mean we're stopping.
             pq.push(job)
           end
         end
 
-        # If we passed the maximum queue size, drop the lowest sort keys and
+        # If we passed the maximum buffer size, drop the lowest sort keys and
         # return their ids to be unlocked.
-        overage = -cache_space
+        overage = -_buffer_space
         pop(overage) if overage > 0
       end
     end
@@ -72,22 +79,16 @@ module Que
     end
 
     def shift_job(priority = nil)
-      sync do
-        if stopping?
-          false
-        elsif (job = @array.first) && job.priority_sufficient?(priority)
-          @array.shift
-        end
-      end
+      sync { _shift_job(priority) }
     end
 
     def accept?(metajobs)
-      return [] if stopping?
-
       metajobs.sort!
 
       sync do
-        start_index = cache_space
+        return [] if _stopping?
+
+        start_index = _buffer_space
         final_index = metajobs.length - 1
 
         return metajobs if start_index > final_index
@@ -126,7 +127,7 @@ module Que
         count = pq.waiting_count
 
         if lowest_priority
-          count += cache_space
+          count += buffer_space
           lowest_priority = false
         end
 
@@ -136,14 +137,12 @@ module Que
       hash
     end
 
-    def cache_space
-      sync do
-        maximum_size - size
-      end
+    def buffer_space
+      sync { _buffer_space }
     end
 
     def size
-      sync { @array.size }
+      sync { _size }
     end
 
     def to_a
@@ -156,40 +155,60 @@ module Que
     end
 
     def clear
-      sync { pop(size) }
+      sync { pop(_size) }
     end
 
     def stopping?
-      sync { !!@stop }
+      sync { _stopping? }
     end
 
     private
 
+    def _buffer_space
+      maximum_size - _size
+    end
+
     def pop(count)
       @array.pop(count)
     end
 
-    def sync
-      @monitor.synchronize { yield }
+    def _shift_job(priority)
+      if _stopping?
+        false
+      elsif (job = @array.first) && job.priority_sufficient?(priority)
+        @array.shift
+      end
+    end
+
+    def _size
+      @array.size
+    end
+
+    def _stopping?
+      !!@stop
+    end
+
+    def sync(&block)
+      @mutex.synchronize(&block)
     end
 
     # A queue object dedicated to a specific worker priority. It's basically a
     # Queue object from the standard library, but it's able to reach into the
-    # JobCache's cache in order to satisfy a pop.
+    # JobBuffer's buffer in order to satisfy a pop.
     class PriorityQueue
-      attr_reader :job_cache, :priority
+      attr_reader :job_buffer, :priority, :mutex
 
       def initialize(
-        job_cache:,
+        job_buffer:,
         priority:
       )
-        @job_cache = job_cache
-        @priority  = priority
-        @waiting   = 0
-        @stopping  = false
-        @items     = [] # Items pending distribution to waiting threads.
-        @monitor   = Monitor.new
-        @cv        = Monitor::ConditionVariable.new(@monitor)
+        @job_buffer = job_buffer
+        @priority   = priority
+        @waiting    = 0
+        @stopping   = false
+        @items      = [] # Items pending distribution to waiting threads.
+        @mutex      = Mutex.new
+        @cv         = ConditionVariable.new
       end
 
       def pop
@@ -201,11 +220,11 @@ module Que
               return item
             end
 
-            job = job_cache.shift_job(priority)
+            job = job_buffer.shift_job(priority)
             return job unless job.nil? # False means we're stopping.
 
             @waiting += 1
-            @cv.wait
+            @cv.wait(mutex)
             @waiting -= 1
           end
         end
@@ -232,8 +251,8 @@ module Que
 
       private
 
-      def sync
-        @monitor.synchronize { yield }
+      def sync(&block)
+        mutex.synchronize(&block)
       end
     end
   end

data/lib/que/locker.rb

@@ -19,35 +19,17 @@ module Que
     %{
       DELETE FROM public.que_lockers
       WHERE pid = pg_backend_pid()
-      OR pid NOT IN (SELECT pid FROM pg_stat_activity)
+      OR NOT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = public.que_lockers.pid)
     }
 
   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[]
-      )
+      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
+    attr_reader :thread, :workers, :job_buffer, :locks
 
     MESSAGE_RESOLVERS = {}
     RESULT_RESOLVERS  = {}
@@ -55,31 +37,31 @@ module Que
     MESSAGE_RESOLVERS[:job_available] =
       -> (messages) {
         metajobs = messages.map { |key| Metajob.new(key) }
-        push_jobs(lock_jobs(job_cache.accept?(metajobs)))
+        push_jobs(lock_jobs(job_buffer.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
+    DEFAULT_POLL_INTERVAL       = 5.0
+    DEFAULT_WAIT_PERIOD         = 50
+    DEFAULT_MINIMUM_BUFFER_SIZE = 2
+    DEFAULT_MAXIMUM_BUFFER_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
+      queues:              [Que.default_queue],
+      connection_url:      nil,
+      listen:              true,
+      poll:                true,
+      poll_interval:       DEFAULT_POLL_INTERVAL,
+      wait_period:         DEFAULT_WAIT_PERIOD,
+      maximum_buffer_size: DEFAULT_MAXIMUM_BUFFER_SIZE,
+      minimum_buffer_size: DEFAULT_MINIMUM_BUFFER_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
@@ -96,27 +78,29 @@ module Que
 
       all_worker_priorities = worker_priorities.values_at(0...worker_count)
 
-      # We use a JobCache to track jobs and pass them to workers, and a
+      # We use a JobBuffer 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,
+      @job_buffer = JobBuffer.new(
+        maximum_size: maximum_buffer_size,
+        minimum_size: minimum_buffer_size,
         priorities:   all_worker_priorities.uniq,
       )
 
       @result_queue = ResultQueue.new
 
+      @stop = false
+
       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,
+          queues:              queues,
+          listen:              listen,
+          poll:                poll,
+          poll_interval:       poll_interval,
+          wait_period:         wait_period,
+          maximum_buffer_size: maximum_buffer_size,
+          minimum_buffer_size: minimum_buffer_size,
+          worker_count:        worker_count,
+          worker_priorities:   worker_priorities,
         }
       end
 
@@ -135,7 +119,7 @@ module Que
         all_worker_priorities.map do |priority|
           Worker.new(
             priority:       priority,
-            job_cache:      @job_cache,
+            job_buffer:     @job_buffer,
             result_queue:   @result_queue,
             start_callback: on_worker_start,
           )
@@ -144,18 +128,39 @@ module Que
       # 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
+        break if job_buffer.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) }
+      # If we weren't passed a specific connection_url, borrow a connection from
+      # the pool and derive the connection string from it.
+      connection_args =
+        if connection_url
+          uri = URI.parse(connection_url)
+
+          {
+            host:     uri.host,
+            user:     uri.user,
+            password: uri.password,
+            port:     uri.port || 5432,
+            dbname:   uri.path[1..-1],
+          }.merge(Hash[uri.query.split("&").map{|s| s.split('=')}.map{|a,b| [a.to_sym, b]}])
         else
-          Que.pool
+          Que.pool.checkout do |conn|
+            c = conn.wrapped_connection
+
+            {
+              host:     c.host,
+              user:     c.user,
+              password: c.pass,
+              port:     c.port,
+              dbname:   c.db,
+            }
+          end
         end
 
+      @connection = Que::Connection.wrap(PG::Connection.open(connection_args))
+
       @thread =
         Thread.new do
           # An error causing this thread to exit is a bug in Que, which we want
@@ -165,47 +170,35 @@ module Que
           # 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(
+          begin
+            unless connection_args.has_key?(:application_name)
+              @connection.execute(
                 "SELECT set_config('application_name', $1, false)",
-                ["Que Locker: #{connection.backend_pid}"]
+                ["Que Locker: #{@connection.backend_pid}"]
               )
+            end
 
-              Poller.setup(connection)
+            Poller.setup(@connection)
 
+            @listener =
               if listen
-                @listener = Listener.new(connection: connection)
+                Listener.new(connection: @connection)
               end
 
+            @pollers =
               if poll
-                @pollers =
-                  queues.map do |queue, interval|
-                    Poller.new(
-                      connection:    connection,
-                      queue:         queue,
-                      poll_interval: interval || poll_interval,
-                    )
-                  end
+                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
+            work_loop
+          ensure
+            @connection.wrapped_connection.close
           end
         end
     end
@@ -215,7 +208,7 @@ module Que
     end
 
     def stop
-      @job_cache.stop
+      @job_buffer.stop
       @stop = true
     end
 
@@ -249,17 +242,7 @@ module Que
       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('","')}\"}",
-        ]
+        startup
 
         {} while cycle
 
@@ -268,11 +251,7 @@ module Que
           event: :locker_stop,
         )
 
-        unlock_jobs(@job_cache.clear)
-
-        @workers.each(&:wait_until_stopped)
-
-        handle_results
+        shutdown
       ensure
         connection.execute :clean_lockers
 
@@ -280,6 +259,20 @@ module Que
       end
     end
 
+    def startup
+      # 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('","')}\"}",
+      ]
+    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.
@@ -300,31 +293,70 @@ module Que
       !@stop
     end
 
+    def shutdown
+      unlock_jobs(@job_buffer.clear)
+      wait_for_shutdown
+      handle_results
+    end
+
+    def wait_for_shutdown
+      @workers.each(&:wait_until_stopped)
+    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?
+      return unless pollers && job_buffer.jobs_needed?
 
-      pollers.each do |poller|
-        priorities = job_cache.available_priorities
-        break if priorities.empty?
+      # Figure out what job priorities we have to fill.
+      priorities = job_buffer.available_priorities
+      all_metajobs = []
 
-        Que.internal_log(:locker_polling, self) { {priorities: priorities, held_locks: @locks.to_a, queue: poller.queue} }
+      pollers.each do |poller|
+        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.sort!
+          all_metajobs.concat(metajobs)
+
+          # Update the desired priorities list to take the priorities that we
+          # just retrieved into account.
           metajobs.each do |metajob|
-            mark_id_as_locked(metajob.id)
+            job_priority = metajob.job.fetch(:priority)
+
+            priorities.each do |priority, count|
+              if job_priority <= priority
+                new_priority = count - 1
+
+                if new_priority <= 0
+                  priorities.delete(priority)
+                else
+                  priorities[priority] = new_priority
+                end
+
+                break
+              end
+            end
           end
 
-          push_jobs(metajobs)
+          break if priorities.empty?
         end
       end
+
+      all_metajobs.each { |metajob| mark_id_as_locked(metajob.id) }
+      push_jobs(all_metajobs)
     end
 
     def wait
-      if @listener
-        @listener.wait_for_grouped_messages(@wait_period).each do |type, messages|
+      if l = @listener
+        l.wait_for_grouped_messages(@wait_period).each do |type, messages|
           if resolver = MESSAGE_RESOLVERS[type]
             instance_exec messages, &resolver
           else
@@ -353,7 +385,7 @@ module Que
       metajobs.reject! { |m| @locks.include?(m.id) }
       return metajobs if metajobs.empty?
 
-      ids = metajobs.map{|m| m.id.to_i}
+      ids = metajobs.map { |m| m.id.to_i }
 
       Que.internal_log :locker_locking, self do
         {
@@ -365,9 +397,7 @@ module Que
       jobs =
         connection.execute \
           <<-SQL
-            WITH jobs AS (
-              SELECT * FROM que_jobs WHERE id IN (#{ids.join(', ')})
-            )
+            WITH jobs AS (SELECT * FROM que_jobs WHERE id IN (#{ids.join(', ')}))
             SELECT * FROM jobs WHERE pg_try_advisory_lock(id)
           SQL
 
@@ -408,12 +438,12 @@ module Que
 
       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)
+      # Need to unlock any low-importance jobs the new ones may displace.
+      if displaced = @job_buffer.push(*good)
+        bad.concat(displaced)
       end
+
+      unlock_jobs(bad)
     end
 
     def finish_jobs(metajobs)