ZkStatusService implements StatusService { @Override public ApplicationLock lockApplication(OrchestratorContext context, ApplicationInstanceReference reference) throws UncheckedTimeoutException { Runnable onRegistryClose; 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) { onRegistryClose.run(); throw t; } } @Inject  ZkStatusService(
            @Component Curator curator,
            @Component Metric metric,
            @Component Timer timer,
            @Component AntiServiceMonitor antiServiceMonitor);   ZkStatusService(Curator curator, Metric metric, Timer timer, HostInfosCache hostInfosCache,
                    AntiServiceMonitor antiServiceMonitor); @Override Set<ApplicationInstanceReference> getAllSuspendedApplications(); @Override Function<ApplicationInstanceReference, HostInfos> getHostInfosByApplicationResolver(); @Override ApplicationLock lockApplication(OrchestratorContext context, ApplicationInstanceReference reference); @Override HostInfo getHostInfo(ApplicationInstanceReference reference, HostName hostName); @Override ApplicationInstanceStatus getApplicationInstanceStatus(ApplicationInstanceReference reference); @Override void onApplicationActivate(ApplicationInstanceReference reference, Set<HostName> hostnames); @Override void onApplicationRemove(ApplicationInstanceReference reference);  }
@Test public void setting_host_state_is_idempotent() throws Exception { when(timer.currentTime()).thenReturn( Instant.ofEpochMilli((1)), Instant.ofEpochMilli((3)), Instant.ofEpochMilli(6)); try (ApplicationLock lock = statusService.lockApplication(context, TestIds.APPLICATION_INSTANCE_REFERENCE)) { 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); assertThat(lock.getHostInfos().getOrNoRemarks(TestIds.HOST_NAME1).status(), is(hostStatus)); } } } }
@Test public void locks_are_exclusive() throws Exception { ZkStatusService zkStatusService2 = new ZkStatusService(curator, mock(Metric.class), new TestTimer(), antiServiceMonitor); final CompletableFuture<Void> lockedSuccessfullyFuture; try (ApplicationLock lock = statusService .lockApplication(context, TestIds.APPLICATION_INSTANCE_REFERENCE)) { lockedSuccessfullyFuture = CompletableFuture.runAsync(() -> { try (ApplicationLock lock2 = zkStatusService2 .lockApplication(context, TestIds.APPLICATION_INSTANCE_REFERENCE)) { } }); try { lockedSuccessfullyFuture.get(3, TimeUnit.SECONDS); fail("Both zookeeper host status services locked simultaneously for the same application instance"); } catch (TimeoutException ignored) { } } lockedSuccessfullyFuture.get(1, TimeUnit.MINUTES); }
@Test public void failing_to_get_lock_closes_SessionFailRetryLoop() throws Exception { ZkStatusService zkStatusService2 = new ZkStatusService(curator, mock(Metric.class), new TestTimer(), antiServiceMonitor); try (ApplicationLock lock = statusService .lockApplication(context, TestIds.APPLICATION_INSTANCE_REFERENCE)) { CompletableFuture<Void> resultOfZkOperationAfterLockFailure = CompletableFuture.runAsync(() -> { try { zkStatusService2.lockApplication(context, TestIds.APPLICATION_INSTANCE_REFERENCE); fail("Both zookeeper host status services locked simultaneously for the same application instance"); } catch (RuntimeException e) { } killSession(curator.framework(), testingServer); lock.getHostInfos().getOrNoRemarks(TestIds.HOST_NAME1); }); assertThat(resultOfZkOperationAfterLockFailure, notHoldsException()); } }