ZkUtils { public void close() { try { zkClient.close(); } catch (ZkInterruptedException e) { LOG.warn("Interrupted when closing zkClient. Clearing the interrupted status and retrying.", e); Thread.interrupted(); zkClient.close(); Thread.currentThread().interrupt(); } } ZkUtils(ZkKeyBuilder zkKeyBuilder, ZkClient zkClient, int connectionTimeoutMs, int sessionTimeOutMs, MetricsRegistry metricsRegistry); void incGeneration(); int getGeneration(); void connect(); synchronized void unregister(); ZkClient getZkClient(); ZkKeyBuilder getKeyBuilder(); synchronized String registerProcessorAndGetId(final ProcessorData data); void timeDelay(int sleepTimeInMillis); synchronized void deleteProcessorNode(String processorId); void writeTaskLocality(TaskName taskName, LocationId locationId); Map<TaskName, LocationId> readTaskLocality(); List<String> getSortedActiveProcessorsZnodes(); List<String> getSortedActiveProcessorsIDs(); List<String> getActiveProcessorsIDs(List<String> znodeIds); void unsubscribeDataChanges(String path, IZkDataListener dataListener); void subscribeDataChanges(String path, IZkDataListener dataListener); void subscribeChildChanges(String path, IZkChildListener listener); void unsubscribeChildChanges(String path, IZkChildListener childListener); void writeData(String path, Object object); boolean exists(String path); void close(); void subscribeToJobModelVersionChange(GenerationAwareZkDataListener dataListener); void publishJobModel(String jobModelVersion, JobModel jobModel); JobModel getJobModel(String jobModelVersion); String getJobModelVersion(); String getNextJobModelVersion(String currentJobModelVersion); void publishJobModelVersion(String oldVersion, String newVersion); void validateZkVersion(); void validatePaths(String[] paths); void subscribeToProcessorChange(IZkChildListener listener); void cleanupZK(int numVersionsToLeave);  }
@Test public void testCloseShouldRetryOnceOnInterruptedException() { ZkClient zkClient = Mockito.mock(ZkClient.class); ZkUtils zkUtils = new ZkUtils(KEY_BUILDER, zkClient, CONNECTION_TIMEOUT_MS, SESSION_TIMEOUT_MS, new NoOpMetricsRegistry()); Mockito.doThrow(new ZkInterruptedException(new InterruptedException())) .doAnswer(invocation -> null) .when(zkClient).close(); zkUtils.close(); Mockito.verify(zkClient, Mockito.times(2)).close(); }
@Test public void testCloseShouldTearDownZkConnectionOnInterruptedException() throws Exception { CountDownLatch latch = new CountDownLatch(1); ZkClient zkClient = new ZkClient("127.0.0.1:" + zkServer.getPort()); ZkUtils zkUtils = new ZkUtils(KEY_BUILDER, zkClient, CONNECTION_TIMEOUT_MS, SESSION_TIMEOUT_MS, new NoOpMetricsRegistry()); Thread threadToInterrupt = new Thread(() -> { try { latch.await(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } zkUtils.close(); }); threadToInterrupt.start(); Field field = ZkClient.class.getDeclaredField("_closed"); field.setAccessible(true); Assert.assertFalse(field.getBoolean(zkClient)); threadToInterrupt.interrupt(); threadToInterrupt.join(); Assert.assertTrue(field.getBoolean(zkClient)); }