基于ZooKeeper的分布式锁实现

非阻塞式锁和阻塞式锁

非阻塞式锁

对于一个资源（可以是方法、属性等），如果线程造访时判断锁已经被其他线程占有，则不会停留，直接返回

阻塞式锁

和非阻塞式锁的区别在于，发现锁被占有之后，会等待锁的释放，直到自己获得了锁

分布式锁

说明

一般的非阻塞式锁和阻塞式锁，指的是单一应用实例下的场景，而为了提高服务的可靠性，通常会将服务部署到多台机器上，这时的服务实例就是多个了，相应的，原来的无论是非阻塞式还是阻塞式锁，也要升级为对应的分布式锁才能满足多实例场景

非阻塞分布式锁

应用场景

如kafka生产者，在多实例下要确保数据不会重复发送，即数据只能有一个实例在生产，这种场景使用非阻塞分布式锁就显得十分契合了

实现思路

在所有实例启动时，通过让他们去创建一个master临时节点，创建成功的即为master，这时程序中只要让判断为master的服务实例去生产消息，就能避免数据的重复发送了，同时利用zookeeper临时节点的特点，即当服务实例和zookeeper的连接断开时，它创建的对应的临时节点会被删除，让未创建成功的去监听当前master，一旦被删除，都再次尝试去创建master

阻塞分布式锁

应用场景

如果代码中有对状态的维护，且需保证状态的修改的先后顺序，保证在被一个线程修改时其他线程不允许操作，同时又不能丢弃之后的修改操作时，就可以使用阻塞分布式锁了

实现思路

利用zookeeper的临时有序节点的特点，有序性指的是根据创建节点的先后顺序，生成节点的序列号将会依次增加，所以，可以将要锁住的资源对应为一个临时有序节点，所有的调用者中序号最小的，说明创建的最早，将其判定为锁的owner，然后在它的业务逻辑结束之后，主动去释放锁，即删除它创建的节点，而在此之前其他序列号的对应线程实例会监听该owner节点，所以在owner释放锁之后，再来一轮序列号大小判断又可得出新的owner

注意事项

阻塞式锁最需要关心的是锁的释放，一旦未能释放，将引起永久阻塞，所以在代码中要注意异常的处理，因为异常下很可能释放锁的代码没有被执行，最好提供一个锁释放的保底操作，比如全局异常捕捉并进行锁释放，或者提供一个 spring aop 的后置操作，并在其中进行锁的释放

代码

ZooKeeperConfig

package zk;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.io.IOException;
import java.util.concurrent.CountDownLatch;


@Slf4j
@Configuration
public class ZooKeeperConfig {
    @Value("${zookeeper.address}")
    private String zkAddress;
    @Value("${zookeeper.session.timeout}")
    private int zkTimeout;

    @Bean
    public ZooKeeper zooKeeper() {
        ZooKeeper zooKeeper;
        try {
            final CountDownLatch latch = new CountDownLatch(1);
            zooKeeper = new ZooKeeper(
                    this.zkAddress, 
                    this.zkTimeout, 
                    watchedEvent -> {
                        if (Watcher.Event.KeeperState.SyncConnected == watchedEvent.getState()) {
                            latch.countDown();
                        }
                    }
            );
            latch.await();
            log.info("Connect zookeeper finished, the state is {}", zooKeeper.getState());
        }catch (IOException | InterruptedException i) {
            log.error("Connect zookeeper failed.", i);
            throw new ZKConnectionException();
        }
        return zooKeeper;
    }
}

ZooKeeperService

package zk;

import java.util.Optional;

public interface ZooKeperService {
    /**
     * 概述：当前实例是否是master
     * 类型：非阻塞分布式锁
     * @return 布尔
     */
    boolean isMaster();

    /**
     * 概述：一直等到自己成为owner才结束本次调用
     * 类型：阻塞式分布式锁
     * @param node 需要分布式锁进行锁定的资源 对应的锁节点目录 如锁定test方法，给他在zk中创建一个 /test 节点
     * @return 本次调用所创建的临时有序的锁节点，如 /foo/bar/locks/test/lock0000000001
     */
    String becomeOwner(String node);

