org.apache.curator.framework.recipes.locks.InterProcessMutex Java Examples

@Test
public void shouldRunRegisteredSetupSteps() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    SetupStep setupStep2 = mock(SetupStep.class);
    steps.add(setupStep1);
    steps.add(setupStep2);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    setupServerIdSelectionMocks();
    setupSetupInProgressPathMocks(ZooDefs.Ids.OPEN_ACL_UNSAFE);

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    verify(setupStep1).run();
    verify(setupStep2).run();
}
 

@Test
public void shouldCreateSetupInProgressNode() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    when(configuration.getString(HAConfiguration.HA_ZOOKEEPER_ACL)).thenReturn("digest:user:pwd");

    List<ACL> aclList = Arrays.asList(new ACL(ZooDefs.Perms.ALL, new Id("digest", "user:pwd")));
    setupServerIdSelectionMocks();
    CreateBuilder createBuilder = setupSetupInProgressPathMocks(aclList).getLeft();

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    verify(createBuilder).withACL(aclList);
    verify(createBuilder).forPath(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT+SetupSteps.SETUP_IN_PROGRESS_NODE,
            "id2".getBytes(Charsets.UTF_8));
}
 

private static void inMutex(CuratorFramework curator, String mutexPath, Runnable work) {
    final InterProcessMutex mutex = new InterProcessMutex(curator, mutexPath);
    try {
        // try to acquire mutex for index within flush period
        if (mutex.acquire(LOCK_ACQUIRE_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)) {
            try {
                work.run();
            } finally {
                mutex.release();
            }
        } else {
            _log.warn("could not acquire index lock after {} millis!!", LOCK_ACQUIRE_TIMEOUT.toMillis());
        }
    } catch (Exception e) {
        throw Throwables.propagate(e);
    }
}
 

@Test
public void shouldDeleteSetupInProgressNodeAfterCompletion() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    when(configuration.getString(HAConfiguration.HA_ZOOKEEPER_ACL)).thenReturn("digest:user:pwd");

    List<ACL> aclList = Arrays.asList(new ACL(ZooDefs.Perms.ALL, new Id("digest", "user:pwd")));
    setupServerIdSelectionMocks();
    DeleteBuilder deleteBuilder = setupSetupInProgressPathMocks(aclList).getRight();

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    verify(deleteBuilder).forPath(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT+SetupSteps.SETUP_IN_PROGRESS_NODE);
}
 

@Override
public boolean tryLock(String key) {
    try {
        InterProcessMutex mutex = locks.computeIfAbsent(key,
                __ -> new InterProcessMutex(curatorFramework, String.join("/", basePath, key)));


        boolean owned = mutex.isAcquiredInThisProcess();
        if(owned) {
            return true;
        } else {
            mutex.acquire(timeout, TimeUnit.MILLISECONDS);
        }
        return mutex.isAcquiredInThisProcess();
    } catch (Exception e) {
        return false;
    }
}
 

public void manualBuy() {
    String lockPath = "/buy";
    log.info("try to buy sth.");
    try {
        InterProcessMutex lock = new InterProcessMutex(zkClient, lockPath);
        try {
            if (lock.acquire(1, TimeUnit.MINUTES)) {
                doBuy();
                log.info("buy successfully!");
            }
        } finally {
            lock.release();
        }
    } catch (Exception e) {
        log.error("zk error");
    }
}
 

@Override
public void truncateEverything() {
	String path = zkRoot;
	InterProcessMutex lock = new InterProcessMutex(_curator, path);
	try{
		lock.acquire();
		if(_curator.checkExists().forPath(path) != null){
			_curator.delete().deletingChildrenIfNeeded().forPath(path);
		}
	}catch(Exception ex){
		LOG.error("fail truncating verything", ex);
		throw new RuntimeException(ex);
	}
	finally {
       	try{
       		lock.release();
       	}catch(Exception e){
       		LOG.error("fail releasing lock", e);
       		throw new RuntimeException(e);
       	}
	}
}
 

@Override
public void stateChanged(CuratorFramework client, ConnectionState newState) {

    switch (newState) {
    case LOST:
    case SUSPENDED:

        Collection<InterProcessMutex> oldLocks = new ArrayList<>(locks.values());
        locks.clear();

        oldLocks.stream().parallel().forEach(lock -> {
            try {
                lock.release();
            } catch (Exception e) {
                logger.trace("Can't release lock on " + newState);
            }
        });
        break;
    default:
    }
}
 

