【從入門到放棄-ZooKeeper】ZooKeeper實戰-分佈式鎖-升級版

前言

上文【從入門到放棄-ZooKeeper】ZooKeeper實戰-分佈式鎖中，我們通過利用ZooKeeper的臨時節點特性，實現了一個分佈式鎖。
但是是通過輪詢的方式去判斷不斷嘗試獲取鎖，空轉對於CPU還是有一定消耗的，同時，對於多個線程競爭鎖激烈的時候，很容易出現羊羣效應。

爲了解決上面兩個問題。本文來看一下如何實現一個升級版的分佈式鎖。

設計

我們依然實現java.util.concurrent.locks.Lock接口。
和上一文中實現方式不同的是，我們使用ZooKeeper的EPHEMERAL_SEQUENTIAL臨時順序節點。
當首次獲取鎖時，會創建一個臨時節點，如果這個臨時節點末尾數字是當前父節點下同名節點中最小的，則獲取鎖成功。
否則，則監聽上一個數字較大的節點，直到上一個節點被釋放，則再次嘗試獲取鎖成功。這樣可以避免多個線程同時獲取一把鎖造成的競爭。
同時使用了ZooKeeper提供的watch功能，避免了輪詢帶來的CPU空轉。
獲取鎖後使用一個volatile int類型的state進行計數，來實現鎖的可重入機制。

DistributedFairLock

public class DistributedFairLock implements Lock {
    private static Logger logger = LoggerFactory.getLogger(DistributedFairLock.class);

    //ZooKeeper客戶端，進行ZooKeeper操作
    private ZooKeeper zooKeeper;

    //根節點名稱
    private String dir;

    //加鎖節點
    private String node;

    //ZooKeeper鑑權信息
    private List<ACL> acls;

    //要加鎖節點
    private String fullPath;

    //加鎖標識，爲0時表示未獲取到鎖，每獲取一次鎖則加一，釋放鎖時減一。減到0時斷開連接，刪除臨時節點。
    private volatile int state;

    //當前鎖創建的節點id
    private String id;

    //通過CountDownLatch阻塞，直到監聽上一節點被取消，再進行後續操作
    private CountDownLatch countDownLatch;

    /**
     * Constructor.
     *
     * @param zooKeeper the zoo keeper
     * @param dir       the dir
     * @param node      the node
     * @param acls      the acls
     */
    public DistributedFairLock(ZooKeeper zooKeeper, String dir, String node, List<ACL> acls) {
        this.zooKeeper = zooKeeper;
        this.dir = dir;
        this.node = node;
        this.acls = acls;
        this.fullPath = dir.concat("/").concat(this.node);
        init();
    }

    private void init() {
        try {
            Stat stat = zooKeeper.exists(dir, false);
            if (stat == null) {
                zooKeeper.create(dir, null, acls, CreateMode.PERSISTENT);
            }
        } catch (Exception e) {
            logger.error("[DistributedFairLock#init] error : " + e.toString(), e);
        }
    }
}

lock

public void lock() {
    try {
        //加鎖
        synchronized (this) {
            //如果當前未持有鎖
            if (state <= 0) {
                //創建節點
                if (id == null) {
                    id = zooKeeper.create(fullPath, null, acls, CreateMode.EPHEMERAL_SEQUENTIAL);
                }

                //獲取當前路徑下所有的節點
                List<String> nodes = zooKeeper.getChildren(dir, false);
                SortedSet<String> sortedSet = new TreeSet<>();
                for (String node : nodes) {
                    sortedSet.add(dir.concat("/").concat(node));
                }

                //獲取所有id小於當前節點順序的節點
                SortedSet<String> lessSet = ((TreeSet<String>) sortedSet).headSet(id);

                if (!lessSet.isEmpty()) {
                    //監聽上一個節點，就是通過這裏避免多鎖競爭和CPU空轉，實現公平鎖的  
                    Stat stat = zooKeeper.exists(lessSet.last(), new LockWatcher());
                    if (stat != null) {
                        countDownLatch = new CountDownLatch(1);
                        countDownLatch.await();
                    }

                }
            }

            state++;
        }
    } catch (InterruptedException e) {
        logger.error("[DistributedFairLock#lock] error : " + e.toString(), e);
        Thread.currentThread().interrupt();
    } catch (KeeperException ke) {
        logger.error("[DistributedFairLock#lock] error : " + ke.toString(), ke);
        if (!KeeperException.Code.NODEEXISTS.equals(ke.code())) {
            Thread.currentThread().interrupt();
        }
    }
}

