Java Code Examples for org.apache.flink.runtime.state.StreamStateHandle#openInputStream()

/**
 * Reads Flink's state meta data file from the state handle.
 */
private KeyedBackendSerializationProxy<K> readMetaData(StreamStateHandle metaStateHandle) throws Exception {

	FSDataInputStream inputStream = null;

	try {
		inputStream = metaStateHandle.openInputStream();
		cancelStreamRegistry.registerCloseable(inputStream);
		DataInputView in = new DataInputViewStreamWrapper(inputStream);
		return readMetaData(in);
	} finally {
		if (cancelStreamRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}
	}
}
 

/**
 * Reads Flink's state meta data file from the state handle.
 */
private KeyedBackendSerializationProxy<K> readMetaData(StreamStateHandle metaStateHandle) throws Exception {

	InputStream inputStream = null;

	try {
		inputStream = metaStateHandle.openInputStream();
		cancelStreamRegistry.registerCloseable(inputStream);
		DataInputView in = new DataInputViewStreamWrapper(inputStream);
		return readMetaData(in);
	} finally {
		if (cancelStreamRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}
	}
}
 

@Test
public void testTaskOwnedStateStream() throws Exception {
	final List<String> state = Arrays.asList("Flopsy", "Mopsy", "Cotton Tail", "Peter");

	final MemoryBackendCheckpointStorage storage = new MemoryBackendCheckpointStorage(
			new JobID(), null, null, DEFAULT_MAX_STATE_SIZE);

	StreamStateHandle stateHandle;

	try (CheckpointStateOutputStream stream = storage.createTaskOwnedStateStream()) {
		assertTrue(stream instanceof MemoryCheckpointOutputStream);

		new ObjectOutputStream(stream).writeObject(state);
		stateHandle = stream.closeAndGetHandle();
	}

	try (ObjectInputStream in = new ObjectInputStream(stateHandle.openInputStream())) {
		assertEquals(state, in.readObject());
	}
}
 

/**
 * Reads Flink's state meta data file from the state handle.
 */
private KeyedBackendSerializationProxy<K> readMetaData(StreamStateHandle metaStateHandle) throws Exception {

	FSDataInputStream inputStream = null;

	try {
		inputStream = metaStateHandle.openInputStream();
		cancelStreamRegistry.registerCloseable(inputStream);
		DataInputView in = new DataInputViewStreamWrapper(inputStream);
		return readMetaData(in);
	} finally {
		if (cancelStreamRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}
	}
}
 

@Test
public void testStateStream() throws Exception {
	HashMap<String, Integer> state = new HashMap<>();
	state.put("hey there", 2);
	state.put("the crazy brown fox stumbles over a sentence that does not contain every letter", 77);

	CheckpointStateOutputStream outStream = new MemoryCheckpointOutputStream(MemoryStateBackend.DEFAULT_MAX_STATE_SIZE);
	ObjectOutputStream oos = new ObjectOutputStream(outStream);
	oos.writeObject(state);
	oos.flush();

	StreamStateHandle handle = outStream.closeAndGetHandle();
	assertNotNull(handle);

	try (ObjectInputStream ois = new ObjectInputStream(handle.openInputStream())) {
		assertEquals(state, ois.readObject());
		assertTrue(ois.available() <= 0);
	}
}
 

@Test
public void testTaskOwnedStateStream() throws Exception {
	final List<String> state = Arrays.asList("Flopsy", "Mopsy", "Cotton Tail", "Peter");

	// we chose a small size threshold here to force the state to disk
	final FsCheckpointStorage storage = new FsCheckpointStorage(
			Path.fromLocalFile(tmp.newFolder()),  null, new JobID(), 10, WRITE_BUFFER_SIZE);

	final StreamStateHandle stateHandle;

	try (CheckpointStateOutputStream stream = storage.createTaskOwnedStateStream()) {
		assertTrue(stream instanceof FsCheckpointStateOutputStream);

		new ObjectOutputStream(stream).writeObject(state);
		stateHandle = stream.closeAndGetHandle();
	}

	// the state must have gone to disk
	FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;

	// check that the state is in the correct directory
	String parentDirName = fileStateHandle.getFilePath().getParent().getName();
	assertEquals(FsCheckpointStorage.CHECKPOINT_TASK_OWNED_STATE_DIR, parentDirName);

	// validate the contents
	try (ObjectInputStream in = new ObjectInputStream(stateHandle.openInputStream())) {
		assertEquals(state, in.readObject());
	}
}
 

