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

Go to download
Show more of this group Show more artifacts with this name
Show all versions of zookeeper-utils Show documentation
Bluejeans Zookeeper utilities
The newest version!
package com.bluejeans.utils.zookeeper;

import static java.util.Collections.emptyList;

import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

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 path, ZKLock zkLock);

        void lockReleased(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();
        }
    }

    /**
     * @return the zkClient
     */
    public CuratorFramework getZkClient() {
        return zkClient;
    }

    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 cancelAll(final Collection locks) {
        locks.forEach(l -> MetaUtil.runNestedMethodsSilently(l,
                "$lock..$semaphore..$lock..$internals..~notifyFromWatcher;;$lock..release"));
    }

    public void lockSomeAsync(final Map locks, final int count, final LockListener listener) {
        new Thread(() -> lockSome(locks, count, listener)).start();
    }

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

    public Map lockSomeItems(final String pathPrefix, final List items, final int count,
            final LockListener listener) {
        return lockSome(
                items.stream()
                        .collect(Collectors.toMap(i -> pathPrefix + i, i -> new ZKLock(zkClient, pathPrefix + i))),
                count, listener);
    }

    public Map lockSome(final Map locks, final int count, final LockListener listener) {
        final Map requiredLocks = new HashMap();
        final String lock = "lock";
        int locksNeeded = count;
        final AtomicBoolean allDone = new AtomicBoolean();
        if (locks.size() < count) {
            locksNeeded = locks.size();
        }
        final CountDownLatch latch = new CountDownLatch(locksNeeded);
        for (final String path : locks.keySet()) {
            locks.put(path, locks.get(path));
            new Thread() {
                @Override
                public void run() {
                    try {
                        if (latch.getCount() > 0) {
                            locks.get(path).lock();
                            synchronized (lock) {
                                if (allDone.get()) {
                                    locks.get(path).release();
                                } else {
                                    if (locks.get(path).isLockObtained()) {
                                        locks.get(path).setLocked(true);
                                        if (listener != null) {
                                            listener.lockObtained(path, locks.get(path));
                                        }
                                    }
                                }
                            }
                        }
                    } catch (final Exception ex) {
                        ex.printStackTrace();
                    } finally {
                        latch.countDown();
                    }
                }
            }.start();
        }
        try {
            latch.await();
        } catch (final InterruptedException e1) {
            // nothing to do
        }
        synchronized (lock) {
            requiredLocks.putAll(EntryStream.of(locks).filter(e -> e.getValue().isLocked()).limit(locksNeeded).toMap());
            cancelAll(EntryStream.of(locks).filter(e -> !requiredLocks.containsKey(e.getKey())).values()
                    .filter(l -> !l.isLocked()).collect(Collectors.toList()));
            EntryStream.of(locks).filter(e -> !requiredLocks.containsKey(e.getKey())).values().filter(l -> l.isLocked())
                    .forEach(l -> {
                        if (listener != null) {
                            listener.lockReleased(l.getLockPath(), l);
                        }
                        l.release();
                    });
            allDone.set(true);
        }
        return requiredLocks;
    }

}