que 0.0.1 → 0.1.0

data/lib/generators/que/install_generator.rb

@@ -0,0 +1,22 @@
+require 'rails/generators'
+require 'rails/generators/migration'
+require 'active_record'
+
+module Que
+  class InstallGenerator < Rails::Generators::Base
+    include Rails::Generators::Migration
+
+    namespace "que:install"
+    self.source_paths << File.join(File.dirname(__FILE__), 'templates')
+    desc "Generates a migration to add Que's job table."
+
+    def self.next_migration_number(dirname)
+      next_migration_number = current_migration_number(dirname) + 1
+      ActiveRecord::Migration.next_migration_number(next_migration_number)
+    end
+
+    def create_migration_file
+      migration_template 'add_que.rb', 'db/migrate/add_que.rb'
+    end
+  end
+end

data/lib/generators/que/templates/add_que.rb

@@ -0,0 +1,9 @@
+class AddQue < ActiveRecord::Migration
+  def self.up
+    Que.create!
+  end
+
+  def self.down
+    Que.drop!
+  end
+end

data/lib/que.rb

@@ -1,10 +1,60 @@
-require 'que/version'
-
 module Que
-  autoload :Job,    'que/job'
-  autoload :Worker, 'que/worker'
+  autoload :Adapters, 'que/adapters/base'
+  autoload :Job,      'que/job'
+  autoload :SQL,      'que/sql'
+  autoload :Version,  'que/version'
+  autoload :Worker,   'que/worker'
 
   class << self
-    attr_accessor :logger
+    attr_accessor :logger, :error_handler
+    attr_writer :adapter
+
+    def adapter
+      @adapter || raise("Que connection not established!")
+    end
+
+    def connection=(connection)
+      self.adapter = if connection.to_s == 'ActiveRecord'
+        Adapters::ActiveRecord.new
+      else
+        case connection.class.to_s
+          when 'Sequel::Postgres::Database' then Adapters::Sequel.new(connection)
+          when 'ConnectionPool'             then Adapters::ConnectionPool.new(connection)
+          when 'PG::Connection'             then Adapters::PG.new(connection)
+          when 'NilClass'                   then connection
+          else raise "Que connection not recognized: #{connection.inspect}"
+        end
+      end
+    end
+
+    def create!
+      execute SQL[:create_table]
+    end
+
+    def drop!
+      execute "DROP TABLE que_jobs"
+    end
+
+    def clear!
+      execute "DELETE FROM que_jobs"
+    end
+
+    def execute(command, *args)
+      case command
+        when Symbol then adapter.execute_prepared(command, *args)
+        when String then adapter.execute(command, *args)
+      end
+    end
+
+    def log(level, text)
+      logger.send level, "[Que] #{text}" if logger
+    end
+
+    # Duplicate some Worker config methods to the Que module for convenience.
+    [:mode, :mode=, :worker_count=, :sleep_period, :sleep_period=].each do |meth|
+      define_method(meth){|*args| Worker.send(meth, *args)}
+    end
   end
 end
+
+require 'que/railtie' if defined? Rails::Railtie

data/lib/que/adapters/active_record.rb

@@ -0,0 +1,9 @@
+module Que
+  module Adapters
+    class ActiveRecord < Base
+      def checkout
+        ::ActiveRecord::Base.connection_pool.with_connection { |conn| yield conn.raw_connection }
+      end
+    end
+  end
+end

data/lib/que/adapters/base.rb

@@ -0,0 +1,49 @@
+module Que
+  module Adapters
+    autoload :ActiveRecord,   'que/adapters/active_record'
+    autoload :ConnectionPool, 'que/adapters/connection_pool'
+    autoload :PG,             'que/adapters/pg'
+    autoload :Sequel,         'que/adapters/sequel'
+
+    class Base
+      def initialize(thing = nil)
+        @statement_mutex = Mutex.new
+      end
+
+      # The only method that adapters really need to implement. Should lock a
+      # PG::Connection (or something that acts like a PG::Connection) so that
+      # no other threads are using it and yield it to the block.
+      def checkout(&block)
+        raise NotImplementedError
+      end
+
+      def execute(*args)
+        checkout { |conn| conn.async_exec(*args) }
+      end
+
+      def execute_prepared(name, params = [])
+        checkout do |conn|
+          unless statements_prepared(conn)[name]
+            conn.prepare("que_#{name}", SQL[name])
+            statements_prepared(conn)[name] = true
+          end
+
+          conn.exec_prepared("que_#{name}", params)
+        end
+      end
+
+      private
+
+      # Each adapter needs to remember which of its connections have prepared
+      # which statements. This is a shared data structure, so protect it. We
+      # assume that the hash of statements for a particular connection is only
+      # being accessed by the thread that's checked it out, though.
+      def statements_prepared(conn)
+        @statement_mutex.synchronize do
+          @statements_prepared       ||= {}
+          @statements_prepared[conn] ||= {}
+        end
+      end
+    end
+  end
+end

