org.apache.flink.runtime.state.RetrievableStateHandle Java Examples

private static CompletedCheckpoint retrieveCompletedCheckpoint(Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> stateHandlePath) throws FlinkException {
	long checkpointId = pathToCheckpointId(stateHandlePath.f1);

	LOG.info("Trying to retrieve checkpoint {}.", checkpointId);

	try {
		return stateHandlePath.f0.retrieveState();
	} catch (ClassNotFoundException cnfe) {
		throw new FlinkException("Could not retrieve checkpoint " + checkpointId + " from state handle under " +
			stateHandlePath.f1 + ". This indicates that you are trying to recover from state written by an " +
			"older Flink version which is not compatible. Try cleaning the state handle store.", cnfe);
	} catch (IOException ioe) {
		throw new FlinkException("Could not retrieve checkpoint " + checkpointId + " from state handle under " +
			stateHandlePath.f1 + ". This indicates that the retrieved state handle is broken. Try cleaning the " +
			"state handle store.", ioe);
	}
}
 

/**
 * Tests get operation.
 */
@Test
public void testGetAndExists() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String pathInZooKeeper = "/testGetAndExists";
	final Long state = 311222268470898L;

	// Test
	assertEquals(-1, store.exists(pathInZooKeeper));

	store.addAndLock(pathInZooKeeper, state);
	RetrievableStateHandle<Long> actual = store.getAndLock(pathInZooKeeper);

	// Verify
	assertEquals(state, actual.retrieveState());
	assertTrue(store.exists(pathInZooKeeper) >= 0);
}
 

/**
 * Tests get operation.
 */
@Test
public void testGetAndExists() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String pathInZooKeeper = "/testGetAndExists";
	final Long state = 311222268470898L;

	// Test
	assertEquals(-1, store.exists(pathInZooKeeper));

	store.addAndLock(pathInZooKeeper, state);
	RetrievableStateHandle<Long> actual = store.getAndLock(pathInZooKeeper);

	// Verify
	assertEquals(state, actual.retrieveState());
	assertTrue(store.exists(pathInZooKeeper) >= 0);
}
 

private static CompletedCheckpoint retrieveCompletedCheckpoint(Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> stateHandlePath) throws FlinkException {
	long checkpointId = pathToCheckpointId(stateHandlePath.f1);

	LOG.info("Trying to retrieve checkpoint {}.", checkpointId);

	try {
		return stateHandlePath.f0.retrieveState();
	} catch (ClassNotFoundException cnfe) {
		throw new FlinkException("Could not retrieve checkpoint " + checkpointId + " from state handle under " +
			stateHandlePath.f1 + ". This indicates that you are trying to recover from state written by an " +
			"older Flink version which is not compatible. Try cleaning the state handle store.", cnfe);
	} catch (IOException ioe) {
		throw new FlinkException("Could not retrieve checkpoint " + checkpointId + " from state handle under " +
			stateHandlePath.f1 + ". This indicates that the retrieved state handle is broken. Try cleaning the " +
			"state handle store.", ioe);
	}
}
 

@Override
public RetrievableStateHandle<T> store(T state) throws Exception {
	Exception latestException = null;

	for (int attempt = 0; attempt < 10; attempt++) {
		Path filePath = getNewFilePath();

		try (FSDataOutputStream outStream = fs.create(filePath, FileSystem.WriteMode.NO_OVERWRITE)) {
			InstantiationUtil.serializeObject(outStream, state);
			return new RetrievableStreamStateHandle<T>(filePath, outStream.getPos());
		}
		catch (Exception e) {
			latestException = e;
		}
	}

	throw new Exception("Could not open output stream for state backend", latestException);
}
 

private static CompletedCheckpoint retrieveCompletedCheckpoint(Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String> stateHandlePath) throws FlinkException {
	long checkpointId = pathToCheckpointId(stateHandlePath.f1);

	LOG.info("Trying to retrieve checkpoint {}.", checkpointId);

	try {
		return stateHandlePath.f0.retrieveState();
	} catch (ClassNotFoundException cnfe) {
		throw new FlinkException("Could not retrieve checkpoint " + checkpointId + " from state handle under " +
			stateHandlePath.f1 + ". This indicates that you are trying to recover from state written by an " +
			"older Flink version which is not compatible. Try cleaning the state handle store.", cnfe);
	} catch (IOException ioe) {
		throw new FlinkException("Could not retrieve checkpoint " + checkpointId + " from state handle under " +
			stateHandlePath.f1 + ". This indicates that the retrieved state handle is broken. Try cleaning the " +
			"state handle store.", ioe);
	}
}
 

