zk 0.6.4

data/lib/z_k/mongoid.rb

@@ -0,0 +1,172 @@
+module ZK
+  module Mongoid
+    # provides a lock_for_update method based on the current class name
+    # and Mongoid document _id.
+    #
+    # Before use (in one of your Rails initializers, for example) you should
+    # assign either a ZK::Client or ZK::Pool subclass to
+    # ZooKeeperLockMixin.zk_lock_pool.
+    #
+    # this class assumes the availability of a 'logger' method in the mixee
+    #
+    module Locking
+      VALID_MODES = [:exclusive, :shared].freeze
+
+      @@zk_lock_pool = nil unless defined?(@@zk_lock_pool)
+
+      def self.zk_lock_pool
+        @@zk_lock_pool
+      end
+      
+      def self.zk_lock_pool=(pool)
+        @@zk_lock_pool = pool 
+      end
+
+      # Provides a re-entrant zookeeper-based lock of a record.
+      #
+      # This also makes it possible to detect if the record has been locked before
+      # performing a potentially dangerous operation by using the assert_locked_for_update!
+      # instance method
+      #
+      # Locks are re-entrant per-thread, but will work as a mutex between
+      # threads.
+      #
+      # You can optionally provide a 'name' which will act as a sub-lock of
+      # sorts. For example, if you are going to create an embedded document,
+      # and only want one process to be able to create it at a time (without
+      # clobbering one another), but don't want to lock the entire record, you
+      # can specify a name for the lock, that way the same code running
+      # elsewhere will synchronize based on the parent record and the
+      # particular action specified by +name+.
+      #
+      # ==== Example
+      #
+      # use of "name"
+      #
+      #   class Thing
+      #     include Mongoid::Document
+      #     include ZK::Mongoid::Locking
+      #
+      #     embedded_in :parent, :inverse_of => :thing
+      #   end
+      #
+      #   class Parent
+      #     include Mongoid::Document
+      #     include ZK::Mongoid::Locking
+      #
+      #     embeds_one :thing
+      #
+      #     def lets_create_a_thing
+      #       lock_for_update('thing_creation') do
+      #         raise "We already got one! it's very nice!" if thing
+      #
+      #         do_something_that_might_take_a_while
+      #         create_thing
+      #       end
+      #     end
+      #   end
+      #
+      #
+      # Now, while the creation of the Thing is synchronized, other processes
+      # can update other aspects of Parent.
+      #
+      #
+      def lock_for_update(name=nil)
+        if locked_for_update?(name)
+          logger.debug { "we are locked for update, yield to the block" }
+          yield
+        else
+          zk_with_lock(:mode => :exclusive, :name => name) { yield }
+        end
+      end
+      alias :with_exclusive_lock :lock_for_update
+
+      def with_shared_lock(name=nil)
+        if locked_for_share?(name)
+          yield
+        else
+          zk_with_lock(:mode => :shared, :name => name) { yield }
+        end
+      end
+
+      # raises MustBeExclusivelyLockedException if we're not currently inside a
+      # lock (optionally with +name+)
+      def assert_locked_for_update!(name=nil)
+        raise ZK::Exceptions::MustBeExclusivelyLockedException unless locked_for_update?(name)
+      end
+
+      # raises MustBeShareLockedException if we're not currently inside a shared lock
+      # (optionally with +name+)
+      def assert_locked_for_share!(name=nil)
+        raise ZK::Exceptions::MustBeShareLockedException unless locked_for_share?(name)
+      end
+
+      def locked_for_update?(name=nil) #:nodoc:
+        zk_mongoid_lock_registry[:exclusive].include?(zk_lock_name(name))
+      end
+
+      def locked_for_share?(name=nil) #:nodoc:
+        zk_mongoid_lock_registry[:shared].include?(zk_lock_name(name))
+      end
+
+      def zk_lock_name(name=nil) #:nodoc:
+        [self.class.to_s, self.id.to_s, name].compact.join('-')
+      end
+
+      protected
+        def zk_mongoid_lock_registry
+          Thread.current.zk_mongoid_lock_registry ||= { :shared => Set.new, :exclusive => Set.new }
+        end
+
+        def zk_add_path_lock(opts={})
+          mode, name = opts.values_at(:mode, :name)
+
+          raise ArgumentError, "You must specify a :mode option" unless mode
+
+          zk_assert_valid_mode!(mode)
+
+          logger.debug { "adding #{zk_lock_name(name).inspect} to #{mode} lock registry" }
+
+          self.zk_mongoid_lock_registry[mode] << zk_lock_name(name)
+        end
+
+        def zk_remove_path_lock(opts={})
+          mode, name = opts.values_at(:mode, :name)
+
+          raise ArgumentError, "You must specify a :mode option" unless mode
+
+          zk_assert_valid_mode!(mode)
+
+          logger.debug { "removing #{zk_lock_name(name).inspect} from #{mode} lock registry" }
+
+          zk_mongoid_lock_registry[mode].delete(zk_lock_name(name))
+        end
+
+        def zk_with_lock(opts={})
+          mode, name = opts.values_at(:mode, :name)
+
+          zk_assert_valid_mode!(mode)
+
+          zk_lock_pool.with_lock(zk_lock_name(name), :mode => mode) do
+            zk_add_path_lock(opts)
+
+            begin
+              logger.debug { "acquired #{zk_lock_name(name).inspect}" }
+              yield
+            ensure
+              logger.debug { "releasing #{zk_lock_name(name).inspect}" }
+              zk_remove_path_lock(opts)
+            end
+          end
+        end
+
+        def zk_lock_pool
+          @zk_lock_pool ||= ::ZK::Mongoid::Locking.zk_lock_pool
+        end
+
+        def zk_assert_valid_mode!(mode)
+          raise ArgumentError, "#{mode.inspect} is not a valid mode value" unless VALID_MODES.include?(mode)
+        end
+    end
+  end
+end