data/lib/que/adapters/connection_pool.rb

@@ -0,0 +1,14 @@
+module Que
+  module Adapters
+    class ConnectionPool < Base
+      def initialize(pool)
+        @pool = pool
+        super
+      end
+
+      def checkout(&block)
+        @pool.with(&block)
+      end
+    end
+  end
+end

data/lib/que/adapters/pg.rb

@@ -0,0 +1,17 @@
+require 'monitor'
+
+module Que
+  module Adapters
+    class PG < Base
+      def initialize(pg)
+        @pg   = pg
+        @lock = Monitor.new # Must be re-entrant.
+        super
+      end
+
+      def checkout
+        @lock.synchronize { yield @pg }
+      end
+    end
+  end
+end

data/lib/que/adapters/sequel.rb

@@ -0,0 +1,14 @@
+module Que
+  module Adapters
+    class Sequel < Base
+      def initialize(db)
+        @db = db
+        super
+      end
+
+      def checkout(&block)
+        @db.synchronize(&block)
+      end
+    end
+  end
+end

data/lib/que/job.rb

@@ -1,185 +1,164 @@
 require 'json'
 
 module Que
-  class Job < Sequel::Model
-    # The Job priority scale:
-    #   1 = Urgent. Somebody's staring at a spinner waiting on this.
-    #   2 = ASAP. Should happen within a few minutes of the run_at time.
-    #   3 = Time-sensitive. Sooner is better than later.
-    #   4 = Time-insensitive. Shouldn't get delayed forever, though.
-    #   5 = Whenever. Timing doesn't matter. May be a huge backlog of these.
+  class Job
+    def initialize(attrs)
+      @attrs = attrs
+    end
 
-    unrestrict_primary_key
+    # Subclasses should define their own run methods, but keep an empty one
+    # here so we can just do Que::Job.queue in testing.
+    def run(*args)
+    end
 
-    plugin :single_table_inheritance, :type, :key_map   => proc(&:to_s),
-                                             :model_map => proc{|s| const_get "::#{s}"}
+    def _run
+      start = Time.now
 
-    class Retry < StandardError; end
+      run *@attrs[:args]
+      destroy unless @destroyed
 
-    class << self
-      # Default is lowest priority, meaning jobs can be done whenever.
-      def default_priority
-        @default_priority ||= 5
-      end
+      Que.log :info, "Worked job in #{((Time.now - start) * 1000).round(1)} ms: #{inspect}"
+    end
 
+    private
+
+    def destroy
+      Que.execute :destroy_job, [@attrs[:priority], @attrs[:run_at], @attrs[:job_id]]
+      @destroyed = true
+    end
+
+    class << self
       def queue(*args)
-        create values_for_args *args
+        if args.last.is_a?(Hash)
+          options  = args.pop
+          run_at   = options.delete(:run_at)
+          priority = options.delete(:priority)
+          args << options if options.any?
+        end
+
+        attrs = {:job_class => to_s, :args => JSON.dump(args)}
+
+        if t = run_at || @default_run_at && @default_run_at.call
+          attrs[:run_at] = t
+        end
+
+        if p = priority || @default_priority
+          attrs[:priority] = p
+        end
+
+        if Que.mode == :sync
+          run_job(attrs)
+        else
+          Que.execute *insert_sql(attrs)
+        end
       end
 
