zookeeper学习五-zookeeper实现分布式锁

1 手写分布式锁

根据 ZK简介 一文中描述的步骤，我们来手动实现分布式锁

package com.jms.zookeeper;
 
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;
import org.junit.Test;
 
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
 
/**
 * @author huwenlong
 * @date 2021/5/1 20:37
 */
public class DistributeLock {
 
    private CountDownLatch connectedSemaphore = new CountDownLatch(1);
 
    private ZooKeeper zooKeeper;
 
    private ZKWatcher watcher = new ZKWatcher(connectedSemaphore);
 
    private ThreadLocal<String> currentLock = new ThreadLocal<>();
 
    private ReentrantLock reentrantLock = new ReentrantLock();
 
    private Map<String, Condition> pathConditionMap = new HashMap<>();
 
    private final Watcher lockWatcher = event -> {
        if (event.getType() == Watcher.Event.EventType.NodeDeleted) {
            // reentrantLock精确唤醒等待该节点的线程
            reentrantLock.lock();
            Condition condition = pathConditionMap.remove(event.getPath());
            if (Objects.nonNull(condition)) {
                condition.signal();
            }
            reentrantLock.unlock();
        }
    };
 
    private final String ROOT_LOCK_PATH = "/distributed_lock";
 
    public DistributeLock() throws Exception {
        zooKeeper = new ZooKeeper(Constant.ZK_HOST, Constant.SESSION_TIMEOUT, watcher);
        //进行阻塞
        connectedSemaphore.await();
        Stat stat = zooKeeper.exists(ROOT_LOCK_PATH, false);
        if (Objects.isNull(stat)) {
            // 创建根节点
            zooKeeper.create(ROOT_LOCK_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
    }
 
    /**
     * 上锁
     *
     * @param lockName 通常是业务名
     */
    public void lock(String lockName) throws Exception {
        // 创建临时节点
        String currentPath = zooKeeper.create(ROOT_LOCK_PATH + "/" + lockName, new byte[0],
                ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        currentLock.set(currentPath);
        tryLock(currentPath);
    }
 
    private void tryLock(String currentPath) throws Exception {
        // 获取根节点下的所有子节点并排序
        List<String> locks = zooKeeper.getChildren(ROOT_LOCK_PATH, false);
        Collections.sort(locks);
        int index = locks.indexOf(currentPath.substring(ROOT_LOCK_PATH.length() + 1));
        if (index == 0) {
            // 如果当前线程创建的节点是第一个节点则表示获取到锁
            System.out.println(Thread.currentThread().getName() + " get lock, " + currentPath);
            return;
        }
        // 监听上一个节点数据变化并阻塞
        String beforePath = ROOT_LOCK_PATH + "/" + locks.get(index - 1);
        Stat stat = zooKeeper.exists(beforePath, lockWatcher);
        if (Objects.isNull(stat)) {
            // 如果上个节点已删除，再次抢锁
            tryLock(currentPath);
        } else {
            // 否则阻塞等待
            reentrantLock.lock();
            System.out.println(Thread.currentThread().getName() + " miss lock, currentPath:" + currentPath +
                    ", waitPath:" + beforePath);
            Condition condition = reentrantLock.newCondition();
            pathConditionMap.put(beforePath, condition);
            condition.await();
            reentrantLock.unlock();
            // 再次抢锁
            tryLock(currentPath);
        }
    }
 
    /**
     * 解锁
     */
    public void unLock() throws Exception {
        String currentPath = currentLock.get();
        zooKeeper.delete(currentPath, -1);
        System.out.println("unlock path:" + currentPath);
    }
 
    @Test
    public void testNoLock() throws InterruptedException {
        int[] num = new int[1];
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                int currentNum = num[0];
                currentNum++;
                try {
                    // 睡若干毫秒
                    Thread.sleep((long) Math.random() * 10);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                num[0] = currentNum;
            }).start();
        }
        Thread.sleep(1000);
        System.out.println(num[0]);
    }
 
    @Test
    public void testLock() throws Exception {
        DistributeLock lock = new DistributeLock();
        int[] num = new int[1];
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                try {
                    lock.lock("incr");
                    int currentNum = num[0];
                    currentNum++;
                    num[0] = currentNum;
                    // 睡若干毫秒
                    Thread.sleep((long) Math.random() * 10);
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    try {
                        lock.unLock();
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            }).start();
        }
        Thread.sleep(2000);
        System.out.println(num[0]);
    }
}

