tom_queue 0.0.1.dev

data/lib/tom_queue.rb

@@ -0,0 +1,56 @@
+# This is a bit of a library-in-a-library, for the time being
+# 
+# This module manages the interaction with AMQP, handling publishing of
+# work messages, scheduling of work off the AMQP queue, etc.
+#
+##
+#
+# You probably want to start with TomQueue::QueueManager
+#
+module TomQueue
+  require 'tom_queue/logging_helper'
+  
+  require 'tom_queue/queue_manager'
+  require 'tom_queue/work'
+  
+  require 'tom_queue/deferred_work_set'
+  require 'tom_queue/deferred_work_manager'
+
+  require 'tom_queue/external_consumer'
+
+  require 'tom_queue/sorted_array'
+
+  # Public: Sets the bunny instance to use for new QueueManager objects
+  def bunny=(new_bunny)
+    @@bunny = new_bunny
+  end
+  # Public: Returns the current bunny instance
+  #
+  # Returns whatever was passed to TomQueue.bunny = 
+  def bunny
+    defined?(@@bunny) && @@bunny
+  end
+  module_function :bunny, :bunny=
+
+  def default_prefix=(new_default_prefix)
+    @@default_prefix = new_default_prefix
+  end
+  def default_prefix
+    defined?(@@default_prefix) && @@default_prefix
+  end
+  module_function :default_prefix=, :default_prefix
+
+
+  # Public: Set an object to receive notifications if an internal exception
+  # is caught and handled.
+  #
+  # IT should be an object that responds to #notify(exception) and should be 
+  # thread safe as reported exceptions will be from background threads crashing.
+  #
+  class << self
+    attr_accessor :exception_reporter
+    attr_accessor :logger
+  end
+
+
+end

data/lib/tom_queue/deferred_work_manager.rb