@Override
public RetrievableStateHandle<T> store(T state) throws Exception {
	Exception latestException = null;

	for (int attempt = 0; attempt < 10; attempt++) {
		Path filePath = getNewFilePath();

		try (FSDataOutputStream outStream = fs.create(filePath, FileSystem.WriteMode.NO_OVERWRITE)) {
			InstantiationUtil.serializeObject(outStream, state);
			return new RetrievableStreamStateHandle<T>(filePath, outStream.getPos());
		}
		catch (Exception e) {
			latestException = e;
		}
	}

	throw new Exception("Could not open output stream for state backend", latestException);
}
 

/**
 * Tests get operation.
 */
@Test
public void testGetAndExists() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String pathInZooKeeper = "/testGetAndExists";
	final Long state = 311222268470898L;

	// Test
	assertEquals(-1, store.exists(pathInZooKeeper));

	store.addAndLock(pathInZooKeeper, state);
	RetrievableStateHandle<Long> actual = store.getAndLock(pathInZooKeeper);

	// Verify
	assertEquals(state, actual.retrieveState());
	assertTrue(store.exists(pathInZooKeeper) >= 0);
}
 

@Override
public RetrievableStateHandle<T> store(T state) throws Exception {
	Exception latestException = null;

	for (int attempt = 0; attempt < 10; attempt++) {
		Path filePath = getNewFilePath();

		try (FSDataOutputStream outStream = fs.create(filePath, FileSystem.WriteMode.NO_OVERWRITE)) {
			InstantiationUtil.serializeObject(outStream, state);
			return new RetrievableStreamStateHandle<T>(filePath, outStream.getPos());
		}
		catch (Exception e) {
			latestException = e;
		}
	}

	throw new Exception("Could not open output stream for state backend", latestException);
}
 

/**
 * Gets a state handle from ZooKeeper and optionally locks it.
 *
 * @param pathInZooKeeper Path in ZooKeeper to get the state handle from
 * @param lock True if we should lock the node; otherwise false
 * @return The state handle
 * @throws IOException Thrown if the method failed to deserialize the stored state handle
 * @throws Exception Thrown if a ZooKeeper operation failed
 */
@SuppressWarnings("unchecked")
private RetrievableStateHandle<T> get(String pathInZooKeeper, boolean lock) throws Exception {
	checkNotNull(pathInZooKeeper, "Path in ZooKeeper");

	final String path = normalizePath(pathInZooKeeper);

	if (lock) {
		// try to lock the node
		try {
			client.create().withMode(CreateMode.EPHEMERAL).forPath(getLockPath(path));
		} catch (KeeperException.NodeExistsException ignored) {
			// we have already created the lock
		}
	}

	boolean success = false;

	try {
		byte[] data = client.getData().forPath(path);

		try {
			RetrievableStateHandle<T> retrievableStateHandle = InstantiationUtil.deserializeObject(
				data,
				Thread.currentThread().getContextClassLoader());

			success = true;

			return retrievableStateHandle;
		} catch (IOException | ClassNotFoundException e) {
			throw new IOException("Failed to deserialize state handle from ZooKeeper data from " +
				path + '.', e);
		}
	} finally {
		if (!success && lock) {
			// release the lock
			release(path);
		}
	}
}
 

/**
 * Releases the lock for the given state node and tries to remove the state node if it is no longer locked.
 * It returns the {@link RetrievableStateHandle} stored under the given state node if any.
 *
 * @param pathInZooKeeper Path of state handle to remove
 * @return True if the state handle could be released
 * @throws Exception If the ZooKeeper operation or discarding the state handle fails
 */
public boolean releaseAndTryRemove(String pathInZooKeeper) throws Exception {
	checkNotNull(pathInZooKeeper, "Path in ZooKeeper");

	final String path = normalizePath(pathInZooKeeper);

	RetrievableStateHandle<T> stateHandle = null;

	try {
		stateHandle = get(path, false);
	} catch (Exception e) {
		LOG.warn("Could not retrieve the state handle from node {}.", path, e);
	}

	release(pathInZooKeeper);

	try {
		client.delete().forPath(path);
	} catch (KeeperException.NotEmptyException ignored) {
		LOG.debug("Could not delete znode {} because it is still locked.", path);
		return false;
	}

	if (stateHandle != null) {
		stateHandle.discardState();
	}

	return true;
}
 