输出

Thread-2 get lock, /distributed_lock/incr0000000750
Thread-7 miss lock, currentPath:/distributed_lock/incr0000000752, waitPath:/distributed_lock/incr0000000751
Thread-5 miss lock, currentPath:/distributed_lock/incr0000000754, waitPath:/distributed_lock/incr0000000753
Thread-1 miss lock, currentPath:/distributed_lock/incr0000000758, waitPath:/distributed_lock/incr0000000757
Thread-6 miss lock, currentPath:/distributed_lock/incr0000000756, waitPath:/distributed_lock/incr0000000755
Thread-3 miss lock, currentPath:/distributed_lock/incr0000000757, waitPath:/distributed_lock/incr0000000756
Thread-8 miss lock, currentPath:/distributed_lock/incr0000000753, waitPath:/distributed_lock/incr0000000752
Thread-4 miss lock, currentPath:/distributed_lock/incr0000000755, waitPath:/distributed_lock/incr0000000754
Thread-0 miss lock, currentPath:/distributed_lock/incr0000000751, waitPath:/distributed_lock/incr0000000750
Thread-9 miss lock, currentPath:/distributed_lock/incr0000000759, waitPath:/distributed_lock/incr0000000758
unlock path:/distributed_lock/incr0000000750
Thread-0 get lock, /distributed_lock/incr0000000751
unlock path:/distributed_lock/incr0000000751
Thread-7 get lock, /distributed_lock/incr0000000752
unlock path:/distributed_lock/incr0000000752
Thread-8 get lock, /distributed_lock/incr0000000753
unlock path:/distributed_lock/incr0000000753
Thread-5 get lock, /distributed_lock/incr0000000754
unlock path:/distributed_lock/incr0000000754
Thread-4 get lock, /distributed_lock/incr0000000755
unlock path:/distributed_lock/incr0000000755
Thread-6 get lock, /distributed_lock/incr0000000756
unlock path:/distributed_lock/incr0000000756
Thread-3 get lock, /distributed_lock/incr0000000757
unlock path:/distributed_lock/incr0000000757
Thread-1 get lock, /distributed_lock/incr0000000758
unlock path:/distributed_lock/incr0000000758
Thread-9 get lock, /distributed_lock/incr0000000759
unlock path:/distributed_lock/incr0000000759
10

2 利用curator工具类实现分布式锁