/**
 * Call each registered {@link SetupStep} one after the other.
 * @throws SetupException Thrown with any error during running setup, including Zookeeper interactions, and
 *                          individual failures in the {@link SetupStep}.
 */
@PostConstruct
public void runSetup() throws SetupException {
    HAConfiguration.ZookeeperProperties zookeeperProperties = HAConfiguration.getZookeeperProperties(configuration);
    InterProcessMutex lock = curatorFactory.lockInstance(zookeeperProperties.getZkRoot());
    try {
        LOG.info("Trying to acquire lock for running setup.");
        lock.acquire();
        LOG.info("Acquired lock for running setup.");
        handleSetupInProgress(configuration, zookeeperProperties);
        for (SetupStep step : setupSteps) {
            LOG.info("Running setup step: {}", step);
            step.run();
        }
        clearSetupInProgress(zookeeperProperties);
    } catch (SetupException se) {
        LOG.error("Got setup exception while trying to setup", se);
        throw se;
    } catch (Throwable e) {
        LOG.error("Error running setup steps", e);
        throw new SetupException("Error running setup steps", e);
    } finally {
        releaseLock(lock);
        curatorFactory.close();
    }
}
 

@Test
public void shouldRunRegisteredSetupSteps() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    SetupStep setupStep2 = mock(SetupStep.class);
    steps.add(setupStep1);
    steps.add(setupStep2);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    setupServerIdSelectionMocks();
    setupSetupInProgressPathMocks(ZooDefs.Ids.OPEN_ACL_UNSAFE);

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    verify(setupStep1).run();
    verify(setupStep2).run();
}
 

@Inject
public ZookeeperLock(ZookeeperConfiguration config) {
    LOCK_NAMESPACE = config.getProperty("workflow.decider.locking.namespace", "");
    RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
    client = CuratorFrameworkFactory.newClient(
            config.getZkConnection(),
            config.getZkSessiontimeoutMs(),
            config.getZkConnectiontimeoutMs(),
            retryPolicy
    );
    client.start();
    zkLocks = CacheBuilder.newBuilder()
            .maximumSize(CACHE_MAXSIZE)
            .expireAfterAccess(CACHE_EXPIRY_TIME, TimeUnit.MINUTES)
            .build(new CacheLoader<String, InterProcessMutex>() {
                       @Override
                       public InterProcessMutex load(String key) throws Exception {
                           return new InterProcessMutex(client, zkPath.concat(key));
                       }
                   }
            );

    zkPath = StringUtils.isEmpty(LOCK_NAMESPACE)
            ? ("/conductor/")
            : ("/conductor/" + LOCK_NAMESPACE + "/");
}
 

public static void main(String[] args) throws Exception {
    CuratorFramework client = ZKUtils.getClient();
    client.start();

    InterProcessLock lock1 = new InterProcessMutex(client, lockPath1); // 可重入锁
    InterProcessLock lock2 = new InterProcessSemaphoreMutex(client, lockPath2); // 不可重入锁
    // 组锁，多锁
    InterProcessMultiLock lock = new InterProcessMultiLock(Arrays.asList(lock1, lock2));
    if (!lock.acquire(10, TimeUnit.SECONDS)) {
        throw new IllegalStateException("不能获取多锁");
    }
    System.out.println("已获取多锁");
    System.out.println("是否有第一个锁: " + lock1.isAcquiredInThisProcess());
    System.out.println("是否有第二个锁: " + lock2.isAcquiredInThisProcess());
    try {
        resource.use(); // 资源操作
    } finally {
        System.out.println("释放多个锁");
        lock.release(); // 释放多锁
    }
    System.out.println("是否有第一个锁: " + lock1.isAcquiredInThisProcess());
    System.out.println("是否有第二个锁: " + lock2.isAcquiredInThisProcess());
    client.close();
    System.out.println("结束!");
}
 

/**
 * 解锁
 * 注意：1、加锁和解锁必须配对地调用！
 * 2、解锁线程必须与加锁线程是同一个,即不支持那种异步回调解锁！
 *
 * @return true if unlock successfully otherwise false. We don't care how the unlock fails.
 */