-      def work(options = {})
+      def work
+        # Job.work will typically be called in a loop, where we'd sleep when
+        # there's no more work to be done, so its return value should reflect
+        # whether we should hit the database again or not. So, return truthy
+        # if we worked a job or encountered a typical error while working a
+        # job, and falsy if we found nothing to do or hit a connection error.
+
         # Since we're taking session-level advisory locks, we have to hold the
         # same connection throughout the process of getting a job, working it,
         # deleting it, and removing the lock.
-
-        DB.synchronize do
+        Que.adapter.checkout do
           begin
-            return unless job = LOCK.call(:priority => options[:priority] || 5)
-
-            # Edge case: It's possible for the lock statement to have grabbed a
-            # job that's already been worked, if the statement took its MVCC
-            # snapshot while the job was processing (making it appear to still
-            # exist), but didn't actually attempt to lock it until the job was
-            # finished (making it appear to be unlocked). Now that we have the
-            # job lock, we know that a previous worker would have deleted it by
-            # now, so we just make sure it still exists before working it.
-            this = dataset.where(:priority => job[:priority], :run_at => job[:run_at], :job_id => job[:job_id])
-            return if this.empty?
-
-            # Split up model instantiation from the DB query, so that model
-            # instantiation errors can be caught.
-            model = sti_load(job)
-
-            # Track how long different jobs take to process.
-            start = Time.now
-            model.work
-            time = Time.now - start
-            Que.logger.info "Worked job in #{(time * 1000).round(1)} ms: #{model.inspect}" if Que.logger
-
-            # Most jobs destroy themselves transactionally in #work. If not,
-            # take care of them. Jobs that don't destroy themselves run the risk
-            # of being repeated after a crash.
-            model.destroy unless model.destroyed?
-
-            # Make sure to return the finished job.
-            model
-          rescue Retry
-            # Don't destroy the job or mark it as having errored. It can be
-            # retried as soon as it is unlocked.
+            if row = Que.execute(:lock_job).first
+              # Edge case: It's possible to have grabbed a job that's already
+              # been worked, if the SELECT took its MVCC snapshot while the
+              # job was processing, but didn't attempt the advisory lock until
+              # it was finished. Now that we have the job lock, we know that a
+              # previous worker would have deleted it by now, so we just
+              # double check that it still exists before working it.
+
+              # Note that there is currently no spec for this behavior, since
+              # I'm not sure how to reliably commit a transaction that deletes
+              # the job in a separate thread between this lock and check.
+              return true if Que.execute(:check_job, [row['priority'], row['run_at'], row['job_id']]).none?
+
+              run_job(row)
+            else
+              Que.log :info, "No jobs available..."
+              nil
+            end
           rescue => error
-            if job && data = JSON.load(job[:data])
-              count = (data['error_count'] || 0) + 1
+            begin
+              if row
+                # Borrowed the exponential backoff formula and error data format from delayed_job.
+                count   = row['error_count'].to_i + 1
+                run_at  = Time.now + (count ** 4 + 3)
+                message = "#{error.message}\n#{error.backtrace.join("\n")}"
+                Que.execute :set_error, [count, run_at, message, row['priority'], row['run_at'], row['job_id']]
+              end
+            rescue
+              # If we can't reach the DB for some reason, too bad, but don't
+              # let it crash the work loop.
+            end
 
-              this.update :run_at => Time.now + (count ** 4 + 3),
-                          :data   => JSON.dump(:error_count => count, :error_message => error.message, :error_backtrace => error.backtrace.join("\n"))
+            if Que.error_handler
+              Que.error_handler.call(error) rescue nil
             end
 
-            raise
+            # If it's a garden variety error, we can just return true, pick up
+            # another job, no big deal. If it's a PG::Error, though, assume
+            # it's a disconnection or something and that we shouldn't just hit
+            # the database again right away.
+            return !error.is_a?(PG::Error)
           ensure
-            DB.get{pg_advisory_unlock(job[:job_id])} if job
+            # Clear the advisory lock we took when locking the job. Important
+            # to do this so that they don't pile up in the database.
+            Que.execute "SELECT pg_advisory_unlock_all()" if row
           end
         end
       end
 
       private
 
