resque-plugins-zookeeper-lock.rb

module Resque
  module Plugins
    # If you want only one instance of your job queued at a time,
    # extend it with this module.
    #
    # For example:
    #
    # require 'resque/plugins/zookeeper_lock'
    #
    # class UpdateNetworkGraph
    #   extend Resque::Plugins::ZookeeperLock
    #
    #   def self.perform(repo_id)
    #     heavy_lifting
    #   end
    # end
    #
    # No other UpdateNetworkGraph jobs will be placed on the queue,
    # the QueueLock class will check Redis to see if any others are
    # queued with the same arguments before queueing. If another
    # is queued the enqueue will be aborted.
    #
    # If you want to define the key yourself you can override the
    # `lock` class method in your subclass, e.g.
    #
    # class UpdateNetworkGraph
    #   extend Resque::Plugins::ZookeeperLock
    #
    #   # Run only one at a time, regardless of repo_id.
    #   def self.lock(repo_id)
    #     "network-graph"
    #   end
    #
    #   def self.perform(repo_id)
    #     heavy_lifting
    #   end
    # end
    #
    # The above modification will ensure only one job of class
    # UpdateNetworkGraph is running at a time, regardless of the
    # repo_id. Normally a job is locked using a combination of its
    # class name and arguments.
    module ZookeeperLock
      # Assign your ZK instance to 
      # Resque::Plugins::ZookeeperLock.zookeeper
      def self.zookeeper=(zk)
        @zk = zk
      end

      def self.zookeeper
        @zk
      end

      # Override in your job to control where the ZK class comes
      # from
      def zookeeper
        ZookeeperLock.zookeeper
      end

      # Override in your job to control the lock key. It is
      # passed the same arguments as `perform`, that is, your job's
      # payload.
      def lock(*args)
        "resque-#{name}-#{args.map(&:to_s).join('-')}"
      end

      def before_perform_zookeeper_lock(*args)
        @zk_lock = zookeeper.exclusive_locker(lock(*args))
        unless @zk_lock.lock
          raise Resque::Job::DontPerform
        end
      end

      def after_perform_zookeeper_lock(*args)
        # Always clear the lock when we're done, even if there is an
        # error.
        if @zk_lock
          @zk_lock.unlock rescue nil
          @zk_lock = nil
        end
      end
    end
  end
end