qu-scheduler 0.0.1 → 0.1.0

data/lib/qu/scheduler.rb

@@ -0,0 +1,227 @@
+require 'qu'
+require 'qu/scheduler/version'
+require 'qu/extensions/scheduler'
+require 'rufus/scheduler'
+require 'thwait'
+
+module Qu
+  class Scheduler
+    @@scheduled_jobs = {}
+
+    class << self
+      # If set, will try to update the schedule 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
+        set_process_title "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, USR1 and USR2 not supported."
+        end
+      end
+
+      def print_schedule
+        if rufus_scheduler
+          Qu.logger.info("Scheduling Info\tLast Run")
+          scheduler_jobs = rufus_scheduler.all_jobs
+          scheduler_jobs.each do |k, v|
+            Qu.logger.info("#{v.t}\t#{v.last}")
+          end
+        end
+      end
+
+      # Pulls the schedule from Qu.schedule and loads it into the
+      # rufus scheduler instance
+      def load_schedule!
+        Qu.backend.load_schedule!
+      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)
+          Qu.logger.info("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
+              @@scheduled_jobs[name] = rufus_scheduler.send(interval_type, config[interval_type]) do
+                Qu.logger.info("queueing #{config['class']} (#{name})")
+                handle_errors { enqueue_from_config(config) }
+              end
+              interval_defined = true
+              break
+            end
+          end
+          unless interval_defined
+            Qu.logger.warn("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 = Qu.backend.next_delayed_timestamp(at_time)
+          set_process_title "processing delayed jobs"
+          while !timestamp.nil?
+            enqueue_delayed_items_for_timestamp(timestamp)
+            timestamp = Qu.backend.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 = Qu.backend.next_item_for_timestamp(timestamp)
+              Qu.logger.debug("queuing #{item['klass']} [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
+        exit if @shutdown
+        yield
+        exit if @shutdown
+      end
+
+      def handle_errors
+        begin
+          yield
+        rescue Exception => e
+          Qu.logger.fatal("#{e.class.name}: #{e.message}")
+        end
+      end
+
+      # Enqueues a job based on a config hash
+      def enqueue_from_config(job_config)
+        payload = Payload.new(job_config)
+        if job_klass = job_config['custom_job_class']
+          # 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 Qu::Job.create.
+          begin
+            Qu::Payload.new.send(:constantize, job_klass).scheduled(payload.queue, job_klass, *payload.args)
+          rescue NameError
+            # Note that the custom job class (job_config['custom_job_class']) is the one enqueued
+            Qu.backend.enqueue(Payload.new(:klass => job_klass, :args => payload.args))
+          end
+        else
+          Qu.backend.enqueue(payload)
+        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!
+        set_process_title "reloading schedule"
+        clear_schedule!
+        load_schedule!
+      end
+
+      def unschedule_job(name)
+        if scheduled_jobs[name]
+          Qu.logger.debug("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 set_process_title(string)
+        Qu.logger.info(string)
+        $0 = "qu-scheduler-#{Qu::Scheduler::VERSION}: #{string}"
+      end
+    end
+  end
+end

data/lib/qu/scheduler/version.rb

@@ -1,5 +1,5 @@
 module Qu
-  module Scheduler
-    VERSION = "0.0.1"
+  class Scheduler
+    VERSION = "0.1.0"
   end
 end

data/test/delayed_queue_test.rb

@@ -0,0 +1,277 @@
+require 'test_helper'
+
+context "DelayedQueue" do
+
+  setup do
+    Qu.backend.redis.flushall
+  end
+
+  test "enqueue_at adds correct list and zset" do
+    timestamp = Time.now - 1 # 1 second ago (in the past, should come out right away)
+
+    assert_equal(0, Qu.backend.redis.llen("delayed:#{timestamp.to_i}").to_i, "delayed queue should be empty to start")
+
+    Qu.enqueue_at(timestamp, SomeIvarJob, "path")
+
+    # Confirm the correct keys were added
+    assert_equal(1, Qu.backend.redis.llen("delayed:#{timestamp.to_i}").to_i, "delayed queue should have one entry now")
+    assert_equal(1, Qu.backend.redis.zcard(:delayed_queue_schedule), "The delayed_queue_schedule should have 1 entry now")
+
+    read_timestamp = Qu.backend.next_delayed_timestamp
+
+    # Confirm the timestamp came out correctly
+    assert_equal(timestamp.to_i, read_timestamp, "The timestamp we pull out of redis should match the one we put in")
+    item = Qu.backend.next_item_for_timestamp(read_timestamp)
+
+    # Confirm the item came out correctly
+    assert_equal('SomeIvarJob', item['klass'], "Should be the same class that we queued")
+    assert_equal(["path"], item['args'], "Should have the same arguments that we queued")
+
+    # And now confirm the keys are gone
+    assert(!Qu.backend.redis.exists("delayed:#{timestamp.to_i}"))
+    assert_equal(0, Qu.backend.redis.zcard(:delayed_queue_schedule), "delayed queue should be empty")
+  end
+
+  test "a job in the future doesn't come out" do
+    timestamp = Time.now + 600 # 10 minutes from now (in the future, shouldn't come out)
+
+    assert_equal(0, Qu.backend.redis.llen("delayed:#{timestamp.to_i}").to_i, "delayed queue should be empty to start")
+
+    Qu.enqueue_at(timestamp, SomeIvarJob, "path")
+
+    # Confirm the correct keys were added
+    assert_equal(1, Qu.backend.redis.llen("delayed:#{timestamp.to_i}").to_i, "delayed queue should have one entry now")
+    assert_equal(1, Qu.backend.redis.zcard(:delayed_queue_schedule), "The delayed_queue_schedule should have 1 entry now")
+
+    read_timestamp = Qu.backend.next_delayed_timestamp
+
+    assert_nil(read_timestamp, "No timestamps should be ready for queueing")
+  end
+
+  test "a job in the future comes out if you want it to" do
+    timestamp = Time.now + 600 # 10 minutes from now
+
+    Qu.enqueue_at(timestamp, SomeIvarJob, "path")
+
+    read_timestamp = Qu.backend.next_delayed_timestamp(timestamp)
+
+    assert_equal(timestamp.to_i, read_timestamp, "The timestamp we pull out of redis should match the one we put in")
+  end
+
+  test "enqueue_at and enqueue_in are equivelent" do
+    timestamp = Time.now + 60
+
+    Qu.enqueue_at(timestamp, SomeIvarJob, "path")
+    Qu.enqueue_in(timestamp - Time.now, SomeIvarJob, "path")
+
+    assert_equal(1, Qu.backend.redis.zcard(:delayed_queue_schedule), "should have one timestamp in the delayed queue")
+    assert_equal(2, Qu.backend.redis.llen("delayed:#{timestamp.to_i}"), "should have 2 items in the timestamp queue")
+  end
+
+  test "empty delayed_queue_peek returns empty array" do
+    assert_equal([], Qu.backend.delayed_queue_peek(0,20))
+  end
+
+  test "delayed_queue_peek returns stuff" do
+    t = Time.now
+    expected_timestamps = (1..5).to_a.map do |i|
+      (t + 60 + i).to_i
+    end
+
+    expected_timestamps.each do |timestamp|
+      Qu.backend.delayed_push(timestamp, Qu::Payload.new({:class => SomeIvarJob, :args => 'blah1'}))
+    end
+
+    timestamps = Qu.backend.delayed_queue_peek(2,3)
+
+    assert_equal(expected_timestamps[2,3], timestamps)
+  end
+
+  test "delayed_queue_schedule_size returns correct size" do
+    assert_equal(0, Qu.backend.delayed_queue_schedule_size)
+    Qu.enqueue_at(Time.now+60, SomeIvarJob)
+    assert_equal(1, Qu.backend.delayed_queue_schedule_size)
+  end
+
+  test "delayed_timestamp_size returns 0 when nothing is queue" do
+    t = Time.now + 60
+    assert_equal(0, Qu.backend.delayed_timestamp_size(t))
+  end
+
+  test "delayed_timestamp_size returns 1 when one thing is queued" do
+    t = Time.now + 60
+    Qu.enqueue_at(t, SomeIvarJob)
+    assert_equal(1, Qu.backend.delayed_timestamp_size(t))
+  end
+
+  test "delayed_timestamp_peek returns empty array when nothings in it" do
+    t = Time.now + 60
+    assert_equal([], Qu.delayed_timestamp_peek(t, 0, 1), "make sure it's an empty array, not nil")
+  end
+
+  test "delayed_timestamp_peek returns an array containing one job when one thing is queued" do
+    t = Time.now + 60
+    Qu.enqueue_at(t, SomeIvarJob)
+    assert_equal [{'args' => [], 'klass' => 'SomeIvarJob'}], Qu.delayed_timestamp_peek(t, 0, 1)
+  end
+
+  test "delayed_timestamp_peek returns an array of multiple jobs when more than one job is queued" do
+    t = Time.now + 60
+    Qu.enqueue_at(t, SomeIvarJob)
+    Qu.enqueue_at(t, SomeIvarJob)
+    job = {'args' => [], 'klass' => 'SomeIvarJob'}
+    assert_equal([job, job], Qu.delayed_timestamp_peek(t, 0, 2))
+  end
+
+  test "delayed_timestamp_peek only returns an array of one job if only asked for 1" do
+    t = Time.now + 60
+    Qu.enqueue_at(t, SomeIvarJob)
+    Qu.enqueue_at(t, SomeIvarJob)
+    job = {'args' => [], 'klass' => 'SomeIvarJob'}
+    assert_equal([job], Qu.delayed_timestamp_peek(t, 0, 1))
+  end
+
+  test "handle_delayed_items with no items" do
+    Qu::Scheduler.expects(:enqueue).never
+    Qu::Scheduler.handle_delayed_items
+  end
+
+  test "handle_delayed_item with items" do
+    t = Time.now - 60 # in the past
+    Qu.enqueue_at(t, SomeIvarJob)
+    Qu.enqueue_at(t, SomeIvarJob)
+
+    # 2 SomeIvarJob jobs should be created in the "ivar" queue
+    Qu.backend.expects(:enqueue).twice.with(all_of(instance_of(Qu::Payload), responds_with(:klass, SomeIvarJob), responds_with(:args, []), responds_with(:queue, 'ivar')))
+    Qu::Scheduler.handle_delayed_items
+  end
+
+  test "handle_delayed_items with items in the future" do
+    t = Time.now + 60 # in the future
+    Qu.enqueue_at(t, SomeIvarJob)
+    Qu.enqueue_at(t, SomeIvarJob)
+
+    # 2 SomeIvarJob jobs should be created in the "ivar" queue
+    Qu.backend.expects(:enqueue).twice.with(all_of(instance_of(Qu::Payload), responds_with(:klass, SomeIvarJob), responds_with(:args, []), responds_with(:queue, 'ivar')))
+    Qu::Scheduler.handle_delayed_items(t)
+  end
+
+  test "enqueue_delayed_items_for_timestamp creates jobs and empties the delayed queue" do
+    t = Time.now + 60
+
+    Qu.enqueue_at(t, SomeIvarJob)
+    Qu.enqueue_at(t, SomeIvarJob)
+
+    # 2 SomeIvarJob jobs should be created in the "ivar" queue
+    Qu.backend.expects(:enqueue).twice.with(all_of(instance_of(Qu::Payload), responds_with(:klass, SomeIvarJob), responds_with(:args, []), responds_with(:queue, 'ivar')))
+
+    Qu::Scheduler.enqueue_delayed_items_for_timestamp(t)
+
+    # delayed queue for timestamp should be empty
+    assert_equal(0, Qu.delayed_timestamp_peek(t, 0, 3).length)
+  end
+
+  test "handle_delayed_items works with out specifying queue (upgrade case)" do
+    t = Time.now - 60
+    Qu.backend.delayed_push(t, Qu::Payload.new(:klass => 'SomeIvarJob'))
+
+    # Since we didn't specify :queue when calling delayed_push, it will be forced
+    # to load the class to figure out the queue.  This is the upgrade case from 1.0.4
+    # to 1.0.5.
+    Qu.backend.expects(:enqueue).once.with(all_of(instance_of(Qu::Payload), responds_with(:klass, SomeIvarJob), responds_with(:args, []), responds_with(:queue, 'ivar')))
+
+    Qu::Scheduler.handle_delayed_items
+  end
+
+  test "reset_delayed_queue clears the queue" do
+    t = Time.now + 120
+    4.times { Qu.enqueue_at(t, SomeIvarJob) }
+    4.times { Qu.enqueue_at(Time.now + rand(100), SomeIvarJob) }
+
+    Qu.backend.reset_delayed_queue
+    assert_equal(0, Qu.backend.delayed_queue_schedule_size)
+  end
+
+  test "remove_delayed removes job and returns the count" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob)
+
+    assert_equal(1, Qu.backend.remove_delayed(SomeIvarJob))
+  end
+
+  test "remove_delayed doesn't remove things it shouldn't" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    Qu.enqueue_at(t, SomeIvarJob, "bar")
+    Qu.enqueue_at(t, SomeIvarJob, "bar")
+    Qu.enqueue_at(t, SomeIvarJob, "baz")
+
+    assert_equal(0, Qu.backend.remove_delayed(SomeIvarJob))
+  end
+
+  test "remove_delayed respected param" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    Qu.enqueue_at(t, SomeIvarJob, "bar")
+    Qu.enqueue_at(t, SomeIvarJob, "bar")
+    Qu.enqueue_at(t, SomeIvarJob, "baz")
+
+    assert_equal(2, Qu.backend.remove_delayed(SomeIvarJob, "bar"))
+    assert_equal(1, Qu.backend.delayed_queue_schedule_size)
+  end
+
+  test "remove_delayed removes items in different timestamps" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    Qu.enqueue_at(t + 1, SomeIvarJob, "bar")
+    Qu.enqueue_at(t + 2, SomeIvarJob, "bar")
+    Qu.enqueue_at(t + 3, SomeIvarJob, "baz")
+
+    assert_equal(2, Qu.backend.remove_delayed(SomeIvarJob, "bar"))
+    assert_equal(2, Qu.backend.count_all_scheduled_jobs)
+  end
+
+  test "remove_delayed_job_from_timestamp removes instances of jobs at a given timestamp" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    assert_equal 1, Qu.backend.remove_delayed_job_from_timestamp(t, SomeIvarJob, "foo")
+    assert_equal 0, Qu.backend.delayed_timestamp_size(t)
+  end
+
+  test "remove_delayed_job_from_timestamp doesn't remove items from other timestamps" do
+    t1 = Time.now + 120
+    t2 = t1 + 1
+    Qu.enqueue_at(t1, SomeIvarJob, "foo")
+    Qu.enqueue_at(t2, SomeIvarJob, "foo")
+    assert_equal 1, Qu.backend.remove_delayed_job_from_timestamp(t2, SomeIvarJob, "foo")
+    assert_equal 1, Qu.backend.delayed_timestamp_size(t1)
+    assert_equal 0, Qu.backend.delayed_timestamp_size(t2)
+  end
+
+  test "remove_delayed_job_from_timestamp removes nothing if there are no matches" do
+    t = Time.now + 120
+    assert_equal 0, Qu.backend.remove_delayed_job_from_timestamp(t, SomeIvarJob, "foo")
+  end
+
+  test "remove_delayed_job_from_timestamp only removes items that match args" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    Qu.enqueue_at(t, SomeIvarJob, "bar")
+    assert_equal 1, Qu.backend.remove_delayed_job_from_timestamp(t, SomeIvarJob, "foo")
+    assert_equal 1, Qu.backend.delayed_timestamp_size(t)
+  end
+
+  test "remove_delayed_job_from_timestamp returns the number of items removed" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    assert_equal 1, Qu.backend.remove_delayed_job_from_timestamp(t, SomeIvarJob, "foo")
+  end
+
+  test "remove_delayed_job_from_timestamp should cleanup the delayed timestamp list if not jobs are left" do
+    t = Time.now + 120
+    Qu.enqueue_at(t, SomeIvarJob, "foo")
+    assert_equal 1, Qu.backend.remove_delayed_job_from_timestamp(t, SomeIvarJob, "foo")
+    assert !Qu.backend.redis.exists("delayed:#{t.to_i}")
+    assert Qu.backend.delayed_queue_peek(0, 100).empty?
+  end
+end