zk 0.6.4

data/lib/z_k/exceptions.rb

@@ -0,0 +1,101 @@
+module ZK
+  module Exceptions
+    OK                      = 0
+    # System and server-side errors
+    SYSTEMERROR             = -1
+    RUNTIMEINCONSISTENCY    = SYSTEMERROR - 1
+    DATAINCONSISTENCY       = SYSTEMERROR - 2
+    CONNECTIONLOSS          = SYSTEMERROR - 3
+    MARSHALLINGERROR        = SYSTEMERROR - 4
+    UNIMPLEMENTED           = SYSTEMERROR - 5
+    OPERATIONTIMEOUT        = SYSTEMERROR - 6
+    BADARGUMENTS            = SYSTEMERROR - 7
+    # API errors  
+    APIERROR                = -100; 
+    NONODE                  = APIERROR - 1 # Node does not exist
+    NOAUTH                  = APIERROR - 2 # Current operation not permitted
+    BADVERSION              = APIERROR - 3 # Version conflict
+    NOCHILDRENFOREPHEMERALS = APIERROR - 8
+    NODEEXISTS              = APIERROR - 10
+    NOTEMPTY                = APIERROR - 11
+    SESSIONEXPIRED          = APIERROR - 12
+    INVALIDCALLBACK         = APIERROR - 13
+    INVALIDACL              = APIERROR - 14
+    AUTHFAILED              = APIERROR - 15 # client authentication failed
+
+
+    # these errors are returned rather than the driver level errors
+    class KeeperException         < StandardError
+      def self.recognized_code?(code)
+        ERROR_MAP.include?(code)
+      end
+        
+      def self.by_code(code)
+        ERROR_MAP.fetch(code.to_i) { raise "API ERROR: no exception defined for code: #{code}" }
+      end
+    end
+
+    class SystemError             < KeeperException; end
+    class RunTimeInconsistency    < KeeperException; end
+    class DataInconsistency       < KeeperException; end
+    class ConnectionLoss          < KeeperException; end
+    class MarshallingError        < KeeperException; end
+    class Unimplemented           < KeeperException; end
+    class OperationTimeOut        < KeeperException; end
+    class BadArguments            < KeeperException; end
+    class ApiError                < KeeperException; end
+    class NoNode                  < KeeperException; end
+    class NoAuth                  < KeeperException; end
+    class BadVersion              < KeeperException; end
+    class NoChildrenForEphemerals < KeeperException; end
+    class NodeExists              < KeeperException; end
+    class NotEmpty                < KeeperException; end
+    class SessionExpired          < KeeperException; end
+    class InvalidCallback         < KeeperException; end
+    class InvalidACL              < KeeperException; end
+    class AuthFailed              < KeeperException; end
+
+    ERROR_MAP = {
+      SYSTEMERROR             => SystemError,
+      RUNTIMEINCONSISTENCY    => RunTimeInconsistency,
+      DATAINCONSISTENCY       => DataInconsistency,
+      CONNECTIONLOSS          => ConnectionLoss,
+      MARSHALLINGERROR        => MarshallingError,
+      UNIMPLEMENTED           => Unimplemented,
+      OPERATIONTIMEOUT        => OperationTimeOut,
+      BADARGUMENTS            => BadArguments,
+      APIERROR                => ApiError,
+      NONODE                  => NoNode,
+      NOAUTH                  => NoAuth,
+      BADVERSION              => BadVersion,
+      NOCHILDRENFOREPHEMERALS => NoChildrenForEphemerals,
+      NODEEXISTS              => NodeExists,
+      NOTEMPTY                => NotEmpty,
+      SESSIONEXPIRED          => SessionExpired,
+      INVALIDCALLBACK         => InvalidCallback,
+      INVALIDACL              => InvalidACL,
+      AUTHFAILED              => AuthFailed,
+    }
+
+    # base class of ZK generated errors (not driver-level errors)
+    class ZKError < StandardError; end
+
+    class LockFileNameParseError < ZKError; end
+
+    # raised when you try to vote twice in a given leader election
+    class ThisIsNotChicagoError < ZKError; end
+    
+    # raised when close_all! has been called on a pool and some thread attempts a checkout
+    class PoolIsShuttingDownException < ZKError; end
+
+    # raised when defer is called on a threadpool that is not running
+    class ThreadpoolIsNotRunningException < ZKError; end
+
+    # raised when assert_locked_for_update! is called and no exclusive lock is held
+    class MustBeExclusivelyLockedException < ZKError; end
+
+    # raised when assert_locked_for_share! is called and no shared lock is held
+    class MustBeShareLockedException < ZKError; end
+  end
+end
+