@@ -0,0 +1,233 @@
+require 'tom_queue/work'
+module TomQueue
+
+  # Internal: This is an internal class that oversees the delay of "deferred"
+  #  work, that is, work with a future :run_at value.
+  #
+  # There is a singleton object for each prefix value (really, there should only be a single
+  # prefix used. The queue manager ensures this thread is running whenever #pop is called.
+  # Work is also pushed to this maanger by the QueueManager when it needs to be deferred.
+  #
+  # Internally, this class opens a separate AMQP channel (without a prefetch limit) and
+  # leaves all the deferred messages in an un-acked state. An internal timer is maintained
+  # to delay until the next message is ready to be sent, at which point the message is
+  # dispatched back to the QueueManager, and at some point will be processed by a worker.
+  #
+  # If the host process of this manager dies for some reason, the broker will re-queue the
+  # un-acked messages onto the deferred queue, to be re-popped by another worker in the pool.
+  #
+  class DeferredWorkManager
+
+    include LoggingHelper
+
+    # Public: Scoped singleton accessor
+    #
+    # This returns the shared work manager instance, creating it if necessary.
+    #
+    # Returns a DeferredWorkManager instance
+    @@singletons_mutex = Mutex.new
+    @@singletons = {}
+    def self.instance(prefix = nil)
+      prefix ||= TomQueue.default_prefix
+      prefix || raise(ArgumentError, 'prefix is required')
+
+      @@singletons_mutex.synchronize { @@singletons[prefix] ||= self.new(prefix) }
+    end
+
+
+    # Public: Return a hash of all prefixed singletons, keyed on their prefix
+    #
+    # This method really is just a convenience method for testing.
+    #
+    # NOTE: The returned hash is both a dupe and frozen, so should be safe to 
+    # iterate and mutate instances.
+    # 
+    # Returns: { "prefix" => DeferredWorkManager(prefix),  ... }
+    def self.instances
+      @@singletons.dup.freeze
+    end
+
+    # Public: Shutdown all managers and wipe the singleton objects.
+    # This method is really just a hook for testing convenience.
+    #
+    # Returns nil
+    def self.reset!
+      @@singletons_mutex.synchronize do
+        @@singletons.each_pair do |k,v|
+          v.ensure_stopped
+        end
+        @@singletons = {}
+      end
+      nil
+    end
+
+    # Public: Return the AMQP prefix used by this manager
+    #
+    # Returns string
+    attr_reader :prefix
+
+    # Internal: Creates the singleton instance, please use the singleton accessor!
+    #
+    # prefix - the AMQP prefix for this instance
+    #
+    def initialize(prefix)
+      @prefix = prefix
+      @thread = nil
+    end
+
+    # Public: Handle a deferred message
+    #
+    # work - (String) the work payload
+    # opts - (Hash) the options of the message. See QueueManager#publish, but must include:
+    #   :run_at = (Time) when the work should be run
+    #
+    def handle_deferred(work, opts)
+      run_at = opts[:run_at]
+      raise ArgumentError, 'work must be a string' unless work.is_a?(String)
+      raise ArgumentError, ':run_at must be specified' if run_at.nil?
+      raise ArgumentError, ':run_at must be a Time object if specified' unless run_at.is_a?(Time)
+
+      debug "[handle_deferred] Pushing work onto deferred exchange: '#{@prefix}.work.deferred'"
+      # Push this work on to the deferred exchange
+      channel = TomQueue.bunny.create_channel
+      exchange, _ = setup_amqp(channel)
+
+      exchange.publish(work, {
+          :mandatory => true,
+          :headers => opts.merge(:run_at => run_at.to_f)
+
+        })
+      channel.close
+    end
+
+    # Public: Return the Thread associated with this manager
+    #
+    # Returns Ruby Thread object, or nil if it's not running
+    attr_reader :thread
+    
+    # Public: Ensure the thread is running, starting if necessary
+    #
+    def ensure_running
+      @thread_shutdown = false
+      @thread = nil unless @thread && @thread.alive?
+      @thread ||= begin
+        debug "[ensure_running] Starting thread"
+        Thread.new(&method(:thread_main))
+      end
+    end
+
+    # Public: Ensure the thread shuts-down and stops. Blocks until
+    # the thread has actually shut down
+    #
+    def ensure_stopped
+      if @thread
+        @thread_shutdown = true
+        @deferred_set && @deferred_set.interrupt
+        @thread.join
+      end
+    end
+
+    # Internal: Creates the bound exchange and queue for deferred work on the provided channel
+    #
+    # channel - a Bunny channel to use to create the queue / exchange binding. If you intend
+    #           to use the returned channel / exchange objects, they will be accessed via
+    #           this channel.
+    #
+    # Returns [ <exchange object>, <queue object> ]
+    def setup_amqp(channel)
+      exchange = channel.fanout("#{prefix}.work.deferred", 
+          :durable     => true,
+          :auto_delete => false)
+
+      queue = channel.queue("#{prefix}.work.deferred",
+          :durable     => true,
+          :auto_delete => false).bind(exchange.name)
+
+      [exchange, queue]
+    end
+
+    ##### Thread Internals #####
+    #
+    # Cross this barrier with care :)
+    #
+
+    # Internal: This is called on a bunny internal work thread when
+    # a new message arrives on the deferred work queue.
+    #
+    # A given message will be delivered only to one deferred manager
+    # so we simply schedule the message in our DeferredWorkSet (which
+    # helpfully deals with all the cross-thread locking, etc.)
+    #
+    # response - the AMQP response object from Bunny
+    # headers  - (Hash) a hash of headers associated with the message
+    # payload  - (String) the message payload
+    #
+    #
+    def thread_consumer_callback(response, headers, payload)
+      run_at = Time.at(headers[:headers]['run_at'])
+
+      # schedule it in the work set
+      @deferred_set.schedule(run_at, [response, headers, payload])
+    rescue Exception => e
+      r = TomQueue.exception_reporter
+      r && r.notify(e)
+      
+      ### Avoid tight spinning workers by not re-queueing redlivered messages more than once!
+      response.channel.reject(response.delivery_tag, !response.redelivered?)
+    end
+
+    # Internal: The main loop of the thread
+    #
+    def thread_main
+      debug "[thread_main] Deferred thread starting up"
+
+      # Make sure we're low priority!
+      Thread.current.priority = -10
+
+      # Create a dedicated channel, and ensure it's prefetch 
+      # means we'll empty the queue
+      @channel = TomQueue.bunny.create_channel
+      @channel.prefetch(0)
+
+      # Create an exchange and queue
+      _, queue = setup_amqp(@channel)
+
+      @deferred_set = DeferredWorkSet.new
+      
+      @out_manager = QueueManager.new(prefix)
+
+      # This block will get called-back for new messages
+      consumer = queue.subscribe(:ack => true, &method(:thread_consumer_callback)) 
+
+      # This is the core event loop - we block on the deferred set to return messages
+      # (which have been scheduled by the AMQP consumer). If a message is returned
+      # then we re-publish the messages to our internal QueueManager and ack the deferred
+      # message
+      until @thread_shutdown
+        
+        # This will block until work is ready to be returned, interrupt
+        # or the 10-second timeout value.
+        response, headers, payload = @deferred_set.pop(2)
+
+        if response
+          headers[:headers].delete('run_at')
+          @out_manager.publish(payload, headers[:headers])
+          @channel.ack(response.delivery_tag)
+        end
+      end
+
+      consumer.cancel
+
+    rescue
+      reporter = TomQueue.exception_reporter
+      reporter && reporter.notify($!)
+
+    ensure
+      @channel && @channel.close
+      @deferred_set = nil
+      @thread = nil      
+    end
+
+  end
+
+end

data/lib/tom_queue/deferred_work_set.rb