data/lib/z_k/pool.rb

@@ -0,0 +1,254 @@
+module ZK
+  module Pool
+    class Base
+      attr_reader :connections #:nodoc:
+
+      def initialize
+        @state = :init
+
+        @connections = []
+        @connections.extend(MonitorMixin)
+        @checkin_cond = @connections.new_cond
+      end
+
+      # has close_all! been called on this ConnectionPool ?
+      def closed?
+        @state == :closed
+      end
+
+      # is the pool shutting down?
+      def closing?
+        @state == :closing
+      end
+
+      # is the pool initialized and in normal operation?
+      def open?
+        @state == :open
+      end
+
+      # has the pool entered the take-no-prisoners connection closing part of shutdown?
+      def forced?
+        @state == :forced
+      end
+
+      # close all the connections on the pool
+      # @param optional Boolean graceful allow the checked out connections to come back first?
+      def close_all!
+        synchronize do 
+          return unless open?
+          @state = :closing
+
+          @checkin_cond.wait_until { (@pool.size == @connections.length) or closed? }
+
+          force_close!
+        end
+      end
+
+      # calls close! on all connection objects, whether or not they're back in the pool
+      # this is DANGEROUS!
+      def force_close! #:nodoc:
+        synchronize do
+          return if (closed? or forced?)
+          @state = :forced
+
+          @pool.clear
+
+          while cnx = @connections.shift
+            cnx.close!
+          end
+
+          @state = :closed
+
+          # free any waiting 
+          @checkin_cond.broadcast
+        end
+      end
+
+      # yields next available connection to the block
+      #
+      # raises PoolIsShuttingDownException immediately if close_all! has been
+      # called on this pool
+      def with_connection
+        assert_open!
+
+        cnx = checkout(true)
+        yield cnx
+      ensure
+        checkin(cnx)
+      end
+
+      #lock lives on past the connection checkout
+      def locker(path)
+        with_connection do |connection|
+          connection.locker(path)
+        end
+      end
+
+      #prefer this method if you can (keeps connection checked out)
+      def with_lock(name, opts={}, &block)
+        with_connection do |connection|
+          connection.with_lock(name, opts, &block)
+        end
+      end
+
+      # handle all
+      def method_missing(meth, *args, &block)
+        with_connection do |connection|
+          connection.__send__(meth, *args, &block)
+        end
+      end
+
+      def size #:nodoc:
+        @pool.size
+      end
+
+      def pool_state #:nodoc:
+        @state
+      end
+
+      protected
+        def synchronize
+          @connections.synchronize { yield }
+        end
+
+        def assert_open!
+          raise Exceptions::PoolIsShuttingDownException unless open? 
+        end
+
+    end # Base
+
+    # like a Simple pool but has high/low watermarks, and can grow dynamically as needed
+    class Bounded < Base
+      DEFAULT_OPTIONS = {
+        :timeout      => 10,
+        :min_clients  => 1,
+        :max_clients  => 10,
+      }.freeze
+
+      # opts:
+      # * <tt>:timeout</tt>: connection establishement timeout
+      # * <tt>:min_clients</tt>: how many clients should be start out with
+      # * <tt>:max_clients</tt>: the maximum number of clients we will create in response to demand
+      def initialize(host, opts={})
+        super()
+        @host = host
+        @connection_args = opts
+
+        opts = DEFAULT_OPTIONS.merge(opts)
+
+        @min_clients = Integer(opts.delete(:min_clients))
+        @max_clients = Integer(opts.delete(:max_clients))
+        @connection_timeout = opts.delete(:timeout)
+
+        # for compatibility w/ ClientPool we'll use @connections for synchronization
+        @pool = []            # currently available connections
+
+        synchronize do
+          populate_pool!(@min_clients)
+          @state = :open
+        end
+      end
+
+      # returns the current number of allocated clients in the pool (not
+      # available clients)
+      def size
+        @connections.length
+      end
+
+      # clients available for checkout (at time of call)
+      def available_size
+        @pool.length
+      end
+
+      def checkin(connection)
+        synchronize do
+          return if @pool.include?(connection)
+
+          @pool.unshift(connection)
+          @checkin_cond.signal
+        end
+      end
+
+      # number of threads waiting for connections
+      def count_waiters #:nodoc:
+        @checkin_cond.count_waiters
+      end
+
+      def checkout(blocking=true) 
+        raise ArgumentError, "checkout does not take a block, use .with_connection" if block_given?
+        synchronize do
+          while true
+            assert_open!
+
+            if @pool.length > 0
+              cnx = @pool.shift
+              
+              # if the cnx isn't connected? then remove it from the pool and go
+              # through the loop again. when the cnx's on_connected event fires, it
+              # will add the connection back into the pool
+              next unless cnx.connected?
+
+              # otherwise we return the cnx
+              return cnx
+            elsif can_grow_pool?
+              add_connection!
+              next
+            elsif blocking
+              @checkin_cond.wait_while { @pool.empty? and open? }
+              next
+            else
+              return false
+            end
+          end
+        end
+      end
+
+      protected
+        def populate_pool!(num_cnx)
+          num_cnx.times { add_connection! }
+        end
+
+        def add_connection!
+          synchronize do
+            cnx = create_connection
+            @connections << cnx 
+
+            cnx.on_connected { checkin(cnx) }
+          end
+        end
+
+        def can_grow_pool?
+          synchronize { @connections.size < @max_clients }
+        end
+
+        def create_connection
+          ZK.new(@host, @connection_timeout, @connection_args)
+        end
+    end # Bounded
+
+    # create a connection pool useful for multithreaded applications
+    #
+    # Will spin up +number_of_connections+ at creation time and remain fixed at
+    # that number for the life of the pool.
+    #
+    # ==== Example
+    #   pool = ZK::Pool::Simple.new("localhost:2181", 10)
+    #   pool.checkout do |zk|
+    #     zk.create("/mynew_path")
+    #   end
+    class Simple < Bounded
+      # initialize a connection pool using the same optons as ZK.new
+      # @param String host the same arguments as ZK.new
+      # @param Integer number_of_connections the number of connections to put in the pool
+      # @param optional Hash opts Options to pass on to each connection
+      # @return ZK::ClientPool
+      def initialize(host, number_of_connections=10, opts = {})
+        opts = opts.dup
+        opts[:max_clients] = opts[:min_clients] = number_of_connections.to_i
+
+        super(host, opts)
+      end
+    end # Simple
+  end   # Pool
+end     # ZK
+