2.1 代码实现

    @Test
    public void testDistributeLock() throws InterruptedException {
        String lockPath = "/lock";
        InterProcessMutex lock = new InterProcessMutex(curatorFramework, lockPath);
 
        ExecutorService executor = Executors.newCachedThreadPool();
        for (int i = 0; i < 5; i++) {
            //创建锁，为可重入锁，即是获锁后，还可以再次获取，本例以此为例
            executor.execute(() -> {
                try {
                    //获取锁
                    lock.acquire();
                    log.info(Thread.currentThread() + " acquire lock");
                } catch (Exception ignore) {
                } finally {
                    //释放锁
                    try {
                        lock.release();
                        log.info(Thread.currentThread() + " release lock");
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            });
        }
        Thread.sleep(10000L);
    }
 
    @Test
    public void testDistributeReadWriteLock() throws InterruptedException {
        String lockPath = "/lock";
        InterProcessReadWriteLock lock = new InterProcessReadWriteLock(curatorFramework, lockPath);
        ExecutorService executor = Executors.newCachedThreadPool();
        for (int i = 0; i < 5; i++) {
            //创建锁，为可重入锁，即是获锁后，还可以再次获取，本例以此为例
            executor.execute(() -> {
                try {
                    //获取锁
                    lock.writeLock().acquire();
                    log.info(Thread.currentThread() + " acquire write lock");
                } catch (Exception ignore) {
                } finally {
                    //释放锁
                    try {
                        lock.writeLock().release();
                        log.info(Thread.currentThread() + " release write lock");
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            });
        }
        Thread.sleep(10000L);
    }

2.2 源码分析

有了手写分布式锁的基础，这里再看InterProcessMutex源码则不难理解

2.2.1 获取锁

（1）主流程

    public void acquire() throws Exception
    {
        if ( !internalLock(-1, null) )
        {
            throw new IOException("Lost connection while trying to acquire lock: " + basePath);
        }
    }


   private boolean internalLock(long time, TimeUnit unit) throws Exception
    {
        /*
           Note on concurrency: a given lockData instance
           can be only acted on by a single thread so locking isn't necessary
        */

        Thread currentThread = Thread.currentThread();
        // 获取当前线程的锁信息，threadData：ConcurrentMap<Thread, LockDate>
        LockData lockData = threadData.get(currentThread);
        if ( lockData != null )
        {
            // 如果当前锁信息存在则表示重入，重入次数+1，抢锁成功
            lockData.lockCount.incrementAndGet();
            return true;
        }
        // 获取锁path
        String lockPath = internals.attemptLock(time, unit, getLockNodeBytes());
        if ( lockPath != null )
        {
            // 获取锁成功后给添加当前线程锁信息
            LockData newLockData = new LockData(currentThread, lockPath);
            threadData.put(currentThread, newLockData);
            return true;
        }

        return false;
    }

（2）获取锁path

    String attemptLock(long time, TimeUnit unit, byte[] lockNodeBytes) throws Exception
    {
        final long      startMillis = System.currentTimeMillis();
        final Long      millisToWait = (unit != null) ? unit.toMillis(time) : null;
        final byte[]    localLockNodeBytes = (revocable.get() != null) ? new byte[0] : lockNodeBytes;
        int             retryCount = 0;

        String          ourPath = null;
        boolean         hasTheLock = false;
        boolean         isDone = false;
        while ( !isDone )
        {
            isDone = true;

            try
            {
                // 创建并返回临时子节点
                ourPath = driver.createsTheLock(client, path, localLockNodeBytes);
                // 利用子节点循环抢锁，原理和手写分布式锁类似，有兴趣的同学可以研究一下源码做了什么优化
                hasTheLock = internalLockLoop(startMillis, millisToWait, ourPath);
            }
            catch ( KeeperException.NoNodeException e )
            {
                // gets thrown by StandardLockInternalsDriver when it can't find the lock node
                // this can happen when the session expires, etc. So, if the retry allows, just try it all again
                if ( client.getZookeeperClient().getRetryPolicy().allowRetry(retryCount++, System.currentTimeMillis() - startMillis, RetryLoop.getDefaultRetrySleeper()) )
                {
                    isDone = false;
                }
                else
                {
                    throw e;
                }
            }
        }

        if ( hasTheLock )
        {
            return ourPath;
        }

        return null;
    }

2.2.2 释放锁

    public void release() throws Exception
    {
        Thread currentThread = Thread.currentThread();
        // 获取当前线程锁信息
        LockData lockData = threadData.get(currentThread);
        if ( lockData == null )
        {
            throw new IllegalMonitorStateException("You do not own the lock: " + basePath);
        }
        // 当前线程锁重入次数-1
        int newLockCount = lockData.lockCount.decrementAndGet();
        if ( newLockCount > 0 )
        {
            // 如果还有重入次数直接返回
            return;
        }
        if ( newLockCount < 0 )
        {
            throw new IllegalMonitorStateException("Lock count has gone negative for lock: " + basePath);
        }
        try
        {
            // 如果重入次数=0则释放锁(实际上就是删除临时节点)
            internals.releaseLock(lockData.lockPath);
        }
        finally
        {
            // 内存中移除当前线程信息
            threadData.remove(currentThread);
        }
    }