zk 0.6.4

data/lib/z_k/election.rb

@@ -0,0 +1,411 @@
+module ZK
+  # ==== Overview
+  #
+  # This module implements the "leader election" protocols described
+  # {here}[http://hadoop.apache.org/zookeeper/docs/current/recipes.html#sc_leaderElection].
+  #
+  # There are Candidates and Observers. Candidates take part in elections and
+  # all have equal ability and chance to become the leader. When a leader is
+  # decided, they hold onto the leadership role until they die. When the leader
+  # dies, an election is held and the winner has its +on_winning_election+
+  # callbacks fired, and the losers have their +on_losing_election+ callbacks
+  # fired. When all of the +on_winning_election+ callbacks have completed
+  # (completing whatever steps are necessary to assume the leadership role),
+  # the leader will "acknowledge" that it has taken over by creating an
+  # ephemeral node at a known location (with optional data that the Observers
+  # can then read and take action upon). Note that when this node is created,
+  # it means the *leader* has finished taking over, but it does *not* mean that
+  # all the slaves have completed *their* tasks.
+  #
+  # Observers are interested parties in the election, the "constituents" of the
+  # process. They can register callbacks to be fired when a new leader has been
+  # elected and when a leader has died. The new leader callbacks will only fire
+  # once the leader has acknowledged its role, so they can be sure that the
+  # leader is ready to perform its duties.
+  #
+  # ==== Use Case / Example
+  #
+  # One problem this pattern can be used to solve is failover between two
+  # database nodes. Candidates set up callbacks to both take over as master
+  # and to follow the master if they lose the election. On the client side,
+  # Obesrvers are set up to follow the "leader ack" node. The leader writes its
+  # connection info to the "leader ack" node, and the clients can reconnect to
+  # the currently active leader.
+  #
+  #
+  #   def server
+  #     candidate = @zk.election_candidate("database_election", "dbhost2.fqdn.tld:4567", :follow => :leader)
+  #     candidate.on_winning_election { become_master_node! }
+  #     candidate.on_losing_election { become_slave_of_master! }
+  #
+  #     @zk.on_connected do
+  #       candidate.vote!
+  #     end
+  #   end
+  #
+  # Note that as soon as vote! is called, either the on_winning_election or
+  # on_losing_election callbacks will be called. 
+  #
+  #
+  module Election
+    VOTE_PREFIX = 'ballot'.freeze
+    ROOT_NODE = '/_zkelection'.freeze
+
+    VALID_FOLLOW_OPTIONS = [:next_node, :leader].freeze
+
+    DEFAULT_OPTS = {
+      :root_election_node => ROOT_NODE,
+    }.freeze
+ 
+    class Base
+     include Logging
+
+     attr_reader :zk, :vote_path, :root_election_node
+
+      def initialize(client, name, opts={})
+        @zk = client
+        @name = name
+        opts = DEFAULT_OPTS.merge(opts)
+        @root_election_node = opts[:root_election_node]
+        @mutex = Monitor.new
+      end
+
+      # holds the ephemeral nodes of this election
+      def root_vote_path #:nodoc:
+        @root_vote_path ||= "#{@root_election_node}/#{@name.gsub('/', '__')}"
+      end
+
+      # this znode will be created as an acknowledgement by the leader 
+      # that it's aware of its status as the new leader and has run its 
+      # procedures to become master
+      def leader_ack_path
+        @leader_ack_path ||= "#{root_vote_path}/leader_ack"
+      end
+     
+      def cast_ballot!(data)
+        return if @vote_path
+        create_root_path!
+        @vote_path = @zk.create("#{root_vote_path}/#{VOTE_PREFIX}", data, :mode => :ephemeral_sequential)
+      rescue Exceptions::NoNode
+        retry
+      end
+      
+      # has the leader acknowledged their role?
+      def leader_acked?(watch=false)
+        @zk.exists?(leader_ack_path, :watch => watch)
+      end
+      
+      # return the data from the current leader or nil if there is no current leader
+      def leader_data
+        @zk.get(leader_ack_path).first
+      rescue Exceptions::NoNode
+      end
+
+      # Asynchronously call the block when the leader has acknowledged its
+      # role. The given block will *always* be called on a background thread. 
+      def on_leader_ack(&block)
+        creation_sub = @zk.watcher.register(leader_ack_path) do |event|
+          case event.type
+          when Zookeeper::ZOO_CREATED_EVENT, Zookeeper::ZOO_CHANGED_EVENT
+            begin
+              logger.debug { "in #{leader_ack_path} watcher, got creation event, notifying" }
+              block.call
+            ensure
+              creation_sub.unregister
+            end
+          else
+            if @zk.exists?(leader_ack_path, :watch => true)
+              begin
+                logger.debug { "in #{leader_ack_path} watcher, node created behind our back, notifying" }
+                block.call
+              ensure
+                creation_sub.unregister
+              end
+            else
+              logger.debug { "in #{leader_ack_path} watcher, got non-creation event, re-watching" }
+            end
+          end
+        end
+
+        @zk.defer do
+          if @zk.exists?(leader_ack_path, :watch => true)
+            logger.debug { "on_leader_ack, #{leader_ack_path} exists, calling block" }
+            begin
+              block.call
+            ensure
+              creation_sub.unregister if creation_sub
+            end
+          end
+        end
+      end
+
+      protected
+        def create_root_path!
+          @zk.mkdir_p(root_vote_path)
+        end
+
+        def vote_basename
+          vote_path and File.basename(vote_path)
+        end
+
+        def digit(path)
+          path[/\d+$/].to_i
+
+        end
+
+        def safe_call(*callbacks)
+          callbacks.each do |cb|
+            begin
+              cb.call
+            rescue Exception => e
+              logger.error { "Error caught in user supplied callback" }
+              logger.error { e.to_std_format }
+            end
+          end
+        end
+    end
+
+    # This class is for registering candidates in the leader election. This instance will
+    # participate in votes for becoming the leader and will be notified in the
+    # case where it needs to take over.
+    #
+    # if data is given, it will be used as the content of both our ballot and
+    # the leader acknowledgement node if and when we become the leader.
+    class Candidate < Base
+      def initialize(client, name, opts={})
+        super(client, name, opts)
+        opts = DEFAULT_OPTS.merge(opts)
+
+        @leader     = nil 
+        @data       = opts[:data] || ''
+        @vote_path  = nil
+       
+        @winner_callbacks = []
+        @loser_callbacks = []
+
+        @next_node_ballot_sub = nil # the subscription for next-node failure
+      end
+
+      def leader?
+        false|@leader
+      end
+
+      # true if leader has been determined at least once (used in tests)
+      def voted? #:nodoc:
+        !@leader.nil?
+      end
+      
+      # When we win the election, we will call the procs registered using this
+      # method.
+      def on_winning_election(&block)
+        @winner_callbacks << block
+      end
+
+      # When we lose the election and are relegated to the shadows, waiting for
+      # the leader to make one small misstep, where we can finally claim what
+      # is rightfully ours! MWUAHAHAHAHAHA(*cough*)
+      def on_losing_election(&block)
+        @loser_callbacks << block
+      end
+
+      # These procs should be run in the case of an error when trying to assume
+      # the leadership role. This should *probably* be a "hara-kiri" or STONITH
+      # type procedure (i.e. kill the candidate)
+      #
+      def on_takeover_error #:nodoc:
+        raise NotImplementedError
+      end
+
+      # volunteer to become the leader. if we win, on_winning_election blocks will
+      # be called, otherwise, wait for next election
+      #
+      # +data+ will be placed in the znode representing our vote
+      def vote!
+        @mutex.synchronize do
+          clear_next_node_ballot_sub!
+          cast_ballot!(@data) unless @vote_path
+          check_election_results!
+        end
+      end
+
+      protected
+        # the inauguration, as it were
+        def acknowledge_win!
+          @zk.create(leader_ack_path, @data, :ephemeral => true) rescue Exceptions::NodeExists
+        end
+
+        # return the list of ephemeral vote nodes
+        def get_ballots
+          @zk.children(root_vote_path).grep(/^ballot/).tap do |ballots|
+            ballots.sort! {|a,b| digit(a) <=> digit(b) }
+          end
+        end
+
+        # if +watch_next+ is true, we register a watcher for the next-lowest
+        # index number in the list of ballots
+        #
+        def check_election_results!
+          #return if leader?         # we already know we're the leader
+          ballots = get_ballots()
+
+          our_idx = ballots.index(vote_basename)
+          
+          if our_idx == 0           # if we have the lowest number
+            logger.info { "ZK: We have become leader, data: #{@data.inspect}" }
+            handle_winning_election
+          else
+            logger.info { "ZK: we are not the leader, data: #{@data.inspect}" }
+            handle_losing_election(our_idx, ballots)
+          end
+        end
+
+        def handle_winning_election
+          @leader = true  
+          fire_winning_callbacks!
+          acknowledge_win!
+        end
+
+        def handle_losing_election(our_idx, ballots)
+          @leader = false
+
+          on_leader_ack do
+            fire_losing_callbacks!
+
+            next_ballot = File.join(root_vote_path, ballots[our_idx - 1])
+
+            logger.info { "ZK: following #{next_ballot} for changes, #{@data.inspect}" }
+
+            @next_node_ballot_sub ||= @zk.watcher.register(next_ballot) do |event| 
+              if event.node_deleted? 
+                logger.debug { "#{next_ballot} was deleted, voting, #{@data.inspect}" }
+                vote! 
+              else
+                # this takes care of the race condition where the leader ballot would
+                # have been deleted before we could re-register to receive updates
+                # if zk.stat returns false, it means the path was deleted
+                unless @zk.exists?(next_ballot, :watch => true)
+                  logger.debug { "#{next_ballot} was deleted (detected on re-watch), voting, #{@data.inspect}" }
+                  vote! 
+                end
+              end
+            end
+
+            # this catches a possible race condition, where the leader has died before
+            # our callback has fired. In this case, retry and do this procedure again
+            unless @zk.stat(next_ballot, :watch => true).exists?
+              logger.debug { "#{@data.inspect}: the node #{next_ballot} did not exist, retrying" }
+              vote!
+            end
+          end
+        end
+
+        def clear_next_node_ballot_sub!
+          if @next_node_ballot_sub
+            @next_node_ballot_sub.unsubscribe 
+            @next_node_ballot_sub = nil
+          end
+        end
+
+        def fire_winning_callbacks!
+          safe_call(*@winner_callbacks)
+        end
+
+        def fire_losing_callbacks!
+          safe_call(*@loser_callbacks)
+        end
+    end
+
+    class Observer < Base
+      def initialize(client, name, opts={})
+        super
+        @leader_death_cbs = []
+        @new_leader_cbs = []
+        @deletion_sub = @creation_sub = nil
+        @leader_alive = nil
+        @observing = false
+      end
+
+      # our current idea about the state of the election
+      def leader_alive #:nodoc:
+        @mutex.synchronize { @leader_alive }
+      end
+
+      # register callbacks that should be fired when a leader dies
+      def on_leaders_death(&blk)
+        @leader_death_cbs << blk
+      end
+
+      # register callbacks for when the new leader has acknowledged their role
+      # returns a subscription object that can be used to cancel further events
+      def on_new_leader(&blk)
+        @new_leader_cbs << blk
+      end
+
+      def observe!
+        @mutex.synchronize do
+          return if @observing 
+          @observing = true
+
+          @leader_ack_sub ||= @zk.watcher.register(leader_ack_path) do |event|
+            if event.node_deleted?
+              the_king_is_dead 
+            elsif event.node_created?
+              long_live_the_king
+            else
+              acked = leader_acked?(true) 
+
+              # If the current state of the system is not what we think it should be
+              # a transition has occurred and we should fire our callbacks
+              if (acked and !@leader_alive)
+                long_live_the_king 
+              elsif (!acked and @leader_alive)
+                the_king_is_dead
+              else
+                # things are how we think they should be, so just wait for the
+                # watch to fire
+              end
+            end
+          end
+
+          leader_acked?(true) ? long_live_the_king : the_king_is_dead
+        end
+      end
+
+      def close
+        @mutex.synchronize do
+          return unless @observing
+
+          @deletion_sub.unregister if @deletion_sub
+          @creation_sub.unregister if @creation_sub
+
+          @deletion_sub = @creation_sub = nil
+
+          @leader_death_cbs.clear
+          @new_leader_cbs.clear
+
+          @leader_alive = nil
+          @observing = false
+        end
+      end
+
+      protected
+        def the_king_is_dead
+          @mutex.synchronize do
+            safe_call(*@leader_death_cbs)
+            @leader_alive = false
+          end
+
+          long_live_the_king if leader_acked?(true)
+        end
+
+        def long_live_the_king
+          @mutex.synchronize do
+            safe_call(*@new_leader_cbs)
+            @leader_alive = true
+          end
+
+          the_king_is_dead unless leader_acked?(true)
+        end
+    end
+  end
+end