data/lib/z_k/threadpool.rb

@@ -0,0 +1,109 @@
+module ZK
+  # a simple threadpool for running blocks of code off the main thread
+  class Threadpool
+    include Logging
+
+    DEFAULT_SIZE = 5
+
+    class << self
+      # size of the ZK.defer threadpool (defaults to 5)
+      attr_accessor :default_size
+      ZK::Threadpool.default_size = DEFAULT_SIZE
+    end
+
+    # the size of this threadpool
+    attr_reader :size
+
+    def initialize(size=nil)
+      @size = size || self.class.default_size
+
+      @threadpool = []
+      @threadqueue = ::Queue.new
+
+      @mutex = Mutex.new
+
+      start!
+    end
+
+    # Queue an operation to be run on an internal threadpool. You may either
+    # provide an object that responds_to?(:call) or pass a block. There is no
+    # mechanism for retrieving the result of the operation, it is purely
+    # fire-and-forget, so the user is expected to make arrangements for this in
+    # their code. 
+    #
+    def defer(callable=nil, &blk)
+      callable ||= blk
+
+      # XXX(slyphon): do we care if the threadpool is not running?
+      raise Exceptions::ThreadpoolIsNotRunningException unless running?
+      raise ArgumentError, "Argument to Threadpool#defer must respond_to?(:call)" unless callable.respond_to?(:call)
+
+      @threadqueue << callable
+      nil
+    end
+
+    def running?
+      @mutex.synchronize { @running }
+    end
+
+    # starts the threadpool if not already running
+    def start!
+      @mutex.synchronize do
+        return false if @running
+        @running = true
+        spawn_threadpool
+      end
+      true
+    end
+
+    # join all threads in this threadpool, they will be given a maximum of +timeout+
+    # seconds to exit before they are considered hung and will be ignored (this is an
+    # issue with threads in general: see 
+    # http://blog.headius.com/2008/02/rubys-threadraise-threadkill-timeoutrb.html for more info)
+    #
+    # the default timeout is 2 seconds per thread
+    # 
+    def shutdown(timeout=2)
+      @mutex.synchronize do
+        return unless @running
+
+        @running = false
+
+        @threadqueue.clear
+        @size.times { @threadqueue << KILL_TOKEN }
+
+        while th = @threadpool.shift
+          begin
+            th.join(timeout)
+          rescue Exception => e
+            logger.error { "Caught exception shutting down threadpool" }
+            logger.error { e.to_std_format }
+          end
+        end
+      end
+
+      nil
+    end
+
+    private
+      def spawn_threadpool #:nodoc:
+        until @threadpool.size == @size.to_i
+          thread = Thread.new do
+            while @running
+              begin
+                op = @threadqueue.pop
+                break if op == KILL_TOKEN
+                op.call
+              rescue Exception => e
+                logger.error { "Exception caught in threadpool" }
+                logger.error { e.to_std_format }
+              end
+            end
+          end
+
+          @threadpool << thread
+        end
+      end
+  end
+end
+

