davidyang-resque-scheduler 1.10.11

data/lib/resque/scheduler.rb

@@ -0,0 +1,236 @@
+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
+      # TODO: clean up (Removing this check as we want to use both dynamic and 
+      # fixed jobs - davidyang
+      
+      # the Rufus::Scheduler jobs that are scheduled
+      def scheduled_jobs
+        @@scheduled_jobs
+      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
+        procline "Loading Schedule"
+        load_schedule!
+
+        # Now start the scheduling part of the loop.
+        loop do
+          handle_delayed_items
+          update_schedule 
+          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') { kill_child }
+          trap('USR2') { reload_schedule! }
+        rescue ArgumentError
+          warn "Signals QUIT and USR1 and USR2 not supported."
+        end
+      end
+
+      # Pulls the schedule from Resque.schedule and loads it into the
+      # rufus scheduler instance
+      def load_schedule!
+        log! "Schedule empty! Set Resque.schedule" if Resque.schedule.empty?
+        
+        @@scheduled_jobs = {}
+        
+        Resque.schedule.each do |name, config|
+          load_schedule_job(name, config)
+        end
+        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} "
+          if !config['cron'].nil? && config['cron'].length > 0
+            @@scheduled_jobs[name] = rufus_scheduler.cron config['cron'] do
+              log! "queuing #{config['class']} (#{name})"
+              enqueue_from_config(config)
+            end
+          else
+            log! "no cron 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
+      def handle_delayed_items
+        item = nil
+        if timestamp = Resque.next_delayed_timestamp
+          procline "Processing Delayed Items"
+          while !timestamp.nil?
+            enqueue_delayed_items_for_timestamp(timestamp)
+            timestamp = Resque.next_delayed_timestamp
+          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]"
+              klass = constantize(item['class'])
+              queue = item['queue'] || Resque.queue_from_class(klass)
+              # Support custom job classes like job with status
+              if (job_klass = item['custom_job_class']) && (job_klass != 'Resque::Job')
+                # custom job classes not supporting the same API calls must implement the #schedule method
+                constantize(job_klass).scheduled(queue, item['class'], *item['args'])
+              else
+                Resque::Job.create(queue, klass, *item['args'])
+              end
+            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(config)
+        args = config['args'] || config[:args]
+        klass_name = config['class'] || config[:class]
+        klass = constantize(klass_name)
+        params = args.nil? ? [] : Array(args)
+        queue = config['queue'] || config[:queue] || Resque.queue_from_class(klass)
+        # Support custom job classes like job with status
+        if (job_klass = config['custom_job_class']) && (job_klass != 'Resque::Job')
+          # custom job classes not supporting the same API calls must implement the #schedule method
+          constantize(job_klass).scheduled(queue, klass_name, *params)
+        else
+          Resque::Job.create(queue, klass, *params)
+        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!
+        Resque.reload_schedule!
+        load_schedule!
+      end
+      
+      def update_schedule
+        if Resque.needs_updating?
+          procline "Updating schedule"
+          # A bit heavy handed here, but unload everything from Rufus and load in the new schedule
+          # since the Resque.schedule (from redis) will always have the true schedule as setup by the user
+          scheduled_jobs.each do |name, config|
+            unschedule_job(name)
+          end
+          
+          Resque.schedule.each do |name, config|
+            load_schedule_job(name, config)
+          end
+        
+          Resque.mark_schedules_as_updated
+          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 5 }
+        @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)
+        $0 = "resque-scheduler-#{ResqueScheduler::Version}: #{string}"
+        log! $0
+      end
+
+    end
+
+  end
+
+end

data/lib/resque_scheduler.rb

@@ -0,0 +1,206 @@
+require 'rubygems'
+require 'resque'
+require 'resque/server'
+require 'resque_scheduler/version'
+require 'resque/scheduler'
+require 'resque_scheduler/server'
+
+module ResqueScheduler
+
+  #
+  # Convenience method for placing a schedule configuration into the redis
+  # server
+  #
+  # 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"},
+  #    ...}
+  #
+  # :name can be anything and is used only to describe the scheduled job
+  # :cron can be any cron scheduling string :job can be any resque job class
+  # :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)
+    #@schedule = schedule_hash
+
+    # put all the jobs from a YAML file into the schedules hash
+    schedule_hash.each do |name, job_spec|
+      set_schedule(name, job_spec)
+    end
+  end
+
+  # Returns the schedule hash
+  def schedule
+    get_schedules || {}
+  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
+  def set_schedule(name, config)
+    redis.hset(:schedules, name, encode(config))
+    redis.set(:schedules_updated, Time.now.to_s)
+  end
+
+  def needs_updating?
+    redis.get(:schedules_updated) ? true : false
+  end
+
+  def mark_schedules_as_updated
+    redis.del(:schedules_updated)
+  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)
+  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)
+    delayed_push(timestamp, job_to_hash(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
+
+  # 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
+    items = redis.zrangebyscore :delayed_queue_schedule, '-inf', 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 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
+
+end
+
+Resque.extend ResqueScheduler
+Resque::Server.class_eval do
+  include ResqueScheduler::Server
+end

data/lib/resque_scheduler/server.rb

@@ -0,0 +1,63 @@
+
+# Extend Resque::Server to add tabs
+module ResqueScheduler
+  
+  module Server
+
+    def self.included(base)
+
+      base.class_eval do
+
+        helpers do
+          def format_time(t)
+            t.strftime("%Y-%m-%d %H:%M:%S")
+          end
+
+          def queue_from_class_name(class_name)
+            Resque.queue_from_class(Resque.constantize(class_name))
+          end
+        end
+
+        get "/schedule" do
+          Resque.reload_schedule! 
+          # Is there a better way to specify alternate template locations with sinatra?
+          erb File.read(File.join(File.dirname(__FILE__), 'server/views/scheduler.erb'))
+        end
+
+        post "/schedule/delete" do
+          Resque.remove_schedule(params['job_name'])
+          redirect url("/schedule")
+        end
+
+        post "/schedule/requeue" do
+          config = Resque.schedule[params['job_name']]
+          Resque::Scheduler.enqueue_from_config(config)
+          redirect url("/overview")
+        end
+        
+        get "/delayed" do
+          # Is there a better way to specify alternate template locations with sinatra?
+          erb File.read(File.join(File.dirname(__FILE__), 'server/views/delayed.erb'))
+        end
+
+        get "/delayed/:timestamp" do
+          # Is there a better way to specify alternate template locations with sinatra?
+          erb File.read(File.join(File.dirname(__FILE__), 'server/views/delayed_timestamp.erb'))
+        end
+        
+        post "/delayed/queue_now" do
+          timestamp = params['timestamp']
+          Resque::Scheduler.enqueue_delayed_items_for_timestamp(timestamp.to_i) if timestamp.to_i > 0
+          redirect url("/overview")
+        end
+
+      end
+
+    end
+
+    Resque::Server.tabs << 'Schedule'
+    Resque::Server.tabs << 'Delayed'
+
+  end
+  
+end