Implementation of distributed locking with ZooKeeper

zk_demo_lock.rb

#!/usr/bin/env ruby
require 'rubygems'
require 'zk'
require './zk_ext'

ZK.open '10.0.0.100:2181' do |zk|
  zk.with_lock('test') do
    zk.explore '/'
  end
end

zk_ext.rb

#!/usr/bin/env ruby
require 'rubygems'
require 'zk'

class ZK::Client::Base
  def explore zpath
    ls(zpath).map do |zcpath|
      explore zcpath
    end
  end

  def ls zpath
    children(zpath).map do |child|
      zcpath = "#{zpath}/#{child}"
      zcpath.gsub! '//', '/'
      puts zcpath
      zcpath
    end
  end

  def suffix_of lock_path
    begin
      lock_path.match(/\-(\d+)$/)[1]
    rescue
      ''
    end
  end

  def lock name
    lock_node = "/_zklocks"
    create lock_node, '', { :mode => :persistent, :ignore => :node_exists }
    lock_path = create "#{lock_node}/#{name}-", '', :mode => :persistent_sequential
    acquire_lock name, lock_node, lock_path
    @locks ||= {}
    @locks[name] = lock_path
    return self
  end

  def acquire_lock name, lock_node, lock_path
    suffix = suffix_of lock_path
    begin
      still_locked = false
      suffixes = children(lock_node).map { |c| suffix_of c }
      lowest_suffix = suffixes.sort.first
      unless suffix.to_i <= lowest_suffix.to_i
        next_path = "#{lock_node}/#{name}-#{lowest_suffix}"
        still_locked = exists? next_path
      end
    end while still_locked
  end

  def unlock name
    delete @locks[name]
    return self
  end
end