org.apache.flink.runtime.checkpoint.OperatorState Java Examples

@Test
public void testReadTime() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new KeyedProcessOperator<>(new StatefulFunctionWithTime()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new TimeReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new TimeReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 1, 2, 2, 3, 3), data);
}
 

private Collection<OperatorState> createOperatorState(OperatorID... operatorIds) {
	Collection<OperatorState> operatorStates = new ArrayList<>(operatorIds.length);

	for (OperatorID operatorId : operatorIds) {
		final OperatorState operatorState = new OperatorState(operatorId, 1, 42);
		final OperatorSubtaskState subtaskState = new OperatorSubtaskState(
			new OperatorStreamStateHandle(
				Collections.emptyMap(),
				new ByteStreamStateHandle("foobar", new byte[0])),
			null,
			null,
			null);
		operatorState.putState(0, subtaskState);
		operatorStates.add(operatorState);
	}

	return operatorStates;
}
 

@Test(expected = IOException.class)
public void testInvalidProcessReaderFunctionFails() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new ReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new InvalidReaderFunction();

	readInputSplit(split, userFunction);

	Assert.fail("KeyedStateReaderFunction did not fail on invalid RuntimeContext use");
}
 

@Test
public void testReadState() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), new Configuration(), new KeyedStateReaderOperator<>(new ReaderFunction(), Types.INT));
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new ReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 2, 3), data);
}
 

@Test
public void testCheckpointWithMasterAndTaskState() throws Exception {
	final Random rnd = new Random();

	final int maxNumMasterStates = 5;
	final int maxTaskStates = 20;
	final int maxNumSubtasks = 20;

	for (int i = 0; i < 100; ++i) {
		final long checkpointId = rnd.nextLong() & 0x7fffffffffffffffL;

		final int numTasks = rnd.nextInt(maxTaskStates) + 1;
		final int numSubtasks = rnd.nextInt(maxNumSubtasks) + 1;
		final Collection<OperatorState> taskStates =
				CheckpointTestUtils.createOperatorStates(rnd, numTasks, numSubtasks);

		final int numMasterStates = rnd.nextInt(maxNumMasterStates) + 1;
		final Collection<MasterState> masterStates =
				CheckpointTestUtils.createRandomMasterStates(rnd, numMasterStates);

		testCheckpointSerialization(checkpointId, taskStates, masterStates);
	}
}
 

@Test
public void testReadState() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new ReaderFunction());
	KeyGroupRangeInputSplit split = format.createInputSplits(1)[0];

	KeyedStateReaderFunction<Integer, Integer> userFunction = new ReaderFunction();

	List<Integer> data = readInputSplit(split, userFunction);

	Assert.assertEquals("Incorrect data read from input split", Arrays.asList(1, 2, 3), data);
}
 

@Test
public void testCheckpointWithOnlyMasterState() throws Exception {
	final Random rnd = new Random();
	final int maxNumMasterStates = 5;

	for (int i = 0; i < 100; ++i) {
		final long checkpointId = rnd.nextLong() & 0x7fffffffffffffffL;

		final Collection<OperatorState> operatorStates = Collections.emptyList();

		final int numMasterStates = rnd.nextInt(maxNumMasterStates) + 1;
		final Collection<MasterState> masterStates = 
				CheckpointTestUtils.createRandomMasterStates(rnd, numMasterStates);

		testCheckpointSerialization(checkpointId, operatorStates, masterStates);
	}
}
 

@Override
protected OperatorState deserializeOperatorState(DataInputStream dis, @Nullable DeserializationContext context) throws IOException {
	final OperatorID jobVertexId = new OperatorID(dis.readLong(), dis.readLong());
	final int parallelism = dis.readInt();
	final int maxParallelism = dis.readInt();

	// this field was "chain length" before Flink 1.3, and it is still part
	// of the format, despite being unused
	dis.readInt();

	// Add task state
	final OperatorState taskState = new OperatorState(jobVertexId, parallelism, maxParallelism);

	// Sub task states
	final int numSubTaskStates = dis.readInt();

	for (int j = 0; j < numSubTaskStates; j++) {
		final int subtaskIndex = dis.readInt();
		final OperatorSubtaskState subtaskState = deserializeSubtaskState(dis, context);
		taskState.putState(subtaskIndex, subtaskState);
	}

	return taskState;
}
 

