Zookeeper实现分布式锁实战

1. 实现原理图

2. 系统环境

JDK1.8 + zookeeper3.8.3 + 依赖如下

<!-- curator-framework -->
<dependency>
    <groupId>org.apache.curator</groupId>
    <artifactId>curator-framework</artifactId>
    <version>5.1.0</version>
</dependency>
<!-- curator-recipes -->
<dependency>
    <groupId>org.apache.curator</groupId>
    <artifactId>curator-recipes</artifactId>
    <version>5.1.0</version>
</dependency>

3. 代码实现

import com.google.common.collect.Lists;
import org.apache.commons.lang3.StringUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.CuratorCache;
import org.apache.curator.framework.recipes.cache.CuratorCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;

import java.io.Closeable;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

/**
 * 分布式锁对象
 */
public class ZkLockClient implements Closeable {
    /**
     * 注意：首先在zookeepe服务里面手动添加"/LOCK"节点
     */
    private static final String LOCK_PATH = "/LOCK";
    /**
     * zookeeper客户端
     */
    private CuratorFramework client;
    /**
     * 当前节点
     */
    private String currentPath = null;
    /**
     * 当前节点的上一个节点
     */
    private String beforePath = null;
    /**
     * 线程同步工具，为了唤醒当前线程使用
     */
    private CountDownLatch countDownLatch;

    /**
     * 对象初始化
     */
    public ZkLockClient() {
        // zookeeper服务连接和启动
        client = CuratorFrameworkFactory.builder()
                .connectString("127.0.0.1:2181")
                .retryPolicy(new ExponentialBackoffRetry(1000, 3))
                .build();
        client.start();
    }

    /**
     * 加锁
     *
     * @throws Exception
     */
    public void getLock() {
        if (currentPath == null) {
            // 在一个LOCK_PATH节点下面添加一个临时顺序节点
            createNode();
        }
        // 获取当前和我竞争的线程
        List<String> tmpPath = getChildren();
        // 重新创建一个集合，因为client.getChildren()获取的是子节点的名称，不是全路径
        List<String> nodeAllPath = Lists.newArrayList();
        for (String oldPath : tmpPath) {
            nodeAllPath.add(LOCK_PATH + "/" + oldPath);
        }
        //如果当前只有我一个或者我就是第一个，直接获取锁
        if (nodeAllPath.size() == 1 || StringUtils.equals(currentPath, nodeAllPath.get(0))) {
            return;
        }
        // 获取自己上一个临时节点
        int currentIndex = nodeAllPath.indexOf(currentPath);
        // 主要为了后面 awaitLock()监听使用的
        beforePath = nodeAllPath.get(currentIndex - 1);
        // 等待锁，并且等待被唤醒
        awaitLock();
        // 递归的在一次去尝试获取锁
        getLock();
    }

    private void awaitLock() {
        CuratorCache curatorCache = CuratorCache.build(client, beforePath);
        curatorCache.listenable().addListener((type, oldData, data) -> {
            // 我的上一个节点已经完成业务逻辑，节点被删除，被我监听到，然后唤醒我去获取锁
            if (type.name().equals(CuratorCacheListener.Type.NODE_DELETED.name())) {
                countDownLatch.countDown();
            }
        });
        curatorCache.start();
        countDownLatch = new CountDownLatch(1);
        try {
            String name = Thread.currentThread().getName();
            System.out.println(name + ": 锁获取失败，等待被唤醒");
            countDownLatch.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    private List<String> getChildren() {
        try {
            List<String> tmpPath = client.getChildren().forPath(LOCK_PATH);
            // 排好序，变为升序，可以看出我们实现的是一个公平锁
            Collections.sort(tmpPath);
            return tmpPath;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private void createNode() {
        try {
            currentPath = client.create().withMode(CreateMode.EPHEMERAL_SEQUENTIAL).forPath(LOCK_PATH + "/", "".getBytes());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 释放锁
     */
    @Override
    public void close() {
        try {
            if (client != null) {
                String name = Thread.currentThread().getName();
                System.out.println(name + ": 释放了锁");
                client.delete().forPath(currentPath);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}

4. 测试

import java.util.concurrent.TimeUnit;

public class Main {
    public static void main(String[] args) {
        for (int i = 0; i < 10; i++) {
            // 同时有10个线程并发获取锁
            new Thread(new MyRunnable()).start();
        }
    }
}

class MyRunnable implements Runnable {
    @Override
    public void run() {
        // 如何使用zookeeper分布式锁
        try (ZkLockClient zkLockClient = new ZkLockClient()) {
            zkLockClient.getLock();
            String name = Thread.currentThread().getName();
            System.out.println(name + ": 获取到了锁");
            // 业务逻辑处理
            TimeUnit.MILLISECONDS.sleep(500);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        // 因为实现了Closeable接口并重写了方法，所以会自动释放锁
    }
}