data/lib/z_k/extensions.rb

@@ -0,0 +1,144 @@
+module ZK
+  module Extensions
+    # some extensions to the ZookeeperCallbacks classes, mainly convenience
+    # interrogators
+    module Callbacks
+      module Callback
+        # allow access to the connection that fired this callback
+        attr_accessor :zk
+
+        def self.included(mod)
+          mod.extend(ZK::Extensions::Callbacks::Callback::ClassMethods)
+        end
+
+        module ClassMethods
+          # allows for easier construction of a user callback block that will be
+          # called with the callback object itself as an argument. 
+          #
+          # *args, if given, will be passed on *after* the callback
+          #
+          # example:
+          #   
+          #   WatcherCallback.create do |cb|
+          #     puts "watcher callback called with argument: #{cb.inspect}"
+          #   end
+          #
+          #   "watcher callback called with argument: #<ZookeeperCallbacks::WatcherCallback:0x1018a3958 @state=3, @type=1, ...>"
+          #
+          #
+          def create(*args, &block)
+            # honestly, i have no idea how this could *possibly* work, but it does...
+            cb_inst = new { block.call(cb_inst) }
+          end
+        end
+      end
+
+      module WatcherCallbackExt
+        include ZookeeperConstants
+
+        EVENT_NAME_MAP = {
+          1   => 'created',
+          2   => 'deleted', 
+          3   => 'changed',
+          4   => 'child',
+          -1  => 'session',
+          -2  => 'notwatching',
+        }.freeze
+
+        STATES = %w[connecting associating connected auth_failed expired_session].freeze unless defined?(STATES)
+
+        EVENT_TYPES = %w[created deleted changed child session notwatching].freeze unless defined?(EVENT_TYPES)
+
+        STATES.each do |state|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def state_#{state}?
+              @state == ZOO_#{state.upcase}_STATE
+            end
+          RUBY
+        end
+
+        EVENT_TYPES.each do |ev|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def node_#{ev}?
+              @type == ZOO_#{ev.upcase}_EVENT
+            end
+          RUBY
+        end
+
+        alias :node_not_watching? :node_notwatching?
+
+        # has this watcher been called because of a change in connection state?
+        def state_event?
+          path.nil? or path.empty?
+        end
+
+        # has this watcher been called because of a change to a zookeeper node?
+        def node_event?
+          path and not path.empty?
+        end
+
+        # cause this watch to be re-registered
+#         def renew_watch!
+#           zk.stat(path, :watch => true)
+#           nil
+#         end
+      end
+    end   # Callbacks
+
+    # aliases for long-names of properties from mb-zookeeper version
+    module Stat
+      [ %w[created_zxid czxid],
+        %w[last_modified_zxid mzxid],
+        %w[created_time ctime],
+        %w[last_modified_time mtime],
+        %w[child_list_version cversion],
+        %w[acl_list_version aversion] ].each do |long, short|
+
+        class_eval <<-RUBY, __FILE__, __LINE__ + 1
+          def #{long}
+            #{short}
+          end
+        RUBY
+      end
+
+      MEMBERS = [:version, :exists, :czxid, :mzxid, :ctime, :mtime, :cversion, :aversion, :ephemeralOwner, :dataLength, :numChildren, :pzxid]
+
+      def self.included(mod)
+        mod.class_eval do
+          unless method_defined?(:exists?)
+            alias :exists? :exists
+          end
+        end
+      end
+
+      def ==(other)
+        MEMBERS.all? { |m| self.__send__(m) == other.__send__(m) }
+      end
+    end
+
+  end     # Extensions
+end       # ZK
+
+# ZookeeperCallbacks::Callback.extend(ZK::Extensions::Callbacks::Callback)
+ZookeeperCallbacks::Callback.send(:include, ZK::Extensions::Callbacks::Callback)
+ZookeeperCallbacks::WatcherCallback.send(:include, ZK::Extensions::Callbacks::WatcherCallbackExt)
+ZookeeperStat::Stat.send(:include, ZK::Extensions::Stat)
+
+class ::Exception
+  unless method_defined?(:to_std_format)
+    def to_std_format
+      "#{self.class}: #{message}\n\t" + backtrace.join("\n\t")
+    end
+  end
+end
+
+class ::Thread
+  def zk_mongoid_lock_registry
+    self[:_zk_mongoid_lock_registry]
+  end
+
+  def zk_mongoid_lock_registry=(obj)
+    self[:_zk_mongoid_lock_registry] = obj
+  end
+end
+