@Override
public RetrievableStateHandle<SubmittedJobGraph> store(SubmittedJobGraph state) throws IOException {
	ByteStreamStateHandle byteStreamStateHandle = new ByteStreamStateHandle(
			String.valueOf(UUID.randomUUID()),
			InstantiationUtil.serializeObject(state));
	return new RetrievableStreamStateHandle<>(byteStreamStateHandle);
}
 

/**
 * Tests that a state handle is replaced.
 */
@Test
public void testReplace() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String pathInZooKeeper = "/testReplace";
	final Long initialState = 30968470898L;
	final Long replaceState = 88383776661L;

	// Test
	store.addAndLock(pathInZooKeeper, initialState);
	store.replace(pathInZooKeeper, 0, replaceState);

	// Verify
	// State handles created
	assertEquals(2, stateHandleProvider.getStateHandles().size());
	assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).retrieveState());
	assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).retrieveState());

	// Path created and is persistent
	Stat stat = ZOOKEEPER.getClient().checkExists().forPath(pathInZooKeeper);
	assertNotNull(stat);
	assertEquals(0, stat.getEphemeralOwner());

	// Data is equal
	@SuppressWarnings("unchecked")
	Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(
			ZOOKEEPER.getClient().getData().forPath(pathInZooKeeper),
			ClassLoader.getSystemClassLoader())).retrieveState();

	assertEquals(replaceState, actual);
}
 

/**
 * FLINK-6612
 *
 * Tests that a concurrent delete operation cannot succeed if another instance holds a lock on the specified
 * node.
 */
@Test
public void testConcurrentDeleteOperation() throws Exception {
	LongStateStorage longStateStorage = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> zkStore1 = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		longStateStorage);

	ZooKeeperStateHandleStore<Long> zkStore2 = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		longStateStorage);

	final String statePath = "/state";

	zkStore1.addAndLock(statePath, 42L);
	RetrievableStateHandle<Long> stateHandle = zkStore2.getAndLock(statePath);

	// this should not remove the referenced node because we are still holding a state handle
	// reference via zkStore2
	zkStore1.releaseAndTryRemove(statePath);

	// sanity check
	assertEquals(42L, (long) stateHandle.retrieveState());

	Stat nodeStat = ZOOKEEPER.getClient().checkExists().forPath(statePath);

	assertNotNull("NodeStat should not be null, otherwise the referenced node does not exist.", nodeStat);

	zkStore2.releaseAndTryRemove(statePath);

	nodeStat = ZOOKEEPER.getClient().checkExists().forPath(statePath);

	assertNull("NodeState should be null, because the referenced node should no longer exist.", nodeStat);
}
 

@Override
public RetrievableStateHandle<Long> store(Long state) throws Exception {
	LongRetrievableStateHandle stateHandle = new LongRetrievableStateHandle(state);
	stateHandles.add(stateHandle);

	return stateHandle;
}
 

/**
 * Tests that the state is returned sorted.
 */
@Test
public void testGetAllSortedByName() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String basePath = "/testGetAllSortedByName";

	final Long[] expected = new Long[] {
			311222268470898L, 132812888L, 27255442L, 11122233124L };

	// Test
	for (long val : expected) {
		final String pathInZooKeeper = String.format("%s%016d", basePath, val);
		store.addAndLock(pathInZooKeeper, val);
	}

	List<Tuple2<RetrievableStateHandle<Long>, String>> actual = store.getAllAndLock();
	assertEquals(expected.length, actual.size());

	// bring the elements in sort order
	Arrays.sort(expected);
	Collections.sort(actual, Comparator.comparing(o -> o.f1));

	for (int i = 0; i < expected.length; i++) {
		assertEquals(expected[i], actual.get(i).f0.retrieveState());
	}
}
 

@Override
public RetrievableStateHandle<Long> store(Long state) throws Exception {
	LongRetrievableStateHandle stateHandle = new LongRetrievableStateHandle(state);
	stateHandles.add(stateHandle);

	return stateHandle;
}
 

/**
 * Tests that the ZooKeeperStateHandleStore can handle corrupted data by releasing and trying to remove the
 * respective ZooKeeper ZNodes.
 */