public static boolean unlock(int innerId) {
    InterProcessMutex mutex = map.remove(innerId);
    if (mutex != null) {
        if (mutex.isAcquiredInThisProcess()) {
            try {
                mutex.release();
                LOG.debug("解锁成功...");
                return true;
            } catch (Exception e) {
                LOG.error("解锁失败！", e);
                return false;
            }
        } else {
            LOG.error(new LockNotOwnedByCurrentThread());
            return false;
        }
    } else {
        //这里暂时不抛出异常，只打印日志
        LOG.error(new RuntimeException("API误用,根本不存在锁，你解什么锁?"));
        return false;
    }
}
 

/**
 * @param client          the client
 * @param leaderPath      the path for this leadership group
 * @param executorService thread pool to use
 * @param listener        listener
 */
public LeaderSelector(CuratorFramework client, String leaderPath, CloseableExecutorService executorService, LeaderSelectorListener listener)
{
    Preconditions.checkNotNull(client, "client cannot be null");
    PathUtils.validatePath(leaderPath);
    Preconditions.checkNotNull(listener, "listener cannot be null");

    this.client = client;
    this.listener = new WrappedListener(this, listener);
    hasLeadership = false;

    this.executorService = executorService;
    mutex = new InterProcessMutex(client, leaderPath)
    {
        @Override
        protected byte[] getLockNodeBytes()
        {
            return (id.length() > 0) ? getIdBytes(id) : null;
        }
    };
}
 

@Test
public void shouldThrowSetupExceptionAndNotDoSetupIfSetupInProgressNodeExists() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    setupServerIdSelectionMocks();
    setupSetupInProgressPathMocks(ZooDefs.Ids.OPEN_ACL_UNSAFE, mock(Stat.class));

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);

    try {
        setupSteps.runSetup();
    } catch (Exception e) {
        assertTrue(e instanceof SetupException);
    }

    verifyZeroInteractions(setupStep1);
}
 

@Test
public void shouldCreateSetupInProgressNode() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    when(configuration.getString(HAConfiguration.HA_ZOOKEEPER_ACL)).thenReturn("digest:user:pwd");

    List<ACL> aclList = Arrays.asList(new ACL(ZooDefs.Perms.ALL, new Id("digest", "user:pwd")));
    setupServerIdSelectionMocks();
    CreateBuilder createBuilder = setupSetupInProgressPathMocks(aclList).getLeft();

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    verify(createBuilder).withACL(aclList);
    verify(createBuilder).forPath(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT+SetupSteps.SETUP_IN_PROGRESS_NODE,
            "id2".getBytes(Charsets.UTF_8));
}
 

@Test
public void shouldDeleteSetupInProgressNodeAfterCompletion() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    when(configuration.getString(HAConfiguration.HA_ZOOKEEPER_ACL)).thenReturn("digest:user:pwd");

    List<ACL> aclList = Arrays.asList(new ACL(ZooDefs.Perms.ALL, new Id("digest", "user:pwd")));
    setupServerIdSelectionMocks();
    DeleteBuilder deleteBuilder = setupSetupInProgressPathMocks(aclList).getRight();

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    verify(deleteBuilder).forPath(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT+SetupSteps.SETUP_IN_PROGRESS_NODE);
}
 

/**
 * Call each registered {@link SetupStep} one after the other.
 * @throws SetupException Thrown with any error during running setup, including Zookeeper interactions, and
 *                          individual failures in the {@link SetupStep}.
 */
@PostConstruct
public void runSetup() throws SetupException {
    HAConfiguration.ZookeeperProperties zookeeperProperties = HAConfiguration.getZookeeperProperties(configuration);
    InterProcessMutex lock = curatorFactory.lockInstance(zookeeperProperties.getZkRoot());
    try {
        LOG.info("Trying to acquire lock for running setup.");
        lock.acquire();
        LOG.info("Acquired lock for running setup.");
        handleSetupInProgress(configuration, zookeeperProperties);
        for (SetupStep step : setupSteps) {
            LOG.info("Running setup step: {}", step);
            step.run();
        }
        clearSetupInProgress(zookeeperProperties);
    } catch (SetupException se) {
        LOG.error("Got setup exception while trying to setup", se);
        throw se;
    } catch (Throwable e) {
        LOG.error("Error running setup steps", e);
        throw new SetupException("Error running setup steps", e);
    } finally {
        releaseLock(lock);
        curatorFactory.close();
    }
}
 