@Override
protected void serializeOperatorState(OperatorState operatorState, DataOutputStream dos) throws IOException {
	// Operator ID
	dos.writeLong(operatorState.getOperatorID().getLowerPart());
	dos.writeLong(operatorState.getOperatorID().getUpperPart());

	// Parallelism
	dos.writeInt(operatorState.getParallelism());
	dos.writeInt(operatorState.getMaxParallelism());

	// Coordinator state
	serializeStreamStateHandle(operatorState.getCoordinatorState(), dos);

	// Sub task states
	final Map<Integer, OperatorSubtaskState> subtaskStateMap = operatorState.getSubtaskStates();
	dos.writeInt(subtaskStateMap.size());
	for (Map.Entry<Integer, OperatorSubtaskState> entry : subtaskStateMap.entrySet()) {
		dos.writeInt(entry.getKey());
		serializeSubtaskState(entry.getValue(), dos);
	}
}
 

@Test(expected = IllegalArgumentException.class)
public void testExistingSavepointEnforceUniqueUIDsWithOldSavepoint() throws IOException {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleStateBootstrapFunction());

	Collection<OperatorState> operatorStates = Collections.singletonList(new OperatorState(
		OperatorIDGenerator.fromUid(UID), 1, 4));

	SavepointMetadata metadata = new SavepointMetadata(4, Collections.emptyList(), operatorStates);

	new ExistingSavepoint(env, metadata, new MemoryStateBackend())
		.withOperator(UID, transformation)
		.write("");
}
 

protected void serializeMetadata(CheckpointMetadata checkpointMetadata, DataOutputStream dos) throws IOException {
	// first: checkpoint ID
	dos.writeLong(checkpointMetadata.getCheckpointId());

	// second: master state
	final Collection<MasterState> masterStates = checkpointMetadata.getMasterStates();
	dos.writeInt(masterStates.size());
	for (MasterState ms : masterStates) {
		serializeMasterState(ms, dos);
	}

	// third: operator states
	Collection<OperatorState> operatorStates = checkpointMetadata.getOperatorStates();
	dos.writeInt(operatorStates.size());

	for (OperatorState operatorState : operatorStates) {
		serializeOperatorState(operatorState, dos);
	}
}
 

/**
 * Read keyed state from an operator in a {@code Savepoint}.
 * @param uid The uid of the operator.
 * @param function The {@link KeyedStateReaderFunction} that is called for each key in state.
 * @param keyTypeInfo The type information of the key in state.
 * @param outTypeInfo The type information of the output of the transform reader function.
 * @param <K> The type of the key in state.
 * @param <OUT> The output type of the transform function.
 * @return A {@code DataSet} of objects read from keyed state.
 * @throws IOException If the savepoint does not contain operator state with the given uid.
 */
public <K, OUT> DataSet<OUT> readKeyedState(
	String uid,
	KeyedStateReaderFunction<K, OUT> function,
	TypeInformation<K> keyTypeInfo,
	TypeInformation<OUT> outTypeInfo) throws IOException {

	OperatorState operatorState = metadata.getOperatorState(uid);
	KeyedStateInputFormat<K, VoidNamespace, OUT> inputFormat = new KeyedStateInputFormat<>(
		operatorState,
		stateBackend,
		env.getConfiguration(),
		new KeyedStateReaderOperator<>(function, keyTypeInfo));

	return env.createInput(inputFormat, outTypeInfo);
}
 