@Test
public void testCorruptedData() throws Exception {
	LongStateStorage stateStorage = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		stateStorage);

	final Collection<Long> input = new HashSet<>();
	input.add(1L);
	input.add(2L);
	input.add(3L);

	for (Long aLong : input) {
		store.addAndLock("/" + aLong, aLong);
	}

	// corrupt one of the entries
	ZOOKEEPER.getClient().setData().forPath("/" + 2, new byte[2]);

	List<Tuple2<RetrievableStateHandle<Long>, String>> allEntries = store.getAllAndLock();

	Collection<Long> expected = new HashSet<>(input);
	expected.remove(2L);

	Collection<Long> actual = new HashSet<>(expected.size());

	for (Tuple2<RetrievableStateHandle<Long>, String> entry : allEntries) {
		actual.add(entry.f0.retrieveState());
	}

	assertEquals(expected, actual);
}
 

/**
 * FLINK-6612
 *
 * Tests that a concurrent delete operation cannot succeed if another instance holds a lock on the specified
 * node.
 */
@Test
public void testConcurrentDeleteOperation() throws Exception {
	LongStateStorage longStateStorage = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> zkStore1 = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		longStateStorage);

	ZooKeeperStateHandleStore<Long> zkStore2 = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		longStateStorage);

	final String statePath = "/state";

	zkStore1.addAndLock(statePath, 42L);
	RetrievableStateHandle<Long> stateHandle = zkStore2.getAndLock(statePath);

	// this should not remove the referenced node because we are still holding a state handle
	// reference via zkStore2
	zkStore1.releaseAndTryRemove(statePath);

	// sanity check
	assertEquals(42L, (long) stateHandle.retrieveState());

	Stat nodeStat = ZOOKEEPER.getClient().checkExists().forPath(statePath);

	assertNotNull("NodeStat should not be null, otherwise the referenced node does not exist.", nodeStat);

	zkStore2.releaseAndTryRemove(statePath);

	nodeStat = ZOOKEEPER.getClient().checkExists().forPath(statePath);

	assertNull("NodeState should be null, because the referenced node should no longer exist.", nodeStat);
}
 

@Override
public RetrievableStateHandle<Long> store(Long state) throws Exception {
	LongRetrievableStateHandle stateHandle = new LongRetrievableStateHandle(state);
	stateHandles.add(stateHandle);

	return stateHandle;
}
 

@Override
public RetrievableStateHandle<SubmittedJobGraph> store(SubmittedJobGraph state) throws IOException {
	ByteStreamStateHandle byteStreamStateHandle = new ByteStreamStateHandle(
			String.valueOf(UUID.randomUUID()),
			InstantiationUtil.serializeObject(state));
	return new RetrievableStreamStateHandle<>(byteStreamStateHandle);
}
 

@Override
public List<MesosWorkerStore.Worker> recoverWorkers() throws Exception {
	synchronized (startStopLock) {
		verifyIsRunning();

		List<Tuple2<RetrievableStateHandle<Worker>, String>> handles = workersInZooKeeper.getAllAndLock();

		if (handles.isEmpty()) {
			return Collections.emptyList();
		}
		else {
			List<MesosWorkerStore.Worker> workers = new ArrayList<>(handles.size());

			for (Tuple2<RetrievableStateHandle<Worker>, String> handle : handles) {
				final Worker worker;

				try {
					worker = handle.f0.retrieveState();
				} catch (ClassNotFoundException cnfe) {
					throw new FlinkException("Could not retrieve Mesos worker from state handle under " +
						handle.f1 + ". This indicates that you are trying to recover from state written by an " +
						"older Flink version which is not compatible. Try cleaning the state handle store.", cnfe);
				} catch (IOException ioe) {
					throw new FlinkException("Could not retrieve Mesos worker from state handle under " +
						handle.f1 + ". This indicates that the retrieved state handle is broken. Try cleaning " +
						"the state handle store.", ioe);
				}

				workers.add(worker);
			}

			return workers;
		}
	}
}
 

/**
 * Replaces a state handle in ZooKeeper and discards the old state handle.
 *
 * @param pathInZooKeeper Destination path in ZooKeeper (expected to exist and start with a '/')
 * @param expectedVersion Expected version of the node to replace
 * @param state           The new state to replace the old one
 * @throws Exception If a ZooKeeper or state handle operation fails
 */
public void replace(String pathInZooKeeper, int expectedVersion, T state) throws Exception {
	checkNotNull(pathInZooKeeper, "Path in ZooKeeper");
	checkNotNull(state, "State");

	final String path = normalizePath(pathInZooKeeper);

	RetrievableStateHandle<T> oldStateHandle = get(path, false);

	RetrievableStateHandle<T> newStateHandle = storage.store(state);

	boolean success = false;

	try {
		// Serialize the new state handle. This writes the state to the backend.
		byte[] serializedStateHandle = InstantiationUtil.serializeObject(newStateHandle);

		// Replace state handle in ZooKeeper.
		client.setData()
				.withVersion(expectedVersion)
				.forPath(path, serializedStateHandle);
		success = true;
	} catch (KeeperException.NoNodeException e) {
		throw new ConcurrentModificationException("ZooKeeper unexpectedly modified", e);
	} finally {
		if (success) {
			oldStateHandle.discardState();
		} else {
			newStateHandle.discardState();
		}
	}

}
 

