com.bluejeans.utils.zookeeper.ZkHelper Maven / Gradle / Ivy

Go to download
package com.bluejeans.utils.zookeeper;

import static java.util.Collections.emptyList;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.bluejeans.utils.MetaUtil;

import one.util.streamex.EntryStream;

public class ZkHelper {

    public static interface LockListener {
        void lockObtained(String prefix, String path, ZKLock zkLock);

        void lockReleased(String prefix, String path, ZKLock zkLock);
    }

    private static Logger logger = LoggerFactory.getLogger(ZkHelper.class);

    private final CuratorFramework zkClient;

    public ZkHelper(final CuratorFramework zkClient) {
        super();
        this.zkClient = zkClient;
        if (zkClient.getState() != CuratorFrameworkState.STARTED) {
            zkClient.start();
        }
    }

    public byte[] getData(final String path) {
        try {
            return zkClient.getData().forPath(path);
        } catch (final Exception e) {
            logger.error("Failed to get data for {}", path, e);
        }
        return null;
    }

    public List getChildren(final String path) {
        try {
            return zkClient.getChildren().forPath(path);
        } catch (final Exception e) {
            logger.error("Failed to get children for {}", path, e);
        }
        return emptyList();
    }

    public void lockSomeAsync(final String pathPrefix, final List items, final int count,
            final LockListener listener) {
        new Thread(() -> lockSome(pathPrefix, items, count, listener)).start();
    }

    public Map lockSome(final String pathPrefix, final List items, final int count,
            final LockListener listener) {
        final Map locks = new HashMap();
        final Map requiredLocks = new HashMap();
        int locksNeeded = count;
        if (items.size() < count) {
            locksNeeded = items.size();
        }
        final CountDownLatch latch = new CountDownLatch(locksNeeded);
        for (final String item : items) {
            locks.put(item, new ZKLock(zkClient, pathPrefix + item));
            new Thread() {
                @Override
                public void run() {
                    try {
                        if (latch.getCount() > 0) {
                            locks.get(item).lock();
                            if (listener != null && locks.get(item).isLockObtained()) {
                                listener.lockObtained(pathPrefix, item, locks.get(item));
                                locks.get(item).setLocked(true);
                            }
                        }
                    } finally {
                        latch.countDown();
                    }
                }
            }.start();
        }
        try {
            latch.await();
        } catch (final InterruptedException e1) {
            // nothing to do
        }
        requiredLocks.putAll(EntryStream.of(locks).filter(e -> e.getValue().isLocked()).limit(locksNeeded).toMap());
        EntryStream.of(locks).filter(e -> !requiredLocks.containsKey(e.getKey())).values().filter(l -> !l.isLocked())
                .forEach(l -> MetaUtil.runNestedMethodsSilently(l,
                        "$lock..$semaphore..$lock..$internals..~notifyFromWatcher;;$lock..release"));
        EntryStream.of(locks).filter(e -> !requiredLocks.containsKey(e.getKey())).values().filter(l -> l.isLocked())
                .forEach(l -> {
                    l.release();
                    listener.lockReleased(pathPrefix, l.getLockPath().substring(l.getLockPath().lastIndexOf("/") + 1),
                            l);
                });
        return requiredLocks;
    }

}