data/lib/z_k/locker.rb

@@ -0,0 +1,254 @@
+module ZK
+  # Implements locking primitives {described here}[http://hadoop.apache.org/zookeeper/docs/current/recipes.html#sc_recipes_Locks]
+  #
+  # There are both shared and exclusive lock implementations.
+  #
+  #
+  # NOTE: These locks are _not_ safe for use across threads. If you want to use
+  # the same Locker class between threads, it is your responsibility to
+  # synchronize operations.
+  #
+  module Locker
+    SHARED_LOCK_PREFIX  = 'sh'.freeze
+    EXCLUSIVE_LOCK_PREFIX = 'ex'.freeze
+
+    def self.shared_locker(zk, name)
+      SharedLocker.new(zk, name)
+    end
+
+    def self.exclusive_locker(zk, name)
+      ExclusiveLocker.new(zk, name)
+    end
+    
+    class NoWriteLockFoundException < StandardError #:nodoc:
+    end
+
+    class WeAreTheLowestLockNumberException < StandardError #:nodoc:
+    end
+
+    class LockerBase
+      include ZK::Logging
+
+      attr_accessor :zk #:nodoc:
+
+      # our absolute lock node path
+      #
+      # ex. '/_zklocking/foobar/__blah/lock000000007'
+      attr_reader :lock_path #;nodoc:
+
+      attr_reader :root_lock_path #:nodoc:
+
+      def self.digit_from_lock_path(path) #:nodoc:
+        path[/0*(\d+)$/, 1].to_i
+      end
+
+      def initialize(zookeeper_client, name, root_lock_node = "/_zklocking") 
+        @zk = zookeeper_client
+        @root_lock_node = root_lock_node
+        @path = name
+        @locked = false
+        @waiting = false
+        @root_lock_path = "#{@root_lock_node}/#{@path.gsub("/", "__")}"
+      end
+      
+      # block caller until lock is aquired, then yield
+      def with_lock
+        lock!(true)
+        yield
+      ensure
+        unlock!
+      end
+
+      # the basename of our lock path
+      #
+      # for the lock_path '/_zklocking/foobar/__blah/lock000000007'
+      # lock_basename is 'lock000000007'
+      #
+      # returns nil if lock_path is not set
+      def lock_basename
+        lock_path and File.basename(lock_path)
+      end
+
+      def locked?
+        false|@locked
+      end
+      
+      def unlock!
+        if @locked
+          cleanup_lock_path!
+          @locked = false
+          true
+        end
+      end
+
+      # returns true if this locker is waiting to acquire lock 
+      def waiting? #:nodoc:
+        false|@waiting
+      end
+
+      protected 
+        def in_waiting_status
+          w, @waiting = @waiting, true
+          yield
+        ensure
+          @waiting = w
+        end
+
+        def digit_from(path)
+          self.class.digit_from_lock_path(path)
+        end
+
+        def lock_children(watch=false)
+          @zk.children(root_lock_path, :watch => watch)
+        end
+
+        def ordered_lock_children(watch=false)
+          lock_children(watch).tap do |ary|
+            ary.sort! { |a,b| digit_from(a) <=> digit_from(b) }
+          end
+        end
+
+        def create_root_path!
+          @zk.mkdir_p(@root_lock_path)
+        end
+
+        # prefix is the string that will appear in front of the sequence num,
+        # defaults to 'lock'
+        def create_lock_path!(prefix='lock')
+          @lock_path = @zk.create("#{root_lock_path}/#{prefix}", "", :mode => :ephemeral_sequential)
+          logger.debug { "got lock path #{@lock_path}" }
+          @lock_path
+        rescue Exceptions::NoNode
+          create_root_path!
+          retry
+        end
+
+        def cleanup_lock_path!
+          logger.debug { "removing lock path #{@lock_path}" }
+          @zk.delete(@lock_path)
+          @zk.delete(root_lock_path) rescue Exceptions::NotEmpty
+        end
+    end
+
+    class SharedLocker < LockerBase
+      def lock!(blocking=false)
+        return true if @locked
+        create_lock_path!(SHARED_LOCK_PREFIX)
+
+        if got_read_lock?      
+          @locked = true
+        elsif blocking
+          in_waiting_status do
+            block_until_read_lock!
+          end
+        else
+          # we didn't get the lock, and we're not gonna wait around for it, so
+          # clean up after ourselves
+          cleanup_lock_path!
+          false
+        end
+      end
+
+      def lock_number #:nodoc:
+        @lock_number ||= (lock_path and digit_from(lock_path))
+      end
+
+      # returns the sequence number of the next lowest write lock node
+      #
+      # raises NoWriteLockFoundException when there are no write nodes with a 
+      # sequence less than ours
+      #
+      def next_lowest_write_lock_num #:nodoc:
+        digit_from(next_lowest_write_lock_name)
+      end
+
+      # the next lowest write lock number to ours
+      #
+      # so if we're "read010" and the children of the lock node are:
+      #
+      #   %w[write008 write009 read010 read011]
+      #
+      # then this method will return write009
+      #
+      # raises NoWriteLockFoundException if there were no write nodes with an
+      # index lower than ours 
+      #
+      def next_lowest_write_lock_name #:nodoc:
+        ary = ordered_lock_children()
+        my_idx = ary.index(lock_basename)   # our idx would be 2
+
+        not_found = lambda { raise NoWriteLockFoundException }
+
+        ary[0..my_idx].reverse.find(not_found) { |n| n =~ /^#{EXCLUSIVE_LOCK_PREFIX}/ }
+      end
+
+      def got_read_lock? #:nodoc:
+        false if next_lowest_write_lock_num 
+      rescue NoWriteLockFoundException
+        true
+      end
+
+      protected
+        # TODO: make this generic, can either block or non-block
+        def block_until_read_lock!
+          begin
+            path = [root_lock_path, next_lowest_write_lock_name].join('/')
+            logger.debug { "SharedLocker#block_until_read_lock! path=#{path.inspect}" }
+            @zk.block_until_node_deleted(path)
+          rescue NoWriteLockFoundException
+            # next_lowest_write_lock_name may raise NoWriteLockFoundException,
+            # which means we should not block as we have the lock (there is nothing to wait for)
+          end
+
+          @locked = true
+        end
+    end # SharedLocker
+
+    # An exclusive lock implementation
+    class ExclusiveLocker < LockerBase
+      def lock!(blocking=false)
+        return true if @locked
+        create_lock_path!(EXCLUSIVE_LOCK_PREFIX)
+
+        if got_write_lock?
+          @locked = true
+        elsif blocking
+          in_waiting_status do
+            block_until_write_lock!
+          end
+        else
+          cleanup_lock_path!
+          false
+        end
+      end
+
+      protected
+        # the node that is next-lowest in sequence number to ours, the one we
+        # watch for updates to
+        def next_lowest_node
+          ary = ordered_lock_children()
+          my_idx = ary.index(lock_basename)
+
+          raise WeAreTheLowestLockNumberException if my_idx == 0
+
+          ary[(my_idx - 1)] 
+        end
+
+        def got_write_lock?
+          ordered_lock_children.first == lock_basename
+        end
+
+        def block_until_write_lock!
+          begin
+            path = [root_lock_path, next_lowest_node].join('/')
+            logger.debug { "SharedLocker#block_until_write_lock! path=#{path.inspect}" }
+            @zk.block_until_node_deleted(path)
+          rescue WeAreTheLowestLockNumberException
+          end
+
+          @locked = true
+        end
+    end # ExclusiveLocker
+  end   # SharedLocker
+end     # ZooKeeper
+

data/lib/z_k/logging.rb

@@ -0,0 +1,15 @@
+module ZK
+  module Logging
+    def self.included(mod)
+      mod.extend(ZK::Logging::Methods)
+      mod.send(:include, ZK::Logging::Methods)
+    end
+    
+    module Methods
+      def logger
+        ZK.logger
+      end
+    end
+  end
+end
+

data/lib/z_k/message_queue.rb

@@ -0,0 +1,143 @@
+module ZK
+  # implements a simple message queue based on Zookeeper recipes
+  # @see http://hadoop.apache.org/zookeeper/docs/r3.0.0/recipes.html#sc_recipes_Queues
+  # these are good for low-volume queues only
+  # because of the way zookeeper works, all message *titles* have to be read into memory
+  # in order to see what message to process next
+  # @example
+  #   queue = zk.queue("somequeue")
+  #   queue.publish(some_string)
+  #   queue.poll! # will return one message
+  #   #subscribe will handle messages as they come in
+  #   queue.subscribe do |title, data|
+  #     #handle message
+  #   end
+  class MessageQueue
+
+    # @private
+    # :nodoc:
+    attr_accessor :zk
+
+    # @private
+    # :nodoc:
+    def initialize(zookeeper_client, queue_name, queue_root = "/_zkqueues")
+      @zk = zookeeper_client
+      @queue = queue_name
+      @queue_root = queue_root
+      @zk.create(@queue_root, "", :mode => :persistent) unless @zk.exists?(@queue_root)
+      @zk.create(full_queue_path, "", :mode => :persistent) unless @zk.exists?(full_queue_path)
+    end
+
+    # publish a message to the queue, you can (optionally) use message titles
+    # to guarantee unique messages in the queue
+    # @param [String] data - any arbitrary string value
+    # @param optional [String] message_title - specify a unique message title for this
+    #   message
+    def publish(data, message_title = nil)
+      mode = :persistent_sequential
+      if message_title
+        mode = :persistent
+      else
+        message_title = "message"
+      end
+      @zk.create("#{full_queue_path}/#{message_title}", data, :mode => mode)
+    rescue KeeperException::NodeExists
+      return false
+    end
+
+    # you barely ever need to actually use this method
+    # but lets you remove a message from the queue by specifying
+    # its title
+    # @param [String] message_title the title of the message to remove
+    def delete_message(message_title)
+      full_path = "#{full_queue_path}/#{message_title}"
+      locker = @zk.locker("#{full_queue_path}/#{message_title}")
+      if locker.lock!
+        begin
+          @zk.delete(full_path)
+          return true
+        ensure
+          locker.unlock!
+        end
+      else
+        return false
+      end
+    end
+
+    # grab one message from the queue
+    # used when you don't want to or can't subscribe
+    # @see ZooKeeper::MessageQueue#subscribe
+    def poll!
+      find_and_process_next_available(messages)
+    end
+
+    # @example
+    #   # subscribe like this:
+    #   subscribe {|title, data| handle_message!; true}
+    #   # returning true in the block deletes the message, false unlocks and requeues
+    # @yield [title, data] yield to your block with the message title and the data of
+    #   the message
+    def subscribe(&block)
+      @subscription_block = block
+      @subscription_reference = @zk.watcher.subscribe(full_queue_path) do |event, zk|
+        find_and_process_next_available(@zk.children(full_queue_path, :watch => true))
+      end
+      find_and_process_next_available(@zk.children(full_queue_path, :watch => true))
+    end
+
+    # stop listening to this queue
+    def unsubscribe
+      @subscription_reference.unsubscribe
+    end
+
+    # a list of the message titles in the queue
+    def messages
+      @zk.children(full_queue_path)
+    end
+
+    # highly destructive method!
+    # WARNING! Will delete the queue and all messages in it
+    def destroy!
+      children = @zk.children(full_queue_path)
+      locks = []
+      children.each do |path|
+        lock = @zk.locker("#{full_queue_path}/#{path}")
+        lock.lock!    # XXX(slyphon): should this be a blocking lock?
+        locks << lock
+      end
+      children.each do |path|
+        @zk.delete("#{full_queue_path}/#{path}") rescue ZK::Exceptions::NoNode
+      end
+      @zk.delete(full_queue_path) rescue ZK::Exceptions::NoNode
+      locks.each do |lock|
+        lock.unlock!
+      end
+    end
+
+  private
+    def find_and_process_next_available(messages)
+      messages.sort! {|a,b| digit_from_path(a) <=> digit_from_path(b)}
+      messages.each do |message_title|
+        message_path = "#{full_queue_path}/#{message_title}"
+        locker = @zk.locker(message_path)
+        if locker.lock!
+          begin
+            data = @zk.get(message_path).first
+            result = @subscription_block.call(message_title, data)
+            @zk.delete(message_path) if result
+          ensure
+            locker.unlock!
+          end
+        end
+      end
+    end
+
+    def full_queue_path
+      @full_queue_path ||= "#{@queue_root}/#{@queue}"
+    end
+
+    def digit_from_path(path)
+      path[/\d+$/].to_i
+    end
+  end
+end