    /**
     * 概述：在becomeOwner基础上添加了等待时间，即如果一定时间内依旧未获得锁，将直接返回一个空结果，而不是继续等待
     * 类型：阻塞式分布式锁
     * @param node 同becomeOwner
     * @param waitSeconds 指定的等待时间，单位秒
     * @return 含义同becomeOwner，只不过可能在未获取锁的情况下而返回一个空结果
     */
    Optional<String> try2becomeOwner(String node, int waitSeconds);

    /**
     * 概述：提供一个主动释放锁的方法
     * @param lockNode 对应becomeOwner/try2becomeOwner返回的节点值
     */
    void releaseLock(String lockNode);
}

ZooKeeperServiceImpl

package zk;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.PropertySource;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.nio.charset.StandardCharsets;
import java.util.Comparator;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

@Slf4j
@Service
@PropertySource(value = "classpath:application.properties")
public class ZooKeeperServiceImpl implements ZooKeperService {
    @Autowired
    private ZooKeeper zooKeeper;

    //主节点，如：/foo/bar/master
    @Value("${zookeeper.master.node}")
    private String masterNode;

    //将所有资源对应的zk节点放在此节点下，方便管理，如：/foo/bar/locks
    @Value("${zookeeper.locks.node}")
    private String locksNode;

    //用于标记不同服务实例（每个服务都不同，安全策略允许下可以用机器ip，便于辨识）
    @Value("${zookeeper.service.id}")
    private String serviceId;

    private boolean isMaster;

    /**
     * 通过创建master临时节点来竞选master
     */
    @PostConstruct
    private void electMaster() {
        this.masterNode = StringUtils.stripEnd(this.masterNode, "/"); //zookeeper节点路径不允许带/
        String[] nodes = StringUtils.split(this.masterNode, "/");
        /*先创建master临时节点的存放目录，如/foo/bar/master 中的 /foo/bar*/
        String path = StringUtils.join(ArrayUtils.subarray(nodes, 0, nodes.length - 1), "/");
        createNode(null, path);
        /*创建master节点 /foo/bar/master*/
        createMasterNode();
    }

    /**
     * 创建 /foo/bar/locks 节点
     */
    @PostConstruct
    private void createLocksNode() {
        this.locksNode = StringUtils.stripEnd(this.locksNode, "/");
        createNode(null, this.locksNode);
    }

    /**
     * 创建永久节点（自动联级创建）
     *
     * @param parentNode 父节点，即在该父节点下进行（联级）创建
     * @param node       待创建的节点
     */
    private void createNode(String parentNode, String node) {
        parentNode = null == parentNode ? StringUtils.EMPTY : parentNode;
        parentNode = StringUtils.stripEnd(parentNode, "/");
        node = StringUtils.strip(node, "/");
        String[] nodes = StringUtils.split(node, "/");
        StringBuilder tmpNode = new StringBuilder(parentNode);
        for (int i = 0; i < nodes.length; i++) {
            this.zooKeeper.create(
                    tmpNode.append("/").append(nodes[i]).toString(),
                    null,
                    ZooDefs.Ids.OPEN_ACL_UNSAFE,
                    CreateMode.PERSISTENT,
                    (code, path, ctx, name) -> {
                        if (KeeperException.Code.get(code) == KeeperException.Code.OK) {
                            log.info("Create node ({}) success.", path);
                        } else if (KeeperException.Code.get(code) == KeeperException.Code.NODEEXISTS) {
                            log.warn("Create failed, the node ({}) already exists.", path);
                        } else {
                            log.error("Create node ({}) failed: {}", path, code);
                        }
                    },
                    null
            );
        }
    }