@Test(expected = IllegalArgumentException.class)
public void testExistingSavepointEnforceUniqueUIDs() throws IOException {
	ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
	env.setParallelism(10);

	DataSource<Integer> input = env.fromElements(0);

	BootstrapTransformation<Integer> transformation = OperatorTransformation
		.bootstrapWith(input)
		.transform(new ExampleStateBootstrapFunction());

	Collection<OperatorState> operatorStates = Collections.singletonList(new OperatorState(
		OperatorIDGenerator.fromUid(UID), 1, 4));

	SavepointMetadata metadata = new SavepointMetadata(4, Collections.emptyList(), operatorStates);

	new ExistingSavepoint(env, metadata, new MemoryStateBackend())
		.withOperator(UID, transformation)
		.withOperator(UID, transformation);
}
 

@Nonnull
private List<Integer> readInputSplit(KeyGroupRangeInputSplit split, KeyedStateReaderFunction<Integer, Integer> userFunction) throws IOException {
	KeyedStateInputFormat<Integer, Integer> format = new KeyedStateInputFormat<>(
		new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4),
		new MemoryStateBackend(),
		Types.INT,
		userFunction);

	List<Integer> data = new ArrayList<>();

	format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));

	format.openInputFormat();
	format.open(split);

	while (!format.reachedEnd()) {
		data.add(format.nextRecord(0));
	}

	format.close();
	format.closeInputFormat();

	data.sort(Comparator.comparingInt(id -> id));
	return data;
}
 

private Collection<OperatorState> createOperatorState(OperatorID... operatorIds) {
	Collection<OperatorState> operatorStates = new ArrayList<>(operatorIds.length);

	for (OperatorID operatorId : operatorIds) {
		final OperatorState operatorState = new OperatorState(operatorId, 1, 42);
		final OperatorSubtaskState subtaskState = new OperatorSubtaskState(
			new OperatorStreamStateHandle(
				Collections.emptyMap(),
				new ByteStreamStateHandle("foobar", new byte[0])),
			null,
			null,
			null);
		operatorState.putState(0, subtaskState);
		operatorStates.add(operatorState);
	}

	return operatorStates;
}
 

@Override
public void dispose() throws Exception {
	for (OperatorState operatorState : operatorStates) {
		operatorState.discardState();
	}
	operatorStates.clear();
	masterStates.clear();
}
 

public SavepointV2(long checkpointId, Collection<OperatorState> operatorStates, Collection<MasterState> masterStates) {
	this(
		checkpointId,
		checkNotNull(operatorStates, "operatorStates"),
		null,
		masterStates
	);
}
 

@Test
public void testReadBroadcastState() throws Exception {
	try (TwoInputStreamOperatorTestHarness<Void, Integer, Void> testHarness = getTestHarness()) {
		testHarness.open();

		testHarness.processElement2(new StreamRecord<>(1));
		testHarness.processElement2(new StreamRecord<>(2));
		testHarness.processElement2(new StreamRecord<>(3));

		OperatorSubtaskState subtaskState = testHarness.snapshot(0, 0);
		OperatorState state = new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4);
		state.putState(0, subtaskState);

		OperatorStateInputSplit split = new OperatorStateInputSplit(subtaskState.getManagedOperatorState(), 0);

		BroadcastStateInputFormat<Integer, Integer> format = new BroadcastStateInputFormat<>(state, descriptor);

		format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
		format.open(split);

		Map<Integer, Integer> results = new HashMap<>(3);

		while (!format.reachedEnd()) {
			Tuple2<Integer, Integer> entry = format.nextRecord(null);
			results.put(entry.f0, entry.f1);
		}

		Map<Integer, Integer> expected = new HashMap<>(3);
		expected.put(1, 1);
		expected.put(2, 2);
		expected.put(3, 3);

		Assert.assertEquals("Failed to read correct list state from state backend", expected, results);
	}
}
 

