nogara-resque-scheduler 2.0.1

data/Rakefile

@@ -0,0 +1,28 @@
+require 'bundler'
+
+Bundler::GemHelper.install_tasks
+
+$LOAD_PATH.unshift 'lib'
+
+task :default => :test
+
+# Tests
+desc "Run tests"
+task :test do
+  Dir['test/*_test.rb'].each do |f|
+    require File.expand_path(f)
+  end
+end
+
+# Documentation Tasks
+begin
+  require 'rdoc/task'
+
+  Rake::RDocTask.new do |rd|
+    rd.main = "README.markdown"
+    rd.rdoc_files.include("README.markdown", "lib/**/*.rb")
+    rd.rdoc_dir = 'doc'
+  end
+rescue LoadError
+end
+

data/lib/resque/scheduler.rb

@@ -0,0 +1,362 @@
+require 'rufus/scheduler'
+require 'thwait'
+
+module Resque
+
+  class Scheduler
+
+    extend Resque::Helpers
+
+    class << self
+
+      LOCK_TIMEOUT = 60 * 5
+
+      # If true, logs more stuff...
+      attr_accessor :verbose
+
+      # If set, produces no output
+      attr_accessor :mute
+
+      # If set, will try to update the schulde in the loop
+      attr_accessor :dynamic
+
+      # Amount of time in seconds to sleep between polls of the delayed
+      # queue.  Defaults to 5
+      attr_writer :poll_sleep_amount
+
+      # the Rufus::Scheduler jobs that are scheduled
+      def scheduled_jobs
+        @@scheduled_jobs
+      end
+
+      def poll_sleep_amount
+        @poll_sleep_amount ||= 5 # seconds
+      end
+
+      # Schedule all jobs and continually look for delayed jobs (never returns)
+      def run
+        $0 = "resque-scheduler: Starting"
+        # trap signals
+        register_signal_handlers
+
+        loop do
+          got_lock = can_lock_scheduler?
+          if got_lock == true
+
+            # Load the schedule into rufus
+            # If dynamic is set, load that schedule otherwise use normal load
+            if dynamic
+              reload_schedule!
+            else
+              load_schedule!
+            end
+
+            first_time = false
+
+            # Now start the scheduling part of the loop.
+
+            30.times do #30 * 5 seconds, it should be less than the timeout defined above
+            # loop do
+              begin
+                handle_delayed_items
+                update_schedule if dynamic
+              rescue Errno::EAGAIN, Errno::ECONNRESET => e
+                warn e.message
+              end
+              poll_sleep
+            end
+
+            unlock_scheduler
+            clear_schedule!
+
+          else
+            puts "Scheduler locked!!!"
+            sleep 5
+          end
+        end
+        # never gets here.
+      end
+
+      # For all signals, set the shutdown flag and wait for current
+      # poll/enqueing to finish (should be almost istant).  In the
+      # case of sleeping, exit immediately.
+      def register_signal_handlers
+        trap("TERM") { shutdown }
+        trap("INT") { shutdown }
+
+        begin
+          trap('QUIT') { shutdown }
+          trap('USR1') { print_schedule }
+          trap('USR2') { reload_schedule! }
+        rescue ArgumentError
+          warn "Signals QUIT and USR1 and USR2 not supported."
+        end
+      end
+
+      def print_schedule
+        if rufus_scheduler
+          log! "Scheduling Info\tLast Run"
+          scheduler_jobs = rufus_scheduler.all_jobs
+          scheduler_jobs.each do |k, v|
+            log! "#{v.t}\t#{v.last}\t"
+          end
+        end
+      end
+
+      # Pulls the schedule from Resque.schedule and loads it into the
+      # rufus scheduler instance
+      def load_schedule!
+        procline "Loading Schedule"
+
+        # Need to load the schedule from redis for the first time if dynamic
+        Resque.reload_schedule! if dynamic
+
+        log! "Schedule empty! Set Resque.schedule" if Resque.schedule.empty?
+
+        @@scheduled_jobs = {}
+
+        Resque.schedule.each do |name, config|
+          load_schedule_job(name, config)
+        end
+        Resque.redis.del(:schedules_changed)
+        procline "Schedules Loaded"
+      end
+
+      # modify interval type value to value with options if options available
+      def optionizate_interval_value(value)
+        args = value
+        if args.is_a?(::Array)
+          return args.first if args.size > 2 || !args.last.is_a?(::Hash)
+          # symbolize keys of hash for options
+          args[1] = args[1].inject({}) do |m, i|
+            key, value = i
+            m[(key.to_sym rescue key) || key] = value
+            m
+          end
+        end
+        args
+      end
+
+      # Loads a job schedule into the Rufus::Scheduler and stores it in @@scheduled_jobs
+      def load_schedule_job(name, config)
+        # If rails_env is set in the config, enforce ENV['RAILS_ENV'] as
+        # required for the jobs to be scheduled.  If rails_env is missing, the
+        # job should be scheduled regardless of what ENV['RAILS_ENV'] is set
+        # to.
+        if config['rails_env'].nil? || rails_env_matches?(config)
+          log! "Scheduling #{name} "
+          interval_defined = false
+          interval_types = %w{cron every}
+          interval_types.each do |interval_type|
+            if !config[interval_type].nil? && config[interval_type].length > 0
+              args = optionizate_interval_value(config[interval_type])
+              @@scheduled_jobs[name] = rufus_scheduler.send(interval_type, *args) do
+                log! "queueing #{config['class']} (#{name})"
+                handle_errors { enqueue_from_config(config) }
+              end
+              interval_defined = true
+              break
+            end
+          end
+          unless interval_defined
+            log! "no #{interval_types.join(' / ')} found for #{config['class']} (#{name}) - skipping"
+          end
+        end
+      end
+
+      # Returns true if the given schedule config hash matches the current
+      # ENV['RAILS_ENV']
+      def rails_env_matches?(config)
+        config['rails_env'] && ENV['RAILS_ENV'] && config['rails_env'].gsub(/\s/,'').split(',').include?(ENV['RAILS_ENV'])
+      end
+
+      # Handles queueing delayed items
+      # at_time - Time to start scheduling items (default: now).
+      def handle_delayed_items(at_time=nil)
+        if timestamp = Resque.next_delayed_timestamp(at_time)
+          procline "Processing Delayed Items"
+          while !timestamp.nil?
+            enqueue_delayed_items_for_timestamp(timestamp)
+            timestamp = Resque.next_delayed_timestamp(at_time)
+          end
+        end
+      end
+
+      # Enqueues all delayed jobs for a timestamp
+      def enqueue_delayed_items_for_timestamp(timestamp)
+        item = nil
+        begin
+          handle_shutdown do
+            if item = Resque.next_item_for_timestamp(timestamp)
+              log "queuing #{item['class']} [delayed]"
+              handle_errors { enqueue_from_config(item) }
+            end
+          end
+        # continue processing until there are no more ready items in this timestamp
+        end while !item.nil?
+      end
+
+      def handle_shutdown
+        begin
+          unlock_scheduler if @shutdown
+        rescue
+        end
+        exit if @shutdown
+        yield
+        begin
+          unlock_scheduler if @shutdown
+        rescue
+        end
+        exit if @shutdown
+      end
+
+      def handle_errors
+        begin
+          yield
+        rescue Exception => e
+          log! "#{e.class.name}: #{e.message}"
+        end
+      end
+
+      # Enqueues a job based on a config hash
+      def enqueue_from_config(job_config)
+        args = job_config['args'] || job_config[:args]
+
+        klass_name = job_config['class'] || job_config[:class]
+        klass = constantize(klass_name) rescue klass_name
+
+        params = args.is_a?(Hash) ? [args] : Array(args)
+        queue = job_config['queue'] || job_config[:queue] || Resque.queue_from_class(klass)
+        # Support custom job classes like those that inherit from Resque::JobWithStatus (resque-status)
+        if (job_klass = job_config['custom_job_class']) && (job_klass != 'Resque::Job')
+          # The custom job class API must offer a static "scheduled" method. If the custom
+          # job class can not be constantized (via a requeue call from the web perhaps), fall
+          # back to enqueing normally via Resque::Job.create.
+          begin
+            constantize(job_klass).scheduled(queue, klass_name, *params)
+          rescue NameError
+            # Note that the custom job class (job_config['custom_job_class']) is the one enqueued
+            Resque::Job.create(queue, job_klass, *params)
+          end
+        else
+          # hack to avoid havoc for people shoving stuff into queues
+          # for non-existent classes (for example: running scheduler in
+          # one app that schedules for another
+          if Class === klass
+            ResqueScheduler::Plugin.run_before_delayed_enqueue_hooks(klass, *params)
+            Resque.enqueue_to(queue, klass, *params)
+          else
+            # This will not run the before_hooks in rescue, but will at least
+            # queue the job.
+            Resque::Job.create(queue, klass, *params)
+          end
+        end
+      end
+
+      def rufus_scheduler
+        @rufus_scheduler ||= Rufus::Scheduler.start_new
+      end
+
+      # Stops old rufus scheduler and creates a new one.  Returns the new
+      # rufus scheduler
+      def clear_schedule!
+        rufus_scheduler.stop
+        @rufus_scheduler = nil
+        @@scheduled_jobs = {}
+        rufus_scheduler
+      end
+
+      def reload_schedule!
+        procline "Reloading Schedule"
+        clear_schedule!
+        load_schedule!
+      end
+
+      def update_schedule
+        if Resque.redis.scard(:schedules_changed) > 0
+          procline "Updating schedule"
+          Resque.reload_schedule!
+          while schedule_name = Resque.redis.spop(:schedules_changed)
+            if Resque.schedule.keys.include?(schedule_name)
+              unschedule_job(schedule_name)
+              load_schedule_job(schedule_name, Resque.schedule[schedule_name])
+            else
+              unschedule_job(schedule_name)
+            end
+          end
+          procline "Schedules Loaded"
+        end
+      end
+
+      def unschedule_job(name)
+        if scheduled_jobs[name]
+          log "Removing schedule #{name}"
+          scheduled_jobs[name].unschedule
+          @@scheduled_jobs.delete(name)
+        end
+      end
+
+      # Sleeps and returns true
+      def poll_sleep
+        @sleeping = true
+        handle_shutdown { sleep poll_sleep_amount }
+        @sleeping = false
+        true
+      end
+
+      # Sets the shutdown flag, exits if sleeping
+      def shutdown
+        @shutdown = true
+        exit if @sleeping
+      end
+
+      def log!(msg)
+        puts "#{Time.now.strftime("%Y-%m-%d %H:%M:%S")} #{msg}" unless mute
+      end
+
+      def log(msg)
+        # add "verbose" logic later
+        log!(msg) if verbose
+      end
+
+      def procline(string)
+        log! string
+        $0 = "resque-scheduler-#{ResqueScheduler::VERSION}: #{string}"
+      end
+
+      def lock_timeout
+        Time.now.utc.to_i + LOCK_TIMEOUT + 1
+      end
+
+      def can_lock_scheduler?
+        #using logic from http://redis.io/commands/getset
+        got_lock = Resque.redis.setnx('scheduler:lock', lock_timeout)
+        puts "First get lock #{got_lock}"
+        unless got_lock
+          timestamp = Resque.redis.get('scheduler:lock').to_i
+          puts "Timestamp: #{timestamp}"
+          timestamp_now = Time.now.utc.to_i
+          puts "Timestamp Now: #{timestamp_now}"
+          if timestamp_now > timestamp
+            timestamp_old = Resque.redis.getset('scheduler:lock', lock_timeout).to_i
+           puts "Timestamp Old: #{timestamp_old}"
+            if timestamp_old < timestamp_now
+              puts "Got lock here"
+              got_lock = true
+            end
+          end
+        end
+        puts "Second get lock #{got_lock}"
+        got_lock
+      end
+
+      def unlock_scheduler
+        puts "Unlocking scheduler lock"
+        Resque.redis.del('scheduler:lock')
+      end
+
+    end
+
+  end
+
+end