@Test
public void shouldThrowSetupExceptionAndNotDoSetupIfSetupInProgressNodeExists() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    setupServerIdSelectionMocks();
    setupSetupInProgressPathMocks(ZooDefs.Ids.OPEN_ACL_UNSAFE, mock(Stat.class));

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);

    try {
        setupSteps.runSetup();
    } catch (Exception e) {
        assertTrue(e instanceof SetupException);
    }

    verifyZeroInteractions(setupStep1);
}
 

private SafeLock safeLock(String executionPath) {
    final InterProcessMutex mtx = mutexMap.computeIfAbsent(
            executionPath, k -> new InterProcessMutex(curator, absolutePath(LOCK_PATH, executionPath)));

    try {
        mtx.acquire();
    } catch (Exception e) {
        logger.error("Failed to acquire a lock for {}; entering read-only mode", executionPath, e);
        stopLater();
        throw new ReplicationException("failed to acquire a lock for " + executionPath, e);
    }

    return () -> {
        try {
            mtx.release();
        } catch (Exception ignored) {
            // Ignore.
        }
    };
}
 

/**
 * 使用分布式锁执行任务
 *
 * @param path
 * @param getLockTimeout 获取锁超时时间（单位ms)
 * @param task
 * @auth anduo 2015年5月8日
 */
public void distributeLock(String path, int getLockTimeout, Runnable task) {
    InterProcessMutex lock = new InterProcessMutex(client, path);
    try {
        LOGGER.debug("尝试获取锁。。。");
        if (lock.acquire(getLockTimeout, TimeUnit.MILLISECONDS)) {
            try {
                LOGGER.debug("获得锁，开始执行任务。。。");
                task.run();
            } finally {
                lock.release();
                LOGGER.debug("释放锁,path:" + path);
            }
        } else {
            LOGGER.info("任务执行失败，在时间：" + getLockTimeout + "ms内，未获得分布式锁!");
        }
    } catch (Exception e) {
        LOGGER.error("执行分布式锁任务异常。", e);
    }

}
 

public void manualBuy() {
    String lockPath = "/buy";
    log.info("try to buy sth.");
    try {
        InterProcessMutex lock = new InterProcessMutex(zkClient, lockPath);
        try {
            if (lock.acquire(1, TimeUnit.MINUTES)) {
                doBuy();
                log.info("buy successfully!");
            }
        } finally {
            lock.release();
        }
    } catch (Exception e) {
        log.error("zk error");
    }
}
 

/**
 * 环绕操作
 *
 * @param point 切入点
 * @return 原方法返回值
 * @throws Throwable 异常信息
 */
@Around("doLock()")
public Object around(ProceedingJoinPoint point) throws Throwable {
    MethodSignature signature = (MethodSignature) point.getSignature();
    Method method = signature.getMethod();
    Object[] args = point.getArgs();
    ZooLock zooLock = method.getAnnotation(ZooLock.class);
    if (StrUtil.isBlank(zooLock.key())) {
        throw new RuntimeException("分布式锁键不能为空");
    }
    String lockKey = buildLockKey(zooLock, method, args);
    InterProcessMutex lock = new InterProcessMutex(zkClient, lockKey);
    try {
        // 假设上锁成功，以后拿到的都是 false
        if (lock.acquire(zooLock.timeout(), zooLock.timeUnit())) {
            return point.proceed();
        } else {
            throw new RuntimeException("请勿重复提交");
        }
    } finally {
        lock.release();
    }
}
 

/**
 * 环绕操作
 *
 * @param point 切入点
 * @return 原方法返回值
 * @throws Throwable 异常信息
 */
@Around("doLock()")
public Object around(ProceedingJoinPoint point) throws Throwable {
    MethodSignature signature = (MethodSignature) point.getSignature();
    Method method = signature.getMethod();
    Object[] args = point.getArgs();
    ZooLock zooLock = method.getAnnotation(ZooLock.class);
    if (StrUtil.isBlank(zooLock.key())) {
        throw new RuntimeException("分布式锁键不能为空");
    }
    String lockKey = buildLockKey(zooLock, method, args);
    InterProcessMutex lock = new InterProcessMutex(zkClient, lockKey);
    try {
        // 假设上锁成功，以后拿到的都是 false
        if (lock.acquire(zooLock.timeout(), zooLock.timeUnit())) {
            return point.proceed();
        } else {
            throw new RuntimeException("请勿重复提交");
        }
    } finally {
        lock.release();
    }
}
 

