Java：使用ZooKeeper实现分布式锁

本文提供了使用ZooKeeper实现分布式锁的一个实现。与网上其它实现方式的主要区别在于，可以使用同一个客户端创建针对多个不同路径的锁。

1、锁管理器

锁管理器内部包含一个Zookeeper实例，创建完成后可以根据不同路径创建不同的锁。

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

public class ReentrantZKLockMgr {

  private ZooKeeper zooKeeper = null;
  private static byte[] data = new byte[0];

  public ReentrantZKLockMgr(String connectString, int sessionTimeout) throws Exception {
    zooKeeper = new ZooKeeper(connectString, sessionTimeout, new SessionWatcher());
  }

  public ReentrantZKLock makeLock(String categoryPath) throws Exception {
    Stat stat = zooKeeper.exists(categoryPath, false);
    if (stat == null) {
      zooKeeper.create(categoryPath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
    }

    return new ReentrantZKLock(zooKeeper, categoryPath);
  }

  public void close() throws Exception {
    zooKeeper.close();
  }
}

class SessionWatcher implements Watcher {

  public void process(WatchedEvent event) {

    System.out.println("State=" + event.getState());
    System.out.println("Type=" + event.getType());
  }

}

2、锁监视器

监听节点删除事件，当客户端监听的前一个节点被删除时，唤醒当前客户端。

import java.util.concurrent.CountDownLatch;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;

public class ReentrantZKLockWatcher implements Watcher {
  private CountDownLatch latch;

  public ReentrantZKLockWatcher(CountDownLatch latch) {
    this.latch = latch;
  }

  public void process(WatchedEvent event) {
    if (Watcher.Event.EventType.NodeDeleted == event.getType()) {
      latch.countDown();
    }
  }
}

3、同步锁

提供锁定与解锁功能

import java.util.List;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class ReentrantZKLock {

  private static final Logger LOG = LoggerFactory.getLogger(ReentrantZKLock.class);

  private static byte[] data = new byte[0];
  private static String lockNode = "lock_";

  private ZooKeeper zooKeeper;
  private String categoryPath;
  private String nodePath;

  private CountDownLatch latch = new CountDownLatch(1);

  public ReentrantZKLock(ZooKeeper zooKeeper, String categoryPath) {
    this.zooKeeper = zooKeeper;
    this.categoryPath = categoryPath;
  }

  public void lock() {
    try {
      String rawNodePath = categoryPath + "/" + lockNode;
      nodePath = zooKeeper.create(rawNodePath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
      System.out.println("nodePath=" + nodePath);

      List<String> nodeIdList = zooKeeper.getChildren(categoryPath, true);
      System.out.println(nodeIdList);

      TreeSet<String> nodeIdSet = new TreeSet<String>(nodeIdList);
      String firstNodeId = nodeIdSet.first();

      String nodeId = nodePath.substring((categoryPath + "/").length());
      if (firstNodeId == nodeId) {
        return;
      }

      String preNodeId = nodeIdSet.lower(nodeId);
      if (preNodeId != null) {
        String preNodePath = categoryPath + "/" + preNodeId;
        Stat stat = zooKeeper.exists(preNodePath, new ReentrantZKLockWatcher(latch));

        if (stat != null) {
          latch.await();

          System.out.println("nodeId=" + nodeId + ", preNodeId=" + preNodeId);
        }
      }
    } catch (Exception e) {
      LOG.error("lock " + categoryPath, e);
    }
  }

  public void unlock() {
    try {
      zooKeeper.delete(nodePath, -1);
    } catch (Exception e) {
      LOG.error("unlock " + categoryPath, e);
    }
  }
}

4、测试代码

import com.jjhgame.test.ReentrantZKLock;
import com.jjhgame.test.ReentrantZKLockMgr;

public class ReentrantZKLockTester implements Runnable {

  private String path = "/test1";
  private ReentrantZKLockMgr lockmgr;

  public ReentrantZKLockTester() {
    String connectString = "127.0.0.1:2181";
    int sessionTimeout = 5000;

    try {
      lockmgr = new ReentrantZKLockMgr(connectString, sessionTimeout);

    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  public void run() {
    try {
      ReentrantZKLock lock = lockmgr.makeLock(path);
      lock.lock();
      try {
        Thread.sleep(5 * 1000);
        System.out.println(Thread.currentThread().getId());
      } finally {
        lock.unlock();
      }
      lockmgr.close();
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  public static void main(String[] args) {
    Thread[] ths = new Thread[3];
    for (int i = 0; i < ths.length; i++) {
      ths[i] = new Thread(new ReentrantZKLockTester());
    }
    for (int i = 0; i < ths.length; i++) {
      ths[i].start();
    }
  }
}

参考文档

使用zookeeper实现分布式锁