data/lib/resque_scheduler.rb

@@ -0,0 +1,293 @@
+require 'rubygems'
+require 'resque'
+require 'resque_scheduler/version'
+require 'resque/scheduler'
+require 'resque_scheduler/plugin'
+
+module ResqueScheduler
+
+  #
+  # Accepts a new schedule configuration of the form:
+  #
+  #   {
+  #     "MakeTea" => {
+  #       "every" => "1m" },
+  #     "some_name" => {
+  #       "cron"        => "5/* * * *",
+  #       "class"       => "DoSomeWork",
+  #       "args"        => "work on this string",
+  #       "description" => "this thing works it"s butter off" },
+  #     ...
+  #   }
+  #
+  # Hash keys can be anything and are used to describe and reference
+  # the scheduled job. If the "class" argument is missing, the key
+  # is used implicitly as "class" argument - in the "MakeTea" example,
+  # "MakeTea" is used both as job name and resque worker class.
+  #
+  # :cron can be any cron scheduling string
+  #
+  # :every can be used in lieu of :cron. see rufus-scheduler's 'every' usage
+  # for valid syntax. If :cron is present it will take precedence over :every.
+  #
+  # :class must be a resque worker class. If it is missing, the job name (hash key)
+  # will be used as :class.
+  #
+  # :args can be any yaml which will be converted to a ruby literal and
+  # passed in a params. (optional)
+  #
+  # :rails_envs is the list of envs where the job gets loaded. Envs are
+  # comma separated (optional)
+  #
+  # :description is just that, a description of the job (optional). If
+  # params is an array, each element in the array is passed as a separate
+  # param, otherwise params is passed in as the only parameter to perform.
+  def schedule=(schedule_hash)
+    schedule_hash = prepare_schedule(schedule_hash)
+
+    if Resque::Scheduler.dynamic
+      schedule_hash.each do |name, job_spec|
+        set_schedule(name, job_spec)
+      end
+    end
+    @schedule = schedule_hash
+  end
+
+  # Returns the schedule hash
+  def schedule
+    @schedule ||= {}
+  end
+
+  # reloads the schedule from redis
+  def reload_schedule!
+    @schedule = get_schedules
+  end
+
+  # gets the schedule as it exists in redis
+  def get_schedules
+    if redis.exists(:schedules)
+      redis.hgetall(:schedules).tap do |h|
+        h.each do |name, config|
+          h[name] = decode(config)
+        end
+      end
+    else
+      nil
+    end
+  end
+
+  # Create or update a schedule with the provided name and configuration.
+  #
+  # Note: values for class and custom_job_class need to be strings,
+  # not constants.
+  #
+  #    Resque.set_schedule('some_job', {:class => 'SomeJob',
+  #                                     :every => '15mins',
+  #                                     :queue => 'high',
+  #                                     :args => '/tmp/poop'})
+  def set_schedule(name, config)
+    existing_config = get_schedule(name)
+    unless existing_config && existing_config == config
+      redis.hset(:schedules, name, encode(config))
+      redis.sadd(:schedules_changed, name)
+    end
+    config
+  end
+
+  # retrive the schedule configuration for the given name
+  def get_schedule(name)
+    decode(redis.hget(:schedules, name))
+  end
+
+  # remove a given schedule by name
+  def remove_schedule(name)
+    redis.hdel(:schedules, name)
+    redis.sadd(:schedules_changed, name)
+  end
+
+  # This method is nearly identical to +enqueue+ only it also
+  # takes a timestamp which will be used to schedule the job
+  # for queueing.  Until timestamp is in the past, the job will
+  # sit in the schedule list.
+  def enqueue_at(timestamp, klass, *args)
+    validate_job!(klass)
+    enqueue_at_with_queue(queue_from_class(klass), timestamp, klass, *args)
+  end
+
+  # Identical to +enqueue_at+, except you can also specify
+  # a queue in which the job will be placed after the
+  # timestamp has passed. It respects Resque.inline option, by
+  # creating the job right away instead of adding to the queue.
+  def enqueue_at_with_queue(queue, timestamp, klass, *args)
+    return false unless Plugin.run_before_schedule_hooks(klass, *args)
+
+    if Resque.inline?
+      # Just create the job and let resque perform it right away with inline.
+      Resque::Job.create(queue, klass, *args)
+    else
+      delayed_push(timestamp, job_to_hash_with_queue(queue, klass, args))
+    end
+
+    Plugin.run_after_schedule_hooks(klass, *args)
+  end
+
+  # Identical to enqueue_at but takes number_of_seconds_from_now
+  # instead of a timestamp.
+  def enqueue_in(number_of_seconds_from_now, klass, *args)
+    enqueue_at(Time.now + number_of_seconds_from_now, klass, *args)
+  end
+
+  # Identical to +enqueue_in+, except you can also specify
+  # a queue in which the job will be placed after the
+  # number of seconds has passed.
+  def enqueue_in_with_queue(queue, number_of_seconds_from_now, klass, *args)
+    enqueue_at_with_queue(queue, Time.now + number_of_seconds_from_now, klass, *args)
+  end
+
+  # Used internally to stuff the item into the schedule sorted list.
+  # +timestamp+ can be either in seconds or a datetime object
+  # Insertion if O(log(n)).
+  # Returns true if it's the first job to be scheduled at that time, else false
+  def delayed_push(timestamp, item)
+    # First add this item to the list for this timestamp
+    redis.rpush("delayed:#{timestamp.to_i}", encode(item))
+
+    # Now, add this timestamp to the zsets.  The score and the value are
+    # the same since we'll be querying by timestamp, and we don't have
+    # anything else to store.
+    redis.zadd :delayed_queue_schedule, timestamp.to_i, timestamp.to_i
+  end
+
+  # Returns an array of timestamps based on start and count
+  def delayed_queue_peek(start, count)
+    Array(redis.zrange(:delayed_queue_schedule, start, start+count-1)).collect{|x| x.to_i}
+  end
+
+  # Returns the size of the delayed queue schedule
+  def delayed_queue_schedule_size
+    redis.zcard :delayed_queue_schedule
+  end
+
+  # Returns the number of jobs for a given timestamp in the delayed queue schedule
+  def delayed_timestamp_size(timestamp)
+    redis.llen("delayed:#{timestamp.to_i}").to_i
+  end
+
+  # Returns an array of delayed items for the given timestamp
+  def delayed_timestamp_peek(timestamp, start, count)
+    if 1 == count
+      r = list_range "delayed:#{timestamp.to_i}", start, count
+      r.nil? ? [] : [r]
+    else
+      list_range "delayed:#{timestamp.to_i}", start, count
+    end
+  end
+
+  # Returns the next delayed queue timestamp
+  # (don't call directly)
+  def next_delayed_timestamp(at_time=nil)
+    items = redis.zrangebyscore :delayed_queue_schedule, '-inf', (at_time || Time.now).to_i, :limit => [0, 1]
+    timestamp = items.nil? ? nil : Array(items).first
+    timestamp.to_i unless timestamp.nil?
+  end
+
+  # Returns the next item to be processed for a given timestamp, nil if
+  # done. (don't call directly)
+  # +timestamp+ can either be in seconds or a datetime
+  def next_item_for_timestamp(timestamp)
+    key = "delayed:#{timestamp.to_i}"
+
+    item = decode redis.lpop(key)
+
+    # If the list is empty, remove it.
+    clean_up_timestamp(key, timestamp)
+    item
+  end
+
+  # Clears all jobs created with enqueue_at or enqueue_in
+  def reset_delayed_queue
+    Array(redis.zrange(:delayed_queue_schedule, 0, -1)).each do |item|
+      redis.del "delayed:#{item}"
+    end
+
+    redis.del :delayed_queue_schedule
+  end
+
+  # Given an encoded item, remove it from the delayed_queue
+  #
+  # This method is potentially very expensive since it needs to scan
+  # through the delayed queue for every timestamp.
+  def remove_delayed(klass, *args)
+    destroyed = 0
+    search = encode(job_to_hash(klass, args))
+    Array(redis.keys("delayed:*")).each do |key|
+      destroyed += redis.lrem key, 0, search
+    end
+    destroyed
+  end
+
+  # Given a timestamp and job (klass + args) it removes all instances and
+  # returns the count of jobs removed.
+  #
+  # O(N) where N is the number of jobs scheduled to fire at the given
+  # timestamp
+  def remove_delayed_job_from_timestamp(timestamp, klass, *args)
+    key = "delayed:#{timestamp.to_i}"
+    count = redis.lrem key, 0, encode(job_to_hash(klass, args))
+    clean_up_timestamp(key, timestamp)
+    count
+  end
+
+  def count_all_scheduled_jobs
+    total_jobs = 0
+    Array(redis.zrange(:delayed_queue_schedule, 0, -1)).each do |timestamp|
+      total_jobs += redis.llen("delayed:#{timestamp}").to_i
+    end
+    total_jobs
+  end
+
+  private
+
+    def job_to_hash(klass, args)
+      {:class => klass.to_s, :args => args, :queue => queue_from_class(klass)}
+    end
+
+    def job_to_hash_with_queue(queue, klass, args)
+      {:class => klass.to_s, :args => args, :queue => queue}
+    end
+
+    def clean_up_timestamp(key, timestamp)
+      # If the list is empty, remove it.
+      redis.watch key
+      if 0 == redis.llen(key).to_i
+        redis.multi do
+          redis.del key
+          redis.zrem :delayed_queue_schedule, timestamp.to_i
+        end
+      else
+        redis.unwatch
+      end
+    end
+    def validate_job!(klass)
+      if klass.to_s.empty?
+        raise Resque::NoClassError.new("Jobs must be given a class.")
+      end
+
+      unless queue_from_class(klass)
+        raise Resque::NoQueueError.new("Jobs must be placed onto a queue.")
+      end
+    end
+
+    def prepare_schedule(schedule_hash)
+      prepared_hash = {}
+      schedule_hash.each do |name, job_spec|
+        job_spec = job_spec.dup
+        job_spec['class'] = name unless job_spec.key?('class') || job_spec.key?(:class)
+        prepared_hash[name] = job_spec
+      end
+      prepared_hash
+    end
+
+end
+
+Resque.extend ResqueScheduler