public void manualBuy() {
    String lockPath = "/buy";
    log.info("try to buy sth.");
    try {
        InterProcessMutex lock = new InterProcessMutex(zkClient, lockPath);
        try {
            if (lock.acquire(1, TimeUnit.MINUTES)) {
                doBuy();
                log.info("buy successfully!");
            }
        } finally {
            lock.release();
        }
    } catch (Exception e) {
        log.error("zk error");
    }
}
 

public static boolean switchWithZK(int id, PhysicalDataHost dh, String subHostName, ManagerConnection mc) {
    CuratorFramework zkConn = ZKUtils.getConnection();
    InterProcessMutex distributeLock = new InterProcessMutex(zkConn, KVPathUtil.getHaLockPath(dh.getHostName()));
    try {
        try {
            if (!distributeLock.acquire(100, TimeUnit.MILLISECONDS)) {
                mc.writeErrMessage(ErrorCode.ER_YES, "Other instance is change the dataHost status");
                return false;
            }
            String result = dh.switchMaster(subHostName, false);
            DataHostDisable.setStatusToZK(KVPathUtil.getHaStatusPath(dh.getHostName()), zkConn, result);
            HaConfigManager.getInstance().haFinish(id, null, result);
        } finally {
            distributeLock.release();
            LOGGER.info("reload config: release distributeLock " + KVPathUtil.getConfChangeLockPath() + " from zk");
        }
    } catch (Exception e) {
        LOGGER.info("reload config using ZK failure", e);
        mc.writeErrMessage(ErrorCode.ER_YES, e.getMessage());
        HaConfigManager.getInstance().haFinish(id, e.getMessage(), null);
        return false;
    }
    return true;
}
 

public void releaseLock(String lockId) {
    if (StringUtils.isEmpty(lockId)) {
        throw new IllegalArgumentException("lockId cannot be NULL or empty: lockId=" + lockId);
    }
    try {
        InterProcessMutex lock = zkLocks.getIfPresent(lockId);
        if (lock != null) {
            lock.release();
        }
    } catch (Exception e) {
        LOGGER.debug("Failed in releaseLock: ", e);
    }
}
 

@Test
public void shouldRunSetupStepsUnderLock() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    SetupStep setupStep2 = mock(SetupStep.class);
    steps.add(setupStep1);
    steps.add(setupStep2);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    setupServerIdSelectionMocks();
    setupSetupInProgressPathMocks(ZooDefs.Ids.OPEN_ACL_UNSAFE);

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    InOrder inOrder = inOrder(lock, setupStep1, setupStep2);

    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    setupSteps.runSetup();

    inOrder.verify(lock).acquire();
    inOrder.verify(setupStep1).run();
    inOrder.verify(setupStep2).run();
    inOrder.verify(lock).release();
}
 

private void releaseLock(InterProcessMutex lock) {
    try {
        lock.release();
        LOG.info("Released lock after running setup.");
    } catch (Exception e) {
        LOG.error("Error releasing acquired lock.", e);
    }
}
 

@Test
public void shouldReleaseLockOnException() throws Exception {
    Set<SetupStep> steps = new LinkedHashSet<>();
    SetupStep setupStep1 = mock(SetupStep.class);
    steps.add(setupStep1);

    when(configuration.
            getString(HAConfiguration.ATLAS_SERVER_HA_ZK_ROOT_KEY, HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT);
    setupServerIdSelectionMocks();
    setupSetupInProgressPathMocks(ZooDefs.Ids.OPEN_ACL_UNSAFE);

    doThrow(new RuntimeException("Simulating setup failure.")).when(setupStep1).run();

    InterProcessMutex lock = mock(InterProcessMutex.class);
    when(curatorFactory.lockInstance(HAConfiguration.ATLAS_SERVER_ZK_ROOT_DEFAULT)).
            thenReturn(lock);
    InOrder inOrder = inOrder(lock, setupStep1);

    SetupSteps setupSteps = new SetupSteps(steps, curatorFactory, configuration);
    try {
        setupSteps.runSetup();
    } catch (Exception e) {
        assertTrue(e instanceof SetupException);
    }

    inOrder.verify(lock).acquire();
    inOrder.verify(setupStep1).run();
    inOrder.verify(lock).release();
}