zookeeper_mutex

gistfile1.rb

require 'zookeeper'

class Lock  
  def initialize(host, root="/my-app")    
    @zk = Zookeeper.new(host)
    @root = root
  end
  
  def with_lock(app, timeout, timeout_callback, &block)
    new_lock_res = @zk.create(:path => "#{@root}/#{app}-", :sequence => true, :ephemeral => true)
    unique_lock_path = new_lock_res[:path]
    if get_lock(unique_lock_path, timeout)
      yield
      @zk.delete(:path => unique_lock_path)
    else
      timeout_callback.call
    end
  end

  private
  def get_lock(unique_lock_path, timeout)
    lock_key = unique_lock_path.gsub(/^#{Regexp.quote(@root)}\//, '')

    (0..4).each do
      children = @zk.get_children(:path => @root)[:children].sort
      watcher = Zookeeper::WatcherCallback.new {}
      if (children.first == lock_key)
        return true
      else
        less_than_path_idx = (children.index {|p| p == lock_key}) - 1
        stat_res = @zk.stat(:path => "#{@root}/#{children[less_than_path_idx]}", 
                           :watcher => watcher)
        if stat_res[:stat].exists
          success = wait_until(timeout) { watcher.completed? }
          if !success
            return false            
          end
        end
      end
    end
    return false
  end

  def wait_until(timeout=10, &block)
    time_to_stop = Time.now + timeout
    until yield do
      if Time.now > time_to_stop
        return false
      end
      sleep 0.1
    end
    return true
  end
end