ealdent-resque-scheduler 2.0.0.e

data/Rakefile

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

data/lib/resque/scheduler.rb

@@ -0,0 +1,285 @@
+require 'rufus/scheduler'
+require 'thwait'
+
+module Resque
+
+  class Scheduler
+
+    extend Resque::Helpers
+
+    class << self
+
+      # 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
+
+        # Load the schedule into rufus
+        # If dynamic is set, load that schedule otherwise use normal load
+        if dynamic
+          reload_schedule!
+        else
+          load_schedule!
+        end
+
+        # Now start the scheduling part of the loop.
+        loop do
+          begin
+            handle_delayed_items
+            update_schedule if dynamic
+          rescue Errno::EAGAIN, Errno::ECONNRESET => e
+            warn e.message
+          end
+          poll_sleep
+        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
+
+      # 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
+              begin
+                @@scheduled_jobs[name] = rufus_scheduler.send(interval_type, config[interval_type]) do
+                  log! "queueing #{config['class']} (#{name})"
+                  enqueue_from_config(config)
+                end
+              rescue Exception => e
+                log! "#{e.class.name}: #{e.message}"
+              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)
+        item = 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]"
+              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
+        exit if @shutdown
+        yield
+        exit if @shutdown
+      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
+            Resque.enqueue(klass, *params)
+          else
+            Resque::Job.create(queue, klass, *params)
+          end
+        end
+      rescue
+        log! "Failed to enqueue #{klass_name}:\n #{$!}"
+      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
+
+    end
+
+  end
+
+end

data/lib/resque_scheduler.rb

@@ -0,0 +1,248 @@
+require 'rubygems'
+require 'resque'
+require 'resque/server'
+require 'resque_scheduler/version'
+require 'resque/scheduler'
+require 'resque_scheduler/server'
+
+module ResqueScheduler
+
+  #
+  # Accepts a new schedule configuration of the form:
+  #
+  #   {'some_name' => {"cron" => "5/* * * *",
+  #                  "class" => "DoSomeWork",
+  #                  "args" => "work on this string",
+  #                  "description" => "this thing works it"s butter off"},
+  #    ...}
+  #
+  # 'some_name' can be anything and is used only to describe and reference 
+  # the scheduled job
+  #
+  # :cron can be any cron scheduling string :job can be any resque job class
+  #
+  # :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
+  #
+  # :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)
+    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)
+    delayed_push(timestamp, job_to_hash(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.
+  def enqueue_at_with_queue(queue, timestamp, klass, *args)
+    delayed_push(timestamp, job_to_hash_with_queue(queue, 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)).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
+  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
+
+  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.
+      if 0 == redis.llen(key).to_i
+        redis.del key
+        redis.zrem :delayed_queue_schedule, timestamp.to_i
+      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
+
+end
+
+Resque.extend ResqueScheduler
+Resque::Server.class_eval do
+  include ResqueScheduler::Server
+end