tryLock

public boolean tryLock() {
    try {
        synchronized (this) {
            if (state <= 0) {
                if (id == null) {
                    id = zooKeeper.create(fullPath, null, acls, CreateMode.EPHEMERAL_SEQUENTIAL);
                }

                List<String> nodes = zooKeeper.getChildren(dir, false);
                SortedSet<String> sortedSet = new TreeSet<>();
                for (String node : nodes) {
                    sortedSet.add(dir.concat("/").concat(node));
                }


                SortedSet<String> lessSet = ((TreeSet<String>) sortedSet).headSet(id);

                if (!lessSet.isEmpty()) {
                    return false;
                }
            }
            state++;
        }
    } catch (InterruptedException e) {
        logger.error("[DistributedFairLock#tryLock] error : " + e.toString(), e);
        return false;
    } catch (KeeperException ke) {
        logger.error("[DistributedFairLock#tryLock] error : " + ke.toString(), ke);
        if (!KeeperException.Code.NODEEXISTS.equals(ke.code())) {
            return false;
        }
    }
    return true;
}

@Override
public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
    try {
        synchronized (this) {
            if (state <= 0) {
                if (id == null) {
                    id = zooKeeper.create(fullPath, null, acls, CreateMode.EPHEMERAL_SEQUENTIAL);
                }

                List<String> nodes = zooKeeper.getChildren(dir, false);
                SortedSet<String> sortedSet = new TreeSet<>();
                for (String node : nodes) {
                    sortedSet.add(dir.concat("/").concat(node));
                }


                SortedSet<String> lessSet = ((TreeSet<String>) sortedSet).headSet(id);

                if (!lessSet.isEmpty()) {
                    Stat stat = zooKeeper.exists(lessSet.last(), new LockWatcher());
                    if (stat != null) {
                        countDownLatch = new CountDownLatch(1);
                        countDownLatch.await(time, unit);
                    }

                }
            }

            state++;
        }
    } catch (InterruptedException e) {
        logger.error("[DistributedFairLock#tryLock] error : " + e.toString(), e);
        return false;
    } catch (KeeperException ke) {
        logger.error("[DistributedFairLock#tryLock] error : " + ke.toString(), ke);
        if (!KeeperException.Code.NODEEXISTS.equals(ke.code())) {
            return false;
        }
    }
    return true;
}

unlock

public void unlock() {
    synchronized (this) {
        if (state > 0) {
            state--;
        }
        //當不再持有鎖時，刪除創建的臨時節點
        if (state == 0 && zooKeeper != null) {
            try {
                zooKeeper.delete(id, -1);
                id = null;
            } catch (Exception e) {
                logger.error("[DistributedFairLock#unlock] error : " + e.toString(), e);
            }
        }
    }
}

LockWatcher

private class LockWatcher implements Watcher {
    @Override
    public void process(WatchedEvent event) {
        synchronized (this) {
            if (countDownLatch != null) {
                countDownLatch.countDown();
            }
        }
    }
}

總結

上面就是我們改良後，通過臨時順序節點和watch機制實現的公平可重入分佈式鎖。
源代碼可見：aloofJr
通過watch機制避免輪詢帶來的CPU空轉。
通過順序臨時節點避免了羊羣效應。

如果對以上方式有更好的優化方案，歡迎一起討論。

更多文章

見我的博客：https://nc2era.com

written by AloofJr，轉載請註明出處