/**
 * Copies the file from a single state handle to the given path.
 */
private void downloadDataForStateHandle(
	Path restoreFilePath,
	StreamStateHandle remoteFileHandle,
	CloseableRegistry closeableRegistry) throws IOException {

	FSDataInputStream inputStream = null;
	FSDataOutputStream outputStream = null;

	try {
		FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
		inputStream = remoteFileHandle.openInputStream();
		closeableRegistry.registerCloseable(inputStream);

		outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
		closeableRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (closeableRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

@Test
public void testTaskOwnedStateStream() throws Exception {
	final List<String> state = Arrays.asList("Flopsy", "Mopsy", "Cotton Tail", "Peter");

	// we chose a small size threshold here to force the state to disk
	final FsCheckpointStorage storage = new FsCheckpointStorage(
			Path.fromLocalFile(tmp.newFolder()),  null, new JobID(), 10);

	final StreamStateHandle stateHandle;

	try (CheckpointStateOutputStream stream = storage.createTaskOwnedStateStream()) {
		assertTrue(stream instanceof FsCheckpointStateOutputStream);

		new ObjectOutputStream(stream).writeObject(state);
		stateHandle = stream.closeAndGetHandle();
	}

	// the state must have gone to disk
	FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;

	// check that the state is in the correct directory
	String parentDirName = fileStateHandle.getFilePath().getParent().getName();
	assertEquals(FsCheckpointStorage.CHECKPOINT_TASK_OWNED_STATE_DIR, parentDirName);

	// validate the contents
	try (ObjectInputStream in = new ObjectInputStream(stateHandle.openInputStream())) {
		assertEquals(state, in.readObject());
	}
}
 

@Test
public void testTaskOwnedStateStream() throws Exception {
	final List<String> state = Arrays.asList("Flopsy", "Mopsy", "Cotton Tail", "Peter");

	// we chose a small size threshold here to force the state to disk
	final FsCheckpointStorage storage = new FsCheckpointStorage(
			Path.fromLocalFile(tmp.newFolder()),  null, new JobID(), 10, WRITE_BUFFER_SIZE);

	final StreamStateHandle stateHandle;

	try (CheckpointStateOutputStream stream = storage.createTaskOwnedStateStream()) {
		assertTrue(stream instanceof FsCheckpointStateOutputStream);

		new ObjectOutputStream(stream).writeObject(state);
		stateHandle = stream.closeAndGetHandle();
	}

	// the state must have gone to disk
	FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;

	// check that the state is in the correct directory
	String parentDirName = fileStateHandle.getFilePath().getParent().getName();
	assertEquals(FsCheckpointStorage.CHECKPOINT_TASK_OWNED_STATE_DIR, parentDirName);

	// validate the contents
	try (ObjectInputStream in = new ObjectInputStream(stateHandle.openInputStream())) {
		assertEquals(state, in.readObject());
	}
}
 

public static KeyedBackendSerializationProxy<?> getKeyedBackendSerializationProxy(
		StreamStateHandle streamStateHandle) {
	KeyedBackendSerializationProxy<Integer> serializationProxy = new KeyedBackendSerializationProxy<>(
			StateMetadataUtils.class.getClassLoader());
	try (FSDataInputStream is = streamStateHandle.openInputStream()) {
		DataInputViewStreamWrapper iw = new DataInputViewStreamWrapper(is);
		serializationProxy.read(iw);
		return serializationProxy;
	} catch (IOException e) {
		throw new RuntimeException(e);
	}
}
 

private void copyStateDataHandleData(
		Path restoreFilePath,
		StreamStateHandle remoteFileHandle) throws IOException {

	FileSystem restoreFileSystem = restoreFilePath.getFileSystem();

	FSDataInputStream inputStream = null;
	FSDataOutputStream outputStream = null;

	try {
		inputStream = remoteFileHandle.openInputStream();
		cancelStreamRegistry.registerCloseable(inputStream);

		outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
		cancelStreamRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (cancelStreamRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (cancelStreamRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

/**
 * Copies the file from a single state handle to the given path.
 */
private void downloadDataForStateHandle(
	Path restoreFilePath,
	StreamStateHandle remoteFileHandle,
	CloseableRegistry closeableRegistry) throws IOException {

	FSDataInputStream inputStream = null;
	OutputStream outputStream = null;

	try {
		inputStream = remoteFileHandle.openInputStream();
		closeableRegistry.registerCloseable(inputStream);

		Files.createDirectories(restoreFilePath.getParent());
		outputStream = Files.newOutputStream(restoreFilePath);
		closeableRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (closeableRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

/**
 * Copies the file from a single state handle to the given path.
 */
private void downloadDataForStateHandle(
	Path restoreFilePath,
	StreamStateHandle remoteFileHandle,
	CloseableRegistry closeableRegistry) throws IOException {

	FSDataInputStream inputStream = null;
	FSDataOutputStream outputStream = null;

	try {
		FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
		inputStream = remoteFileHandle.openInputStream();
		closeableRegistry.registerCloseable(inputStream);

		outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
		closeableRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (closeableRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

private void runTest(int numBytes, int bufferSize, int threshold, boolean expectFile) throws Exception {
	FsCheckpointStreamFactory.CheckpointStateOutputStream stream =
		new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
				Path.fromLocalFile(tempDir.newFolder()), FileSystem.getLocalFileSystem(), bufferSize, threshold);

	Random rnd = new Random();
	byte[] original = new byte[numBytes];
	byte[] bytes = new byte[original.length];

	rnd.nextBytes(original);
	System.arraycopy(original, 0, bytes, 0, original.length);

	// the test writes a mixture of writing individual bytes and byte arrays
	int pos = 0;
	while (pos < bytes.length) {
		boolean single = rnd.nextBoolean();
		if (single) {
			stream.write(bytes[pos++]);
		}
		else {
			int num = rnd.nextInt(Math.min(10, bytes.length - pos));
			stream.write(bytes, pos, num);
			pos += num;
		}
	}

	StreamStateHandle handle = stream.closeAndGetHandle();
	if (expectFile) {
		assertTrue(handle instanceof FileStateHandle);
	} else {
		assertTrue(handle instanceof ByteStreamStateHandle);
	}

	// make sure the writing process did not alter the original byte array
	assertArrayEquals(original, bytes);

	try (InputStream inStream = handle.openInputStream()) {
		byte[] validation = new byte[bytes.length];

		DataInputStream dataInputStream = new DataInputStream(inStream);
		dataInputStream.readFully(validation);

		assertArrayEquals(bytes, validation);
	}

	handle.discardState();
}
 

@Override
public void notifyCheckpointComplete(long checkpointId) throws Exception {
	super.notifyCheckpointComplete(checkpointId);

	synchronized (pendingCheckpoints) {

		Iterator<PendingCheckpoint> pendingCheckpointIt = pendingCheckpoints.iterator();
		while (pendingCheckpointIt.hasNext()) {

			PendingCheckpoint pendingCheckpoint = pendingCheckpointIt.next();

			long pastCheckpointId = pendingCheckpoint.checkpointId;
			int subtaskId = pendingCheckpoint.subtaskId;
			long timestamp = pendingCheckpoint.timestamp;
			StreamStateHandle streamHandle = pendingCheckpoint.stateHandle;

			if (pastCheckpointId <= checkpointId) {
				try {
					if (!committer.isCheckpointCommitted(subtaskId, pastCheckpointId)) {
						try (FSDataInputStream in = streamHandle.openInputStream()) {
							boolean success = sendValues(
									new ReusingMutableToRegularIteratorWrapper<>(
											new InputViewIterator<>(
													new DataInputViewStreamWrapper(
															in),
													serializer),
											serializer),
									pastCheckpointId,
									timestamp);
							if (success) {
								// in case the checkpoint was successfully committed,
								// discard its state from the backend and mark it for removal
								// in case it failed, we retry on the next checkpoint
								committer.commitCheckpoint(subtaskId, pastCheckpointId);
								streamHandle.discardState();
								pendingCheckpointIt.remove();
							}
						}
					} else {
						streamHandle.discardState();
						pendingCheckpointIt.remove();
					}
				} catch (Exception e) {
					// we have to break here to prevent a new (later) checkpoint
					// from being committed before this one
					LOG.error("Could not commit checkpoint.", e);
					break;
				}
			}
		}
	}
}
 

private static void validateContents(StreamStateHandle handle, byte[] expected) throws IOException {
	try (FSDataInputStream in = handle.openInputStream()) {
		validateContents(in, expected);
	}
}
 

private void runTest(int numBytes, int bufferSize, int threshold, boolean expectFile) throws Exception {
	FsCheckpointStreamFactory.CheckpointStateOutputStream stream =
		new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
				Path.fromLocalFile(tempDir.newFolder()), FileSystem.getLocalFileSystem(), bufferSize, threshold);

	Random rnd = new Random();
	byte[] original = new byte[numBytes];
	byte[] bytes = new byte[original.length];

	rnd.nextBytes(original);
	System.arraycopy(original, 0, bytes, 0, original.length);

	// the test writes a mixture of writing individual bytes and byte arrays
	int pos = 0;
	while (pos < bytes.length) {
		boolean single = rnd.nextBoolean();
		if (single) {
			stream.write(bytes[pos++]);
		}
		else {
			int num = rnd.nextBoolean() ?
				(bytes.length - pos) : rnd.nextInt(bytes.length - pos);
			stream.write(bytes, pos, num);
			pos += num;
		}
	}

	StreamStateHandle handle = stream.closeAndGetHandle();
	if (expectFile) {
		assertTrue(handle instanceof FileStateHandle);
	} else {
		assertTrue(handle instanceof ByteStreamStateHandle);
	}

	// make sure the writing process did not alter the original byte array
	assertArrayEquals(original, bytes);

	try (InputStream inStream = handle.openInputStream()) {
		byte[] validation = new byte[bytes.length];

		DataInputStream dataInputStream = new DataInputStream(inStream);
		dataInputStream.readFully(validation);

		assertArrayEquals(bytes, validation);
	}

	handle.discardState();
}
 

private static void validateContents(StreamStateHandle handle, byte[] expected) throws IOException {
	try (FSDataInputStream in = handle.openInputStream()) {
		validateContents(in, expected);
	}
}
 

public static CompletedCheckpoint loadAndValidateCheckpoint(
		JobID jobId,
		Map<JobVertexID, ExecutionJobVertex> tasks,
		CompletedCheckpointStorageLocation location,
		ClassLoader classLoader,
		boolean allowNonRestoredState) throws IOException {

	checkNotNull(jobId, "jobId");
	checkNotNull(tasks, "tasks");
	checkNotNull(location, "location");
	checkNotNull(classLoader, "classLoader");

	final StreamStateHandle metadataHandle = location.getMetadataHandle();
	final String checkpointPointer = location.getExternalPointer();

	// (1) load the savepoint
	final CheckpointMetadata checkpointMetadata;
	try (InputStream in = metadataHandle.openInputStream()) {
		DataInputStream dis = new DataInputStream(in);
		checkpointMetadata = loadCheckpointMetadata(dis, classLoader, checkpointPointer);
	}

	// generate mapping from operator to task
	Map<OperatorID, ExecutionJobVertex> operatorToJobVertexMapping = new HashMap<>();
	for (ExecutionJobVertex task : tasks.values()) {
		for (OperatorIDPair operatorIDPair : task.getOperatorIDs()) {
			operatorToJobVertexMapping.put(operatorIDPair.getGeneratedOperatorID(), task);
			operatorIDPair.getUserDefinedOperatorID().ifPresent(id -> operatorToJobVertexMapping.put(id, task));
		}
	}

	// (2) validate it (parallelism, etc)
	HashMap<OperatorID, OperatorState> operatorStates = new HashMap<>(checkpointMetadata.getOperatorStates().size());
	for (OperatorState operatorState : checkpointMetadata.getOperatorStates()) {

		ExecutionJobVertex executionJobVertex = operatorToJobVertexMapping.get(operatorState.getOperatorID());

		if (executionJobVertex != null) {

			if (executionJobVertex.getMaxParallelism() == operatorState.getMaxParallelism()
					|| !executionJobVertex.isMaxParallelismConfigured()) {
				operatorStates.put(operatorState.getOperatorID(), operatorState);
			} else {
				String msg = String.format("Failed to rollback to checkpoint/savepoint %s. " +
								"Max parallelism mismatch between checkpoint/savepoint state and new program. " +
								"Cannot map operator %s with max parallelism %d to new program with " +
								"max parallelism %d. This indicates that the program has been changed " +
								"in a non-compatible way after the checkpoint/savepoint.",
						checkpointMetadata,
						operatorState.getOperatorID(),
						operatorState.getMaxParallelism(),
						executionJobVertex.getMaxParallelism());

				throw new IllegalStateException(msg);
			}
		} else if (allowNonRestoredState) {
			LOG.info("Skipping savepoint state for operator {}.", operatorState.getOperatorID());
		} else {
			if (operatorState.getCoordinatorState() != null) {
				throwNonRestoredStateException(checkpointPointer, operatorState.getOperatorID());
			}

			for (OperatorSubtaskState operatorSubtaskState : operatorState.getStates()) {
				if (operatorSubtaskState.hasState()) {
					throwNonRestoredStateException(checkpointPointer, operatorState.getOperatorID());
				}
			}

			LOG.info("Skipping empty savepoint state for operator {}.", operatorState.getOperatorID());
		}
	}

	// (3) convert to checkpoint so the system can fall back to it
	CheckpointProperties props = CheckpointProperties.forSavepoint(false);

	return new CompletedCheckpoint(
			jobId,
			checkpointMetadata.getCheckpointId(),
			0L,
			0L,
			operatorStates,
			checkpointMetadata.getMasterStates(),
			props,
			location);
}
 

@Override
public void notifyCheckpointComplete(long checkpointId) throws Exception {
	super.notifyCheckpointComplete(checkpointId);

	synchronized (pendingCheckpoints) {

		Iterator<PendingCheckpoint> pendingCheckpointIt = pendingCheckpoints.iterator();
		while (pendingCheckpointIt.hasNext()) {

			PendingCheckpoint pendingCheckpoint = pendingCheckpointIt.next();

			long pastCheckpointId = pendingCheckpoint.checkpointId;
			int subtaskId = pendingCheckpoint.subtaskId;
			long timestamp = pendingCheckpoint.timestamp;
			StreamStateHandle streamHandle = pendingCheckpoint.stateHandle;

			if (pastCheckpointId <= checkpointId) {
				try {
					if (!committer.isCheckpointCommitted(subtaskId, pastCheckpointId)) {
						try (FSDataInputStream in = streamHandle.openInputStream()) {
							boolean success = sendValues(
									new ReusingMutableToRegularIteratorWrapper<>(
											new InputViewIterator<>(
													new DataInputViewStreamWrapper(
															in),
													serializer),
											serializer),
									pastCheckpointId,
									timestamp);
							if (success) {
								// in case the checkpoint was successfully committed,
								// discard its state from the backend and mark it for removal
								// in case it failed, we retry on the next checkpoint
								committer.commitCheckpoint(subtaskId, pastCheckpointId);
								streamHandle.discardState();
								pendingCheckpointIt.remove();
							}
						}
					} else {
						streamHandle.discardState();
						pendingCheckpointIt.remove();
					}
				} catch (Exception e) {
					// we have to break here to prevent a new (later) checkpoint
					// from being committed before this one
					LOG.error("Could not commit checkpoint.", e);
					break;
				}
			}
		}
	}
}