-      def values_for_args(*args)
-        opts = args.last.is_a?(Hash) ? args.pop : {}
-
-        result = {}
-        result[:run_at]   = opts.delete(:run_at)   if opts[:run_at]
-        result[:priority] = opts.delete(:priority) if opts[:priority]
-
-        args << opts if opts.any?
-        result[:args] = JSON.dump(args)
+      # Column names are not escaped, so this method should not be called with untrusted hashes.
+      def insert_sql(hash)
+        number       = 0
+        columns      = []
+        placeholders = []
+        values       = []
+
+        hash.each do |key, value|
+          columns      << key
+          placeholders << "$#{number += 1}"
+          values       << value
+        end
 
-        result
+        ["INSERT INTO que_jobs (#{columns.join(', ')}) VALUES (#{placeholders.join(', ')})", values]
       end
-    end
-
-    # Send the args attribute to the perform() method.
-    def work
-      perform(*JSON.parse(args))
-    end
-
-    # Call perform on a job to run it. No perform method means NOOP.
-    def perform(*args)
-    end
 
-    def destroyed?
-      !!@destroyed
-    end
-
-    private
+      def run_job(attrs)
+        attrs = indifferentiate(attrs)
+        attrs[:args] = indifferentiate(JSON.load(attrs[:args]))
+        const_get("::#{attrs[:job_class]}").new(attrs).tap(&:_run)
+      end
 
-    # If we add any more callbacks here, make sure to also special-case them in
-    # queue_array above.
-    def before_create
-      self.priority ||= self.class.default_priority
-
-      # If there's no run_at set, the job needs to be run immediately, so we
-      # need to trigger a worker to work it after it's committed and visible.
-      if run_at.nil?
-        case Worker.state
-          when :sync  then DB.after_commit { Job.work     }
-          when :async then DB.after_commit { Worker.wake! }
+      def indifferentiate(input)
+        case input
+        when Hash
+          h = indifferent_hash
+          input.each { |k, v| h[k] = indifferentiate(v) }
+          h
+        when Array
+          input.map { |v| indifferentiate(v) }
+        else
+          input
         end
       end
 
-      super
-    end
-
-    def after_destroy
-      super
-      @destroyed = true
+      def indifferent_hash
+        # Tiny hack to better support Rails.
+        if {}.respond_to?(:with_indifferent_access)
+          {}.with_indifferent_access
+        else
+          Hash.new { |hash, key| hash[key.to_s] if Symbol === key }
+        end
+      end
     end
-
-    sql = <<-SQL
-      WITH RECURSIVE cte AS (
-        SELECT (job).*, pg_try_advisory_lock((job).job_id) AS locked
-        FROM (
-          SELECT job
-          FROM jobs AS job
-          WHERE ((run_at <= now()) AND (priority <= ?))
-          ORDER BY priority, run_at, job_id
-          LIMIT 1
-        ) AS t1
-        UNION ALL (
-        SELECT (job).*, pg_try_advisory_lock((job).job_id) AS locked
-        FROM (
-          SELECT (
-            SELECT job
-            FROM jobs AS job
-            WHERE ((run_at <= now()) AND (priority <= ?) AND ((priority, run_at, job_id) > (cte.priority, cte.run_at, cte.job_id)))
-            ORDER BY priority, run_at, job_id
-            LIMIT 1
-          ) AS job
-          FROM cte
-          WHERE NOT cte.locked
-          LIMIT 1
-        ) AS t1)
-      )
-      SELECT *
-      FROM cte
-      WHERE locked
-    SQL
-
-    LOCK = DB[sql, :$priority, :$priority].prepare(:first, :lock_job)
-
-    # An alternate scheme using LATERAL, which will arrive in Postgres 9.3.
-    # Basically the same, but benchmark to see if it's faster/just as reliable.
-
-    # with recursive
-    #  t as (select *, pg_try_advisory_lock(s.job_id) as locked
-    #          from (select * from jobs j
-    #                 where run_at >= now()
-    #                 order by priority, run_at, job_id limit 1) s
-    #        union all
-    #        select j.*, pg_try_advisory_lock(j.job_id)
-    #          from (select * from t where not locked) t,
-    #               lateral (select * from jobs
-    #                         where run_at >= now()
-    #                           and (priority,run_at,job_id) > (t.priority,t.run_at,t.job_id)
-    #                         order by priority, run_at, job_id limit 1) j
-    # select * from t where locked;
   end
 end