/**
 * Releases the lock for the given state node and tries to remove the state node if it is no longer locked.
 * It returns the {@link RetrievableStateHandle} stored under the given state node if any.
 *
 * @param pathInZooKeeper Path of state handle to remove
 * @return True if the state handle could be released
 * @throws Exception If the ZooKeeper operation or discarding the state handle fails
 */
public boolean releaseAndTryRemove(String pathInZooKeeper) throws Exception {
	checkNotNull(pathInZooKeeper, "Path in ZooKeeper");

	final String path = normalizePath(pathInZooKeeper);

	RetrievableStateHandle<T> stateHandle = null;

	try {
		stateHandle = get(path, false);
	} catch (Exception e) {
		LOG.warn("Could not retrieve the state handle from node {}.", path, e);
	}

	release(pathInZooKeeper);

	try {
		client.delete().forPath(path);
	} catch (KeeperException.NotEmptyException ignored) {
		LOG.debug("Could not delete znode {} because it is still locked.", path);
		return false;
	}

	if (stateHandle != null) {
		stateHandle.discardState();
	}

	return true;
}
 

/**
 * Gets a state handle from ZooKeeper and optionally locks it.
 *
 * @param pathInZooKeeper Path in ZooKeeper to get the state handle from
 * @param lock True if we should lock the node; otherwise false
 * @return The state handle
 * @throws IOException Thrown if the method failed to deserialize the stored state handle
 * @throws Exception Thrown if a ZooKeeper operation failed
 */
@SuppressWarnings("unchecked")
private RetrievableStateHandle<T> get(String pathInZooKeeper, boolean lock) throws Exception {
	checkNotNull(pathInZooKeeper, "Path in ZooKeeper");

	final String path = normalizePath(pathInZooKeeper);

	if (lock) {
		// try to lock the node
		try {
			client.create().withMode(CreateMode.EPHEMERAL).forPath(getLockPath(path));
		} catch (KeeperException.NodeExistsException ignored) {
			// we have already created the lock
		}
	}

	boolean success = false;

	try {
		byte[] data = client.getData().forPath(path);

		try {
			RetrievableStateHandle<T> retrievableStateHandle = InstantiationUtil.deserializeObject(
				data,
				Thread.currentThread().getContextClassLoader());

			success = true;

			return retrievableStateHandle;
		} catch (IOException | ClassNotFoundException e) {
			throw new IOException("Failed to deserialize state handle from ZooKeeper data from " +
				path + '.', e);
		}
	} finally {
		if (!success && lock) {
			// release the lock
			release(path);
		}
	}
}
 

/**
 * Tests that a state handle is replaced.
 */
@Test
public void testReplace() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String pathInZooKeeper = "/testReplace";
	final Long initialState = 30968470898L;
	final Long replaceState = 88383776661L;

	// Test
	store.addAndLock(pathInZooKeeper, initialState);
	store.replace(pathInZooKeeper, 0, replaceState);

	// Verify
	// State handles created
	assertEquals(2, stateHandleProvider.getStateHandles().size());
	assertEquals(initialState, stateHandleProvider.getStateHandles().get(0).retrieveState());
	assertEquals(replaceState, stateHandleProvider.getStateHandles().get(1).retrieveState());

	// Path created and is persistent
	Stat stat = ZOOKEEPER.getClient().checkExists().forPath(pathInZooKeeper);
	assertNotNull(stat);
	assertEquals(0, stat.getEphemeralOwner());

	// Data is equal
	@SuppressWarnings("unchecked")
	Long actual = ((RetrievableStateHandle<Long>) InstantiationUtil.deserializeObject(
			ZOOKEEPER.getClient().getData().forPath(pathInZooKeeper),
			ClassLoader.getSystemClassLoader())).retrieveState();

	assertEquals(replaceState, actual);
}
 

/**
 * Tests that the state is returned sorted.
 */