private void testCheckpointSerialization(
		long checkpointId,
		Collection<OperatorState> operatorStates,
		Collection<MasterState> masterStates) throws IOException {

	SavepointV2Serializer serializer = SavepointV2Serializer.INSTANCE;

	ByteArrayOutputStreamWithPos baos = new ByteArrayOutputStreamWithPos();
	DataOutputStream out = new DataOutputViewStreamWrapper(baos);

	serializer.serialize(new SavepointV2(checkpointId, operatorStates, masterStates), out);
	out.close();

	byte[] bytes = baos.toByteArray();

	DataInputStream in = new DataInputViewStreamWrapper(new ByteArrayInputStreamWithPos(bytes));
	SavepointV2 deserialized = serializer.deserialize(in, getClass().getClassLoader());

	assertEquals(checkpointId, deserialized.getCheckpointId());
	assertEquals(operatorStates, deserialized.getOperatorStates());

	assertEquals(masterStates.size(), deserialized.getMasterStates().size());
	for (Iterator<MasterState> a = masterStates.iterator(), b = deserialized.getMasterStates().iterator();
			a.hasNext();)
	{
		CheckpointTestUtils.assertMasterStateEquality(a.next(), b.next());
	}
}
 

public static Optional<KeyedBackendSerializationProxy<?>> getKeyedBackendSerializationProxy(OperatorState opState) {
	try {
		KeyedStateHandle firstHandle = opState.getStates().iterator().next().getManagedKeyedState().iterator()
				.next();
		if (firstHandle instanceof IncrementalKeyedStateHandle) {
			return Optional.of(getKeyedBackendSerializationProxy(
					((IncrementalKeyedStateHandle) firstHandle).getMetaStateHandle()));
		} else {
			return Optional.of(getKeyedBackendSerializationProxy((StreamStateHandle) firstHandle));
		}
	} catch (Exception e) {
		return Optional.empty();
	}
}
 

private static byte[] serializeAsCheckpointMetadata(OperatorID id, byte[] coordinatorState) throws IOException {
	final OperatorState state = createOperatorState(id, coordinatorState);
	final CheckpointMetadata metadata = new CheckpointMetadata(
		1337L, Collections.singletonList(state), Collections.emptyList());

	final ByteArrayOutputStream out = new ByteArrayOutputStream();
	Checkpoints.storeCheckpointMetadata(metadata, out);
	return out.toByteArray();
}
 

@Test
public void testCheckpointWithNoState() throws Exception {
	final Random rnd = new Random();

	for (int i = 0; i < 100; ++i) {
		final long checkpointId = rnd.nextLong() & 0x7fffffffffffffffL;
		final Collection<OperatorState> taskStates = Collections.emptyList();
		final Collection<MasterState> masterStates = Collections.emptyList();

		testCheckpointSerialization(checkpointId, taskStates, masterStates);
	}
}
 

public SavepointMetadata(int maxParallelism, Collection<MasterState> masterStates, Collection<OperatorState> initialStates) {
	Preconditions.checkArgument(maxParallelism > 0
			&& maxParallelism <= UPPER_BOUND_MAX_PARALLELISM,
		"Maximum parallelism must be between 1 and " + UPPER_BOUND_MAX_PARALLELISM
			+ ". Found: " + maxParallelism);
	this.maxParallelism = maxParallelism;

	this.masterStates = Preconditions.checkNotNull(masterStates);

	this.operatorStateIndex = new HashMap<>(initialStates.size());
	initialStates.forEach(existingState -> operatorStateIndex.put(
		existingState.getOperatorID(),
		OperatorStateSpec.existing(existingState)));
}
 

/**
 * Simple test of savepoint methods.
 */
@Test
public void testSavepointV2() throws Exception {
	final Random rnd = new Random();

	final long checkpointId = rnd.nextInt(Integer.MAX_VALUE) + 1;
	final int numTaskStates = 4;
	final int numSubtaskStates = 16;
	final int numMasterStates = 7;

	Collection<OperatorState> taskStates =
			CheckpointTestUtils.createOperatorStates(rnd, numTaskStates, numSubtaskStates);

	Collection<MasterState> masterStates =
			CheckpointTestUtils.createRandomMasterStates(rnd, numMasterStates);

	SavepointV2 checkpoint = new SavepointV2(checkpointId, taskStates, masterStates);

	assertEquals(2, checkpoint.getVersion());
	assertEquals(checkpointId, checkpoint.getCheckpointId());
	assertEquals(taskStates, checkpoint.getOperatorStates());
	assertEquals(masterStates, checkpoint.getMasterStates());

	assertFalse(checkpoint.getOperatorStates().isEmpty());
	assertFalse(checkpoint.getMasterStates().isEmpty());

	checkpoint.dispose();

	assertTrue(checkpoint.getOperatorStates().isEmpty());
	assertTrue(checkpoint.getMasterStates().isEmpty());
}
 