    /**
     * 创建主节点
     */
    private void createMasterNode() {
        this.zooKeeper.create(
                this.masterNode,
                this.serviceId.getBytes(StandardCharsets.UTF_8),
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL,
                (code, path, ctx, name) -> {
                    switch (KeeperException.Code.get(code)) {
                        case NODEEXISTS: //说明master已存在
                            byte[] data;
                            try {
                                data = this.zooKeeper.getData(this.masterNode, Boolean.FALSE, new Stat());
                            } catch (KeeperException | InterruptedException e) {
                                log.error("Elect failed, get the master data error.", e);
                                throw new ZKGetDataException();
                            }
                            if (null != data) {
                                log.info("Elect failed, the master already exists: {}", new String(data, StandardCharsets.UTF_8));
                            }
                            //监听当前master节点
                            listenMaster();
                            break;
                        case OK: //说明创建成功，称为master
                            this.isMaster = Boolean.TRUE;
                            log.info("Elect success, {} become the master.", this.serviceId);
                            break;
                        default:
                            log.error("Elect failed, abnormal state ({}) when creating master node.", code);
                    }
                },
                null
        );
    }

    /**
     * 创建锁节点
     *
     * @param lockNode        某资源对应的锁节点目录，如 test方法 /foo/bar/locks/test/lock
     * @param isOwner         当前是否获取锁
     * @param currentLockNode 本次调用所创建的锁节点，如 /foo/bar/locks/test/lock0000000002
     */
    private void createLockNode(String lockNode, AtomicBoolean isOwner, AtomicReference<String> currentLockNode) {
        this.zooKeeper.create(
                lockNode,
                null,
                ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL_SEQUENTIAL,
                (code, path, ctx, name) -> {
                    if (KeeperException.Code.get(code) == KeeperException.Code.OK) { //创建成功
                        log.info("Create lock node ({}) success.", name);
                        currentLockNode.set(name);
                        String currentLockOwner = getLockOwner(lockNode);
                        if (StringUtils.isNotEmpty(currentLockOwner)) {
                            if (StringUtils.equals(name, currentLockOwner)) { //说明自己就是owner
                                isOwner.set(Boolean.TRUE);
                            } else { //需要对当前owner进行监听
                                listenLockOwner(currentLockOwner, lockNode, name, isOwner);
                            }
                        }
                    }
                },
                null
        );
    }

    /**
     * 监听某资源对应的锁节点目录中的owner
     *
     * @param lockOwner       资源对应的锁节点目录中的owner 如 /foo/bar/locks/test/lock0000000001
     * @param lockNode        资源对应的锁节点目录 如 /foo/bar/locks/test/lock
     * @param currentLockNode 本次创建的锁节点 如 /foo/bar/locks/test/lock0000000002
     * @param isOwner         当前是否获取锁
     */
    private void listenLockOwner(String lockOwner, String lockNode, String currentLockNode, AtomicBoolean isOwner) {
        this.zooKeeper.exists(
                lockOwner,
                watchedEvent -> {
                    String currentLockOwner = getLockOwner(lockNode); //lockOwner可能有变动
                    if (watchedEvent.getType() == Watcher.Event.EventType.NodeDeleted) {
                        // 先判断创建的节点是否还存在（可能已经被提前删除）
                        Stat exists = null;
                        try {
                            exists = this.zooKeeper.exists(currentLockNode, Boolean.FALSE);
                            if (null == exists) {
                                log.warn("The lock node ({}) doesn't exists, stop listening for it", currentLockNode);
                            }
                        } catch (KeeperException | InterruptedException e) {
                            log.error("Listen lock owner ({}) failed, there is an error when check node({}) exists", lockOwner, currentLockNode, e);
                        }
                        if (StringUtils.isNotEmpty(currentLockOwner) && null != exists) {
                            if (StringUtils.equals(currentLockNode, currentLockOwner)) {
                                isOwner.set(Boolean.TRUE);
                            } else {
                                listenLockOwner(currentLockOwner, lockNode, currentLockNode, isOwner);
                            }
                        }
                    } else {
                        //监听都是一次性的，所以需要继续监听
                        listenLockOwner(currentLockOwner, lockNode, currentLockNode, isOwner);
                    }
                },
                null,
                null
        );
    }

