ZooKeeperStateHandleStore { public RetrievableStateHandle<T> addAndLock( String pathInZooKeeper, T state) throws Exception { checkNotNull(pathInZooKeeper, "Path in ZooKeeper"); checkNotNull(state, "State"); final String path = normalizePath(pathInZooKeeper); RetrievableStateHandle<T> storeHandle = storage.store(state); boolean success = false; try { byte[] serializedStoreHandle = InstantiationUtil.serializeObject(storeHandle); client.inTransaction().create().withMode(CreateMode.PERSISTENT).forPath(path, serializedStoreHandle) .and().create().withMode(CreateMode.EPHEMERAL).forPath(getLockPath(path)) .and().commit(); success = true; return storeHandle; } catch (KeeperException.NodeExistsException e) { throw new ConcurrentModificationException("ZooKeeper unexpectedly modified", e); } finally { if (!success) { if (storeHandle != null) { storeHandle.discardState(); } } } } ZooKeeperStateHandleStore(
		CuratorFramework client,
		RetrievableStateStorageHelper<T> storage,
		Executor executor); RetrievableStateHandle<T> addAndLock(
			String pathInZooKeeper,
			T state); void replace(String pathInZooKeeper, int expectedVersion, T state); int exists(String pathInZooKeeper); RetrievableStateHandle<T> getAndLock(String pathInZooKeeper); Collection<String> getAllPaths(); @SuppressWarnings("unchecked") List<Tuple2<RetrievableStateHandle<T>, String>> getAllAndLock(); @SuppressWarnings("unchecked") List<Tuple2<RetrievableStateHandle<T>, String>> getAllSortedByNameAndLock(); void releaseAndTryRemove(String pathInZooKeeper); void releaseAndTryRemove(
			String pathInZooKeeper,
			@Nullable final RemoveCallback<T> callback); void releaseAndTryRemoveAll(); void release(String pathInZooKeeper); void releaseAll(); static Logger LOG; }
@Test public void testAddAndLock() throws Exception { LongStateStorage longStateStorage = new LongStateStorage(); ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<Long>( ZOOKEEPER.getClient(), longStateStorage, Executors.directExecutor()); final String pathInZooKeeper = "/testAdd"; final Long state = 1239712317L; store.addAndLock(pathInZooKeeper, state); assertEquals(1, store.getAllAndLock().size()); assertEquals(state, store.getAndLock(pathInZooKeeper).retrieveState()); Stat stat = ZOOKEEPER.getClient().checkExists().forPath(pathInZooKeeper); assertNotNull(stat); assertEquals(0, stat.getEphemeralOwner()); List<String> children = ZOOKEEPER.getClient().getChildren().forPath(pathInZooKeeper); assertEquals(1, children.size()); stat = ZOOKEEPER.getClient().checkExists().forPath(pathInZooKeeper + '/' + children.get(0)); assertNotNull(stat); assertNotEquals(0, stat.getEphemeralOwner()); @SuppressWarnings("unchecked") Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject( ZOOKEEPER.getClient().getData().forPath(pathInZooKeeper), ClassLoader.getSystemClassLoader())).retrieveState(); assertEquals(state, actual); }
@Test public void testLockCleanupWhenClientTimesOut() throws Exception { LongStateStorage longStateStorage = new LongStateStorage(); Configuration configuration = new Configuration(); configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_QUORUM, ZOOKEEPER.getConnectString()); configuration.setInteger(HighAvailabilityOptions.ZOOKEEPER_SESSION_TIMEOUT, 100); configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_ROOT, "timeout"); try (CuratorFramework client = ZooKeeperUtils.startCuratorFramework(configuration); CuratorFramework client2 = ZooKeeperUtils.startCuratorFramework(configuration)) { ZooKeeperStateHandleStore<Long> zkStore = new ZooKeeperStateHandleStore<>( client, longStateStorage, Executors.directExecutor()); final String path = "/state"; zkStore.addAndLock(path, 42L); client.close(); Stat stat = client2.checkExists().forPath(path); assertNotNull(stat); Collection<String> children = client2.getChildren().forPath(path); assertEquals(0, children.size()); } }