@Test
public void testGetAllSortedByName() throws Exception {
	// Setup
	LongStateStorage stateHandleProvider = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
			ZOOKEEPER.getClient(), stateHandleProvider);

	// Config
	final String basePath = "/testGetAllSortedByName";

	final Long[] expected = new Long[] {
			311222268470898L, 132812888L, 27255442L, 11122233124L };

	// Test
	for (long val : expected) {
		final String pathInZooKeeper = String.format("%s%016d", basePath, val);
		store.addAndLock(pathInZooKeeper, val);
	}

	List<Tuple2<RetrievableStateHandle<Long>, String>> actual = store.getAllAndLock();
	assertEquals(expected.length, actual.size());

	// bring the elements in sort order
	Arrays.sort(expected);
	Collections.sort(actual, Comparator.comparing(o -> o.f1));

	for (int i = 0; i < expected.length; i++) {
		assertEquals(expected[i], actual.get(i).f0.retrieveState());
	}
}
 

/**
 * Tests that the ZooKeeperStateHandleStore can handle corrupted data by releasing and trying to remove the
 * respective ZooKeeper ZNodes.
 */
@Test
public void testCorruptedData() throws Exception {
	LongStateStorage stateStorage = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> store = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		stateStorage);

	final Collection<Long> input = new HashSet<>();
	input.add(1L);
	input.add(2L);
	input.add(3L);

	for (Long aLong : input) {
		store.addAndLock("/" + aLong, aLong);
	}

	// corrupt one of the entries
	ZOOKEEPER.getClient().setData().forPath("/" + 2, new byte[2]);

	List<Tuple2<RetrievableStateHandle<Long>, String>> allEntries = store.getAllAndLock();

	Collection<Long> expected = new HashSet<>(input);
	expected.remove(2L);

	Collection<Long> actual = new HashSet<>(expected.size());

	for (Tuple2<RetrievableStateHandle<Long>, String> entry : allEntries) {
		actual.add(entry.f0.retrieveState());
	}

	assertEquals(expected, actual);
}
 

/**
 * FLINK-6612
 *
 * Tests that a concurrent delete operation cannot succeed if another instance holds a lock on the specified
 * node.
 */
@Test
public void testConcurrentDeleteOperation() throws Exception {
	LongStateStorage longStateStorage = new LongStateStorage();

	ZooKeeperStateHandleStore<Long> zkStore1 = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		longStateStorage);

	ZooKeeperStateHandleStore<Long> zkStore2 = new ZooKeeperStateHandleStore<>(
		ZOOKEEPER.getClient(),
		longStateStorage);

	final String statePath = "/state";

	zkStore1.addAndLock(statePath, 42L);
	RetrievableStateHandle<Long> stateHandle = zkStore2.getAndLock(statePath);

	// this should not remove the referenced node because we are still holding a state handle
	// reference via zkStore2
	zkStore1.releaseAndTryRemove(statePath);

	// sanity check
	assertEquals(42L, (long) stateHandle.retrieveState());

	Stat nodeStat = ZOOKEEPER.getClient().checkExists().forPath(statePath);

	assertNotNull("NodeStat should not be null, otherwise the referenced node does not exist.", nodeStat);

	zkStore2.releaseAndTryRemove(statePath);

	nodeStat = ZOOKEEPER.getClient().checkExists().forPath(statePath);

	assertNull("NodeState should be null, because the referenced node should no longer exist.", nodeStat);
}
 

@Override
public List<MesosWorkerStore.Worker> recoverWorkers() throws Exception {
	synchronized (startStopLock) {
		verifyIsRunning();

		List<Tuple2<RetrievableStateHandle<Worker>, String>> handles = workersInZooKeeper.getAllAndLock();

		if (handles.isEmpty()) {
			return Collections.emptyList();
		}
		else {
			List<MesosWorkerStore.Worker> workers = new ArrayList<>(handles.size());

			for (Tuple2<RetrievableStateHandle<Worker>, String> handle : handles) {
				final Worker worker;

				try {
					worker = handle.f0.retrieveState();
				} catch (ClassNotFoundException cnfe) {
					throw new FlinkException("Could not retrieve Mesos worker from state handle under " +
						handle.f1 + ". This indicates that you are trying to recover from state written by an " +
						"older Flink version which is not compatible. Try cleaning the state handle store.", cnfe);
				} catch (IOException ioe) {
					throw new FlinkException("Could not retrieve Mesos worker from state handle under " +
						handle.f1 + ". This indicates that the retrieved state handle is broken. Try cleaning " +
						"the state handle store.", ioe);
				}

				workers.add(worker);
			}

			return workers;
		}
	}
}