    /**
     * 获取对应资源的锁节点目录中的owner
     *
     * @param lockNode 对应资源的锁节点目录 如 /foo/bar/locks/test/lock
     * @return 最小序列号节点
     */
    private String getLockOwner(String lockNode) {
        // /foo/bar/locks/test/lock -> /foo/bar/locks/test
        String lockNodePath = StringUtils.stripEnd(lockNode, "/lock");
        String lockOwner;
        try {
            lockOwner = this.zooKeeper.getChildren(lockNodePath, Boolean.FALSE)
                    .parallelStream()
                    .min(Comparator.comparingInt(name -> Integer.parseInt(StringUtils.stripStart(name, "lock"))))
                    .orElse(StringUtils.EMPTY);
        } catch (KeeperException | InterruptedException e) {
            log.error("Get lock owner ({}) failed, there is an error when get children", lockNodePath, e);
            throw new ZKGetChildrenException();
        }
        if (StringUtils.isEmpty(lockOwner)) { //可能被提前删除，如设置了等待时间但没有获取锁情景下
            log.warn("Get lock owner ({}) failed, the lock owner doesn't exists", lockNodePath);
            return lockOwner;
        } else {
            log.info("Get lock owner ({}) success, the current lock owner is {}", lockNodePath, lockOwner);
            return lockNodePath + "/" + lockOwner;
        }
    }

    /**
     * 监听/foo/bar/master主节点
     */
    private void listenMaster() {
        this.zooKeeper.exists(
                this.masterNode,
                watchedEvent -> {
                    if (watchedEvent.getType() == Watcher.Event.EventType.NodeDeleted) {
                        electMaster();
                    }
                },
                (code, path, ctx, stat) -> {
                    switch (KeeperException.Code.get(code)) {
                        case OK:
                            break;
                        case NONODE:
                            log.warn("The result of listen-master-exists is NONODE, try to elect...");
                            electMaster();
                            break;
                        default:
                            log.error("Listen master-exists failed, there is an abnormal state ({}).", code);
                            throw new ZKListenException();
                    }
                },
                null
        );
    }

    @Override
    public boolean isMaster() {
        return this.isMaster;
    }

    @Override
    public String becomeOwner(String node) {
        Optional<String> lockNode = try2becomeOwner(node, -1);
        if (!lockNode.isPresent()) {
            log.error("Become owner error, after becoming the lock owner, no lock owner information is obtained");
            throw new ZKExistsException();
        }
        return lockNode.get();
    }

    @Override
    public Optional<String> try2becomeOwner(String node, int waitSeconds) {
        //存放当前是否获取锁
        AtomicBoolean isOwner = new AtomicBoolean(Boolean.FALSE);

        //存放本次创建的锁节点
        AtomicReference<String> currentLockNode = new AtomicReference<>(StringUtils.EMPTY);

        //创建资源对应的锁节点目录
        createNode(this.locksNode, node);

        //创建锁节点并等待称为owner
        final CountDownLatch latch = new CountDownLatch(1);
        new Thread(() -> {
            String lockNode = this.locksNode + "/" + StringUtils.strip(node, "/") + "/lock";
            createLockNode(lockNode, isOwner, currentLockNode);
            //等待上位中...
            while (true) {
                if (isOwner.get()) {
                    latch.countDown();
                    break;
                }
            }
        }).start();
        try {
            if (0 < waitSeconds) {
                latch.await(waitSeconds, TimeUnit.SECONDS);
            } else {
                latch.await();
            }
        } catch (InterruptedException i) {
            log.error("Try to become owner failed, CountDownLatch.await error.", i);
            throw new ZKLatchAwaitException();
        }
        if (isOwner.get()) {
            return Optional.of(currentLockNode.get());
        } else {
            releaseLock(currentLockNode.get()); //未获取，则先将创建的锁节点删除
            return Optional.empty();
        }
    }

    @Override
    public void releaseLock(String lockNode) {
        if (StringUtils.isEmpty(lockNode)) return;
        this.zooKeeper.delete(
                lockNode,
                -1,
                (code, path, obj) -> {
                    switch (KeeperException.Code.get(code)) {
                        case OK:
                            log.info("Release lock node ({}) success.", lockNode);
                            break;
                        case NONODE:
                            log.warn("Release lock node ({}) failed: NONODE.", lockNode);
                            break;
                        default:
                            log.error("Release lock node ({}) failed: {}", lockNode, code);
                            throw new ZKDeleteException();

                    }
                },
                null
        );
    }
}

自定义异常模板

public class ZKExistsException extends RuntimeException {
}