com.bluejeans.utils.zookeeper.ZKLock 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 java.util.concurrent.TimeUnit;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.locks.InterProcessSemaphoreMutex;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.slf4j.Logger;

public class ZKLock {

    private static Logger logger = org.slf4j.LoggerFactory.getLogger(ZKLock.class);

    private InterProcessSemaphoreMutex lock;

    private boolean locked;

    private String lockPath;

    private final CuratorFramework zkClient;

    public ZKLock(final CuratorFramework zkClient, final String lockPath) {
        this.lockPath = lockPath;
        synchronized (zkClient) {
            if (zkClient.getState() == CuratorFrameworkState.LATENT) {
                zkClient.start();
            }
        }
        this.zkClient = zkClient;
        this.zkClient.getConnectionStateListenable().addListener((arg0, state) -> {
            if (state == ConnectionState.RECONNECTED) {
                if (lock != null) {
                    try {
                        lock.acquire();
                    } catch (final Exception e) {
                        logger.warn("Error when relocking after reconnecting", e);
                    }
                }
            }
        });
        lock = new InterProcessSemaphoreMutex(zkClient, lockPath);
    }

    public ZKLock(final CuratorFramework zkClient) {
        if (zkClient.getState() == CuratorFrameworkState.LATENT) {
            zkClient.start();
        }
        this.zkClient = zkClient;
    }

    public void lock() {
        logger.info("Trying to lock on " + lockPath);
        try {
            if (lock != null && !lock.isAcquiredInThisProcess()) {
                lock.acquire();
                logger.info("Locked " + lockPath + " " + lock.isAcquiredInThisProcess());
            }
        } catch (final Exception e) {
            throw new RuntimeException(e);
        }
    }

    public boolean tryLock(final String lockPath, final long time, final TimeUnit unit) throws Exception {
        if (lock != null && lock.isAcquiredInThisProcess()) {
            lock.release();
        }
        lock = new InterProcessSemaphoreMutex(zkClient, lockPath);
        this.lockPath = lockPath;
        return lock.acquire(time, unit);
    }

    public void lockNonBlocking() {
        logger.info("Trying nonblocking lock " + lockPath);
        final Thread t = new Thread((Runnable) () -> lock(), "ZKLock");
        t.start();
    }

    public boolean isLockObtained() {
        return lock.isAcquiredInThisProcess();
    }

    public void release() {
        try {
            if (lock != null && lock.isAcquiredInThisProcess()) {
                lock.release();
                logger.info("Released lock " + lockPath);
            }
        } catch (final Exception e) {
            logger.error("Error while releasing the lock", e);
        }
    }

    public static void main(final String[] args) throws Exception {
        final CuratorFramework client = CuratorFrameworkFactory.newClient("localhost:2181",
                new ExponentialBackoffRetry(1000, 10));
        final ZKLock lock = new ZKLock(client);

        boolean locked = false;

        while (!locked) {
            locked = lock.tryLock("/bjn/indigo/locks/queues/queue1", 5, TimeUnit.SECONDS);
        }

        for (int i = 0; i < 10000; i++) {
            System.out.println(lock.isLockObtained());
            Thread.sleep(3000);
        }

        lock.release();
        client.close();

    }

    /**
     * @return the lock
     */
    public InterProcessSemaphoreMutex getLock() {
        return lock;
    }

    /**
     * @return the lockPath
     */
    public String getLockPath() {
        return lockPath;
    }

    /**
     * @return the locked
     */
    public boolean isLocked() {
        return locked;
    }

    /**
     * @param locked
     *            the locked to set
     */
    public void setLocked(final boolean locked) {
        this.locked = locked;
    }

}