data/lib/z_k/version.rb

@@ -0,0 +1,3 @@
+module ZK
+  VERSION = "0.6.4"
+end

data/lib/z_k.rb

@@ -0,0 +1,73 @@
+require 'rubygems'
+require 'bundler/setup'
+
+require 'logger'
+require 'zookeeper'
+require 'forwardable'
+require 'thread'
+require 'monitor'
+require 'set'
+
+require 'z_k/logging'
+require 'z_k/exceptions'
+require 'z_k/threadpool'
+require 'z_k/event_handler_subscription'
+require 'z_k/event_handler'
+require 'z_k/message_queue'
+# require 'z_k/locker_base'
+require 'z_k/locker'
+require 'z_k/extensions'
+require 'z_k/election'
+require 'z_k/mongoid'
+require 'z_k/client'
+require 'z_k/pool'
+
+module ZK
+  ZK_ROOT = File.expand_path('../..', __FILE__)
+
+  KILL_TOKEN = :__kill_token__ #:nodoc:
+
+
+  # The logger used by the ZK library. uses a Logger to +/dev/null+ by default
+  #
+  def self.logger
+    @logger ||= Logger.new('/dev/null')
+  end
+
+  # Assign the Logger instance to be used by ZK
+  def self.logger=(logger)
+    @logger = logger
+  end
+
+  # Create a new ZK::Client instance. If no arguments are given, the default
+  # config of 'localhost:2181' will be used. Otherwise all args will be passed
+  # to ZK::Client#new
+  #
+  # if a block is given, it will be yielded the client *before* the connection
+  # is established, this is useful for registering connected-state handlers.
+  #
+  def self.new(*args, &block)
+    # XXX: might need to do some param parsing here
+   
+    opts = args.pop if args.last.kind_of?(Hash)
+    args = %w[localhost:2181] if args.empty?
+
+    # ignore opts for now
+    Client.new(*args, &block)
+  end
+
+  # Like new, yields a connection to the given block and closes it when the
+  # block returns
+  def self.open(*args)
+    cnx = new(*args)
+    yield cnx
+  ensure
+    cnx.close! if cnx
+  end
+
+  # creates a new ZK::Pool::Bounded with the default options.
+  def self.new_pool(host, opts={})
+    ZK::Pool::Bounded.new(host, opts)
+  end
+end
+

data/lib/zk.rb

@@ -0,0 +1,2 @@
+require File.expand_path('../z_k', __FILE__)
+