@Override
public void reduce(Iterable<OperatorState> values, Collector<CheckpointMetadata> out) {
	CheckpointMetadata metadata =
		new CheckpointMetadata(
			SnapshotUtils.CHECKPOINT_ID,
			StreamSupport.stream(values.spliterator(), false).collect(Collectors.toList()),
			masterStates);

	out.collect(metadata);
}
 

@Test
public void testCreatePartitionedInputSplits() throws Exception {
	OperatorID operatorID = OperatorIDGenerator.fromUid("uid");

	OperatorSubtaskState state = createOperatorSubtaskState(new StreamFlatMap<>(new StatefulFunction()));
	OperatorState operatorState = new OperatorState(operatorID, 1, 128);
	operatorState.putState(0, state);

	KeyedStateInputFormat<?, ?> format = new KeyedStateInputFormat<>(operatorState, new MemoryStateBackend(), Types.INT, new ReaderFunction());
	KeyGroupRangeInputSplit[] splits = format.createInputSplits(4);
	Assert.assertEquals("Failed to properly partition operator state into input splits", 4, splits.length);
}
 

private DataSet<OperatorState> writeOperatorStates(
		List<BootstrapTransformationWithID<?>> newOperatorStates,
		Path savepointWritePath) {
	return newOperatorStates
		.stream()
		.map(newOperatorState -> newOperatorState
			.getBootstrapTransformation()
			.writeOperatorState(newOperatorState.getOperatorID(), stateBackend, metadata.getMaxParallelism(), savepointWritePath))
		.reduce(DataSet::union)
		.orElseThrow(() -> new IllegalStateException("Savepoint's must contain at least one operator"));
}
 

/**
 * Loads an existing savepoint. Useful if you want to query, modify, or extend
 * the state of an existing application.
 *
 * @param env The execution environment used to transform the savepoint.
 * @param path The path to an existing savepoint on disk.
 * @param stateBackend The state backend of the savepoint.
 */
public static ExistingSavepoint load(ExecutionEnvironment env, String path, StateBackend stateBackend) throws IOException {
	CheckpointMetadata metadata = SavepointLoader.loadSavepointMetadata(path);

	int maxParallelism = metadata
		.getOperatorStates()
		.stream()
		.map(OperatorState::getMaxParallelism)
		.max(Comparator.naturalOrder())
		.orElseThrow(() -> new RuntimeException("Savepoint must contain at least one operator state."));

	SavepointMetadata savepointMetadata = new SavepointMetadata(maxParallelism, metadata.getMasterStates(), metadata.getOperatorStates());
	return new ExistingSavepoint(env, savepointMetadata, stateBackend);
}
 

private static KeyGroupRangeInputSplit createKeyGroupRangeInputSplit(
	OperatorState operatorState,
	int maxParallelism,
	KeyGroupRange keyGroupRange,
	Integer index) {

	final List<KeyedStateHandle> managedKeyedState = StateAssignmentOperation.getManagedKeyedStateHandles(operatorState, keyGroupRange);
	final List<KeyedStateHandle> rawKeyedState = StateAssignmentOperation.getRawKeyedStateHandles(operatorState, keyGroupRange);

	return new KeyGroupRangeInputSplit(managedKeyedState, rawKeyedState, maxParallelism, index);
}
 

private DataSet<OperatorState> unionOperatorStates(DataSet<OperatorState> newOperatorStates, List<OperatorState> existingOperators) {
	DataSet<OperatorState> finalOperatorStates;
	if (existingOperators.isEmpty()) {
		finalOperatorStates = newOperatorStates;
	} else {
		DataSet<OperatorState> wrappedCollection = newOperatorStates
			.getExecutionEnvironment()
			.fromCollection(existingOperators);

		finalOperatorStates = newOperatorStates.union(wrappedCollection);
	}
	return finalOperatorStates;
}