data/lib/z_k/event_handler.rb

@@ -0,0 +1,202 @@
+module ZK
+  # this is the default watcher provided by the zookeeper connection
+  # watchers are implemented by adding the :watch => true flag to
+  # any #children or #get or #exists calls
+  # you never really need to initialize this yourself
+  class EventHandler
+    include org.apache.zookeeper.Watcher if defined?(JRUBY_VERSION)
+    include ZK::Logging
+
+    VALID_WATCH_TYPES = [:data, :child].freeze
+
+    ZOOKEEPER_WATCH_TYPE_MAP = {
+      Zookeeper::ZOO_CREATED_EVENT => :data,
+      Zookeeper::ZOO_DELETED_EVENT => :data,
+      Zookeeper::ZOO_CHANGED_EVENT => :data,
+      Zookeeper::ZOO_CHILD_EVENT   => :child,
+    }.freeze
+
+    attr_accessor :zk  # :nodoc:
+
+    # @private
+    # :nodoc:
+    def initialize(zookeeper_client)
+      @zk = zookeeper_client
+      @callbacks = Hash.new { |h,k| h[k] = [] }
+
+      @mutex = Monitor.new
+
+      @outstanding_watches = VALID_WATCH_TYPES.inject({}) do |h,k|
+        h.tap { |x| x[k] = Set.new }
+      end
+    end
+
+    # register a path with the handler
+    # your block will be called with all events on that path.
+    # aliased as #subscribe
+    # @param [String] path the path you want to listen to
+    # @param [Block] block the block to execute when a watch event happpens
+    # @yield [connection, event] We will call your block with the connection the
+    #   watch event occured on and the event object
+    # @return [ZooKeeper::EventHandlerSubscription] the subscription object
+    #   you can use to to unsubscribe from an event
+    # @see ZooKeeper::WatcherEvent
+    # @see ZooKeeper::EventHandlerSubscription
+    def register(path, &block)
+      logger.debug { "EventHandler#register path=#{path.inspect}" }
+      EventHandlerSubscription.new(self, path, block).tap do |subscription|
+        synchronize { @callbacks[path] << subscription }
+      end
+    end
+    alias :subscribe :register
+
+    # registers a "state of the connection" handler
+    #
+    # @param [String] state the state you want to register for
+    # @param [Block] block the block to execute on state changes
+    # @yield [connection, event] yields your block with
+    def register_state_handler(state, &block)
+      register(state_key(state), &block)
+    end
+
+    # @deprecated use #unsubscribe on the subscription object
+    # @see ZooKeeper::EventHandlerSubscription#unsubscribe
+    def unregister_state_handler(*args)
+      if args.first.is_a?(EventHandlerSubscription)
+        unregister(args.first)
+      else
+        unregister(state_key(args.first), args[1])
+      end
+    end
+
+    # @deprecated use #unsubscribe on the subscription object
+    # @see ZooKeeper::EventHandlerSubscription#unsubscribe
+    def unregister(*args)
+      if args.first.is_a?(EventHandlerSubscription)
+        subscription = args.first
+      elsif args.first.is_a?(String) and args[1].is_a?(EventHandlerSubscription)
+        subscription = args[1]
+      else
+        path, index = args[0..1]
+        synchronize { @callbacks[path][index] = nil }
+        return
+      end
+
+      synchronize do
+        ary = @callbacks[subscription.path]
+
+        idx = ary.index(subscription) and ary.delete_at(idx)
+      end
+
+      nil
+    end
+    alias :unsubscribe :unregister
+
+    # called from the client-registered callback when an event fires
+    def process(event) #:nodoc:
+      logger.debug { "EventHandler#process dispatching event: #{event.inspect}" } unless event.type == -1
+      event.zk = @zk
+
+      cb_key = 
+        if event.node_event?
+          event.path
+        elsif event.state_event?
+          state_key(event.state)
+        else
+          raise ZKError, "don't know how to process event: #{event.inspect}"
+        end
+
+      cb_ary = synchronize do 
+        if event.node_event?
+          if watch_type = ZOOKEEPER_WATCH_TYPE_MAP[event.type]
+            logger.debug { "re-allowing #{watch_type.inspect} watches on path #{event.path.inspect}" }
+            
+            # we recieved a watch event for this path, now we allow code to set new watchers
+            @outstanding_watches[watch_type].delete(event.path)
+          end
+        end
+
+        @callbacks[cb_key].dup
+      end
+
+      cb_ary.compact!
+
+      safe_call(cb_ary, event)
+    end
+
+    # used during shutdown to clear registered listeners
+    def clear! #:nodoc:
+      synchronize do
+        @callbacks.clear
+        nil
+      end
+    end
+
+    def synchronize #:nodoc:
+      @mutex.synchronize { yield }
+    end
+
+    def get_default_watcher_block
+      @default_watcher_block ||= lambda do |hash|
+        watcher_callback.tap do |cb|
+          cb.call(hash)
+        end
+      end
+    end
+
+    # implements not only setting up the watcher callback, but deduplicating 
+    # event delivery. Keeps track of in-flight watcher-type+path requests and
+    # doesn't re-register the watcher with the server until a response has been
+    # fired. This prevents one event delivery to *every* callback per :watch => true
+    # argument.
+    def setup_watcher!(watch_type, opts)
+      return unless opts.delete(:watch)
+
+      synchronize do
+        set = @outstanding_watches.fetch(watch_type)
+        path = opts[:path]
+
+        if set.add?(path)
+          # this path has no outstanding watchers, let it do its thing
+          opts[:watcher] = watcher_callback 
+        else
+          # outstanding watch for path and data pair already exists, so ignore
+          logger.debug { "outstanding watch request for path #{path.inspect} and watcher type #{watch_type.inspect}, not re-registering" }
+        end
+      end
+    end
+
+    protected
+      def watcher_callback
+        ZookeeperCallbacks::WatcherCallback.create { |event| process(event) }
+      end
+
+      def state_key(arg)
+        int = 
+          case arg
+          when String, Symbol
+            ZookeeperConstants.const_get(:"ZOO_#{arg.to_s.upcase}_STATE")
+          when Integer
+            arg
+          else
+            raise NameError # ugh lame
+          end
+
+        "state_#{int}"
+      rescue NameError
+        raise ArgumentError, "#{arg} is not a valid zookeeper state", caller
+      end
+
+      def safe_call(callbacks, *args)
+        callbacks.each do |cb|
+          begin
+            cb.call(*args) if cb.respond_to?(:call)
+          rescue Exception => e
+            logger.error { "Error caught in user supplied callback" }
+            logger.error { e.to_std_format }
+          end
+        end
+      end
+  end
+end
+

data/lib/z_k/event_handler_subscription.rb

@@ -0,0 +1,29 @@
+module ZK
+  # the subscription object that is passed back from subscribing
+  # to events.
+  # @see ZooKeeperEventHandler#subscribe
+  class EventHandlerSubscription
+    attr_accessor :event_handler, :path, :callback
+
+    # @private
+    # :nodoc:
+    def initialize(event_handler, path, callback)
+      @event_handler, @path, @callback = event_handler, path, callback
+    end
+
+    # unsubscribe from the path or state you were watching
+    # @see ZooKeeperEventHandler#subscribe
+    def unsubscribe
+      @event_handler.unregister(self)
+    end
+    alias :unregister :unsubscribe
+
+    # @private
+    # :nodoc:
+    def call(event)
+      callback.call(event)
+    end
+
+  end
+end
+