@Override
    public ApplicationLock lockApplication(OrchestratorContext context, ApplicationInstanceReference reference)
            throws UncheckedTimeoutException {

        Runnable onRegistryClose;

        // A multi-application operation, aka batch suspension, will first issue a probe
        // then a non-probe. With "large locks", the lock is not released in between -
        // no lock is taken on the non-probe. Instead, the release is done on the multi-application
        // context close.
        if (context.hasLock(reference)) {
            onRegistryClose = () -> {};
        } else {
            Runnable unlock = acquireLock(context, reference);
            if (context.registerLockAcquisition(reference, unlock)) {
                onRegistryClose = () -> {};
            } else {
                onRegistryClose = unlock;
            }
        }

        try {
            return new ZkApplicationLock(
                    this,
                    curator,
                    onRegistryClose,
                    reference,
                    context.isProbe(),
                    hostInfosCache);
        } catch (Throwable t) {
            // In case the constructor throws an exception.
            onRegistryClose.run();
            throw t;
        }
    }
@Test
    public void setting_host_state_is_idempotent() {
        when(timer.currentTime()).thenReturn(
                Instant.ofEpochMilli((1)),
                Instant.ofEpochMilli((3)),
                Instant.ofEpochMilli(6));

        try (ApplicationLock lock = statusService.lockApplication(context, TestIds.APPLICATION_INSTANCE_REFERENCE)) {

            //shuffling to catch "clean database" failures for all cases.
            for (HostStatus hostStatus: shuffledList(HostStatus.NO_REMARKS, HostStatus.ALLOWED_TO_BE_DOWN)) {
                for (int i = 0; i < 2; i++) {
                    lock.setHostState(TestIds.HOST_NAME1, hostStatus);

                    assertEquals(hostStatus, lock.getHostInfos().getOrNoRemarks(TestIds.HOST_NAME1).status());
                }
            }
        }
    }