@@ -0,0 +1,165 @@
+module TomQueue
+
+  # Internal: This class wraps the pool of work items that are waiting for their run_at 
+  # time to be reached.
+  # 
+  # It also incorporates the logic and coordination required to stall a thread until the 
+  # work is ready to run.
+  #
+  class DeferredWorkSet
+
+    # Internal: A wrapper object to store the run at and the opaque work object inside the @work array.
+    class Element < Struct.new(:run_at, :work)
+      include Comparable
+
+      # Internal: An integer version of run_at, less precise, but faster to compare.
+      attr_reader :fast_run_at
+
+      # Internal: Creates an element. This is called by DeferredWorkSet as work is scheduled so
+      # shouldn't be done directly.
+      #
+      # run_at - (Time) the time when this job should be run
+      # work   - (Object) a payload associated with this work. Just a plain ruby
+      #          object that it is up to the caller to interpret
+      #
+      def initialize(run_at, work)
+        super(run_at, work)
+        @fast_run_at = (run_at.to_f * 1000).to_i
+      end
+
+      # Internal: Comparison function, referencing the scheduled run-time of the element
+      #
+      # NOTE: We don't compare the Time objects directly as this is /dog/ slow, as is comparing
+      # float objects, and this function will be called a /lot/ - so we compare reasonably 
+      # accurate integer values created in the initializer.
+      #
+      def <=> (other)
+        fast_run_at <=> other.fast_run_at
+      end
+
+      # Internal: We need to override this in order for elements with the same run_at not to be deleted
+      # too soon.
+      #
+      # When #<=> is used with `Comparable`, we get #== for free, but when it operates using run_at, it has
+      # the undesirable side-effect that when Array#delete is called with an element, all other elements with 
+      # the same run_at are deleted, too (since #== is used by Array#delete).
+      #
+      def == (other)
+        false
+      end
+    end
+
+    def initialize
+      @mutex = Mutex.new
+      @condvar = ConditionVariable.new
+      @work = TomQueue::SortedArray.new
+    end
+
+    # Public: Returns the integer number of elements in the set
+    #
+    # Returns integer
+    def size
+      @work.size
+    end
+
+    # Public: Block the calling thread until some work is ready to run
+    # or the timeout expires.
+    #
+    # This is intended to be called from a single worker thread, for the
+    # time being, if you try and block on this method concurrently in 
+    # two threads, it will raise an exception!
+    #
+    # timeout - (Fixnum, seconds) how long to wait before timing out
+    #
+    # Returns previously scheduled work, or
+    #         nil if the thread was interrupted or the timeout expired
+    def pop(timeout)
+      timeout_end = Time.now + timeout
+      returned_work = nil
+
+      @interrupt = false
+
+      @mutex.synchronize do
+        raise RuntimeError, 'DeferredWorkSet: another thread is already blocked on a pop' unless @blocked_thread.nil?
+
+        begin
+          @blocked_thread = Thread.current
+
+          begin
+            end_time = [next_run_at, timeout_end].compact.min
+            delay = end_time - Time.now
+            @condvar.wait(@mutex, delay) if delay > 0
+          end while Time.now < end_time and @interrupt == false
+
+          element = earliest_element
+          if element && element.run_at < Time.now
+            @work.delete(element)
+            returned_work = element.work
+          end
+
+        ensure
+          @blocked_thread = nil
+        end
+      end
+
+      returned_work
+    end
+    
+    # Public: Interrupt anything sleeping on this set
+    #
+    # This is "thread-safe" and is designed to be called from threads
+    # to interrupt the work loop thread blocked on a pop.
+    #
+    def interrupt
+      @mutex.synchronize do
+        @interrupt = true
+        @condvar.signal
+      end
+    end
+
+    # Public: Add some work to the set
+    #
+    # This is "threa-safe" in that it can be (and is intended to
+    # be) called from threads other than the one calling pop without
+    # any additional synchronization.
+    #
+    # run_at - (Time) when the work is to be run
+    # work   - the DeferredWork object
+    #
+    def schedule(run_at, work)
+      @mutex.synchronize do
+        new_element = Element.new(run_at, work)
+        @work << new_element
+        @condvar.signal
+      end
+    end
+
+    # Public: Returns the temporally "soonest" element in the set
+    # i.e. the work that is next to be run
+    #
+    # Returns the work Object passed to schedule instance or
+    #         nil if there is no work in the set
+    def earliest
+      e = earliest_element
+      e && e.work
+    end
+
+    # Internal: The next time this thread should next wake up for an element
+    #
+    # If there are elements in the work set, this will correspond to time of the soonest.
+    #
+    # Returns a Time object, or nil if there are no elements stored.
+    def next_run_at
+      e = earliest_element
+      e && e.run_at
+    end
+
+    # Internal: The Element object wrapping the work with the soonest run_at value
+    #
+    # Returns Element object, or nil if the work set is empty
+    def earliest_element
+      @work.first
+    end
+  end
+
+end