org.apache.flink.runtime.state.memory.ByteStreamStateHandle Java Examples

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
public void testRepartitionBroadcastState() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
	metaInfoMap1.put("t-5", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.BROADCAST));
	metaInfoMap1.put("t-6", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.BROADCAST));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[60]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(2);
	metaInfoMap2.put("t-5", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.BROADCAST));
	metaInfoMap2.put("t-6", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.BROADCAST));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[60]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

/**
 * Creates a deep copy of the given {@link OperatorStreamStateHandle}.
 */
public static OperatorStateHandle deepDummyCopy(OperatorStateHandle original) {

	if (original == null) {
		return null;
	}

	ByteStreamStateHandle stateHandleCopy = cloneByteStreamStateHandle((ByteStreamStateHandle) original.getDelegateStateHandle());
	Map<String, OperatorStateHandle.StateMetaInfo> offsets = original.getStateNameToPartitionOffsets();
	Map<String, OperatorStateHandle.StateMetaInfo> offsetsCopy = new HashMap<>(offsets.size());

	for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> entry : offsets.entrySet()) {
		OperatorStateHandle.StateMetaInfo metaInfo = entry.getValue();
		OperatorStateHandle.StateMetaInfo metaInfoCopy =
			new OperatorStateHandle.StateMetaInfo(metaInfo.getOffsets(), metaInfo.getDistributionMode());
		offsetsCopy.put(String.valueOf(entry.getKey()), metaInfoCopy);
	}
	return new OperatorStreamStateHandle(offsetsCopy, stateHandleCopy);
}
 

public static StreamStateHandle deserializeStreamStateHandle(DataInputStream dis) throws IOException {
	final int type = dis.read();
	if (NULL_HANDLE == type) {
		return null;
	} else if (FILE_STREAM_STATE_HANDLE == type) {
		long size = dis.readLong();
		String pathString = dis.readUTF();
		return new FileStateHandle(new Path(pathString), size);
	} else if (BYTE_STREAM_STATE_HANDLE == type) {
		String handleName = dis.readUTF();
		int numBytes = dis.readInt();
		byte[] data = new byte[numBytes];
		dis.readFully(data);
		return new ByteStreamStateHandle(handleName, data);
	} else {
		throw new IOException("Unknown implementation of StreamStateHandle, code: " + type);
	}
}
 

@VisibleForTesting
public static void serializeStreamStateHandle(
		StreamStateHandle stateHandle, DataOutputStream dos) throws IOException {

	if (stateHandle == null) {
		dos.writeByte(NULL_HANDLE);

	} else if (stateHandle instanceof FileStateHandle) {
		dos.writeByte(FILE_STREAM_STATE_HANDLE);
		FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;
		dos.writeLong(stateHandle.getStateSize());
		dos.writeUTF(fileStateHandle.getFilePath().toString());

	} else if (stateHandle instanceof ByteStreamStateHandle) {
		dos.writeByte(BYTE_STREAM_STATE_HANDLE);
		ByteStreamStateHandle byteStreamStateHandle = (ByteStreamStateHandle) stateHandle;
		dos.writeUTF(byteStreamStateHandle.getHandleName());
		byte[] internalData = byteStreamStateHandle.getData();
		dos.writeInt(internalData.length);
		dos.write(byteStreamStateHandle.getData());
	} else {
		throw new IOException("Unknown implementation of StreamStateHandle: " + stateHandle.getClass());
	}

	dos.flush();
}
 

@VisibleForTesting
public static StreamStateHandle deserializeStreamStateHandle(DataInputStream dis) throws IOException {
	int type = dis.read();
	if (NULL_HANDLE == type) {
		return null;
	} else if (FILE_STREAM_STATE_HANDLE == type) {
		long size = dis.readLong();
		String pathString = dis.readUTF();
		return new FileStateHandle(new Path(pathString), size);
	} else if (BYTE_STREAM_STATE_HANDLE == type) {
		String handleName = dis.readUTF();
		int numBytes = dis.readInt();
		byte[] data = new byte[numBytes];
		dis.readFully(data);
		return new ByteStreamStateHandle(handleName, data);
	} else {
		throw new IOException("Unknown implementation of StreamStateHandle, code: " + type);
	}
}
 

@Test
public void testRepartitionUnionState() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
	metaInfoMap1.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[]{0}, OperatorStateHandle.Mode.UNION));
	metaInfoMap1.put("t-4", new OperatorStateHandle.StateMetaInfo(new long[]{22, 44}, OperatorStateHandle.Mode.UNION));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[50]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
	metaInfoMap2.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[]{0}, OperatorStateHandle.Mode.UNION));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[20]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

@Test
public void testRepartitionSplitDistributeStates() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(1);
	metaInfoMap1.put("t-1", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[30]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
	metaInfoMap2.put("t-2", new OperatorStateHandle.StateMetaInfo(new long[]{0, 15}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[40]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

public static KeyGroupsStateHandle generateKeyGroupState(
		KeyGroupRange keyGroupRange,
		List<? extends Serializable> states) throws IOException {

	Preconditions.checkArgument(keyGroupRange.getNumberOfKeyGroups() == states.size());

	Tuple2<byte[], List<long[]>> serializedDataWithOffsets =
			serializeTogetherAndTrackOffsets(Collections.<List<? extends Serializable>>singletonList(states));

	KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupRange, serializedDataWithOffsets.f1.get(0));

	ByteStreamStateHandle allSerializedStatesHandle = new ByteStreamStateHandle(
			String.valueOf(UUID.randomUUID()),
			serializedDataWithOffsets.f0);

	return new KeyGroupsStateHandle(keyGroupRangeOffsets, allSerializedStatesHandle);
}
 

/**
 * Creates a deep copy of the given {@link OperatorStreamStateHandle}.
 */
public static OperatorStateHandle deepDummyCopy(OperatorStateHandle original) {

	if (original == null) {
		return null;
	}

	ByteStreamStateHandle stateHandle = (ByteStreamStateHandle) original.getDelegateStateHandle();
	ByteStreamStateHandle stateHandleCopy = new ByteStreamStateHandle(
		String.valueOf(stateHandle.getHandleName()),
		stateHandle.getData().clone());
	Map<String, OperatorStateHandle.StateMetaInfo> offsets = original.getStateNameToPartitionOffsets();
	Map<String, OperatorStateHandle.StateMetaInfo> offsetsCopy = new HashMap<>(offsets.size());

	for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> entry : offsets.entrySet()) {
		OperatorStateHandle.StateMetaInfo metaInfo = entry.getValue();
		OperatorStateHandle.StateMetaInfo metaInfoCopy =
			new OperatorStateHandle.StateMetaInfo(metaInfo.getOffsets(), metaInfo.getDistributionMode());
		offsetsCopy.put(String.valueOf(entry.getKey()), metaInfoCopy);
	}
	return new OperatorStreamStateHandle(offsetsCopy, stateHandleCopy);
}
 

/**
 * Creates a new test {@link OperatorStreamStateHandle} with a given number of randomly created named states.
 */
public static OperatorStateHandle createNewOperatorStateHandle(int numNamedStates, Random random) {
	Map<String, OperatorStateHandle.StateMetaInfo> operatorStateMetaData = new HashMap<>(numNamedStates);
	byte[] streamData = new byte[numNamedStates * 4];
	random.nextBytes(streamData);
	long off = 0;
	for (int i = 0; i < numNamedStates; ++i) {
		long[] offsets = new long[4];
		for (int o = 0; o < offsets.length; ++o) {
			offsets[o] = off++;
		}
		OperatorStateHandle.StateMetaInfo metaInfo =
			new OperatorStateHandle.StateMetaInfo(offsets, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE);
		operatorStateMetaData.put(String.valueOf(UUID.randomUUID()), metaInfo);
	}
	ByteStreamStateHandle byteStreamStateHandle =
		new ByteStreamStateHandle(String.valueOf(UUID.randomUUID()), streamData);
	return new OperatorStreamStateHandle(operatorStateMetaData, byteStreamStateHandle);
}
 

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;
}
 

private <I, H extends AbstractChannelStateHandle<I>> H createHandle(
		HandleFactory<I, H> handleFactory,
		StreamStateHandle underlying,
		I channelInfo,
		StateContentMetaInfo contentMetaInfo) throws IOException {
	Optional<byte[]> bytes = underlying.asBytesIfInMemory(); // todo: consider restructuring channel state and removing this method: https://issues.apache.org/jira/browse/FLINK-17972
	if (bytes.isPresent()) {
		StreamStateHandle extracted = new ByteStreamStateHandle(
			randomUUID().toString(),
			serializer.extractAndMerge(bytes.get(), contentMetaInfo.getOffsets()));
		return handleFactory.create(
			channelInfo,
			extracted,
			singletonList(serializer.getHeaderLength()),
			extracted.getStateSize());
	} else {
		return handleFactory.create(channelInfo, underlying, contentMetaInfo.getOffsets(), contentMetaInfo.getSize());
	}
}
 

@VisibleForTesting
public static void serializeStreamStateHandle(
		StreamStateHandle stateHandle, DataOutputStream dos) throws IOException {

	if (stateHandle == null) {
		dos.writeByte(NULL_HANDLE);

	} else if (stateHandle instanceof FileStateHandle) {
		dos.writeByte(FILE_STREAM_STATE_HANDLE);
		FileStateHandle fileStateHandle = (FileStateHandle) stateHandle;
		dos.writeLong(stateHandle.getStateSize());
		dos.writeUTF(fileStateHandle.getFilePath().toString());

	} else if (stateHandle instanceof ByteStreamStateHandle) {
		dos.writeByte(BYTE_STREAM_STATE_HANDLE);
		ByteStreamStateHandle byteStreamStateHandle = (ByteStreamStateHandle) stateHandle;
		dos.writeUTF(byteStreamStateHandle.getHandleName());
		byte[] internalData = byteStreamStateHandle.getData();
		dos.writeInt(internalData.length);
		dos.write(byteStreamStateHandle.getData());
	} else {
		throw new IOException("Unknown implementation of StreamStateHandle: " + stateHandle.getClass());
	}

	dos.flush();
}
 

public static StreamStateHandle deserializeStreamStateHandle(DataInputStream dis) throws IOException {
	final int type = dis.read();
	if (NULL_HANDLE == type) {
		return null;
	} else if (FILE_STREAM_STATE_HANDLE == type) {
		long size = dis.readLong();
		String pathString = dis.readUTF();
		return new FileStateHandle(new Path(pathString), size);
	} else if (BYTE_STREAM_STATE_HANDLE == type) {
		String handleName = dis.readUTF();
		int numBytes = dis.readInt();
		byte[] data = new byte[numBytes];
		dis.readFully(data);
		return new ByteStreamStateHandle(handleName, data);
	} else {
		throw new IOException("Unknown implementation of StreamStateHandle, code: " + type);
	}
}
 

@Test
public void testRepartitionUnionState() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
	metaInfoMap1.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[]{0}, OperatorStateHandle.Mode.UNION));
	metaInfoMap1.put("t-4", new OperatorStateHandle.StateMetaInfo(new long[]{22, 44}, OperatorStateHandle.Mode.UNION));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[50]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
	metaInfoMap2.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[]{0}, OperatorStateHandle.Mode.UNION));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[20]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

@VisibleForTesting
public static StreamStateHandle deserializeStreamStateHandle(DataInputStream dis) throws IOException {
	int type = dis.read();
	if (NULL_HANDLE == type) {
		return null;
	} else if (FILE_STREAM_STATE_HANDLE == type) {
		long size = dis.readLong();
		String pathString = dis.readUTF();
		return new FileStateHandle(new Path(pathString), size);
	} else if (BYTE_STREAM_STATE_HANDLE == type) {
		String handleName = dis.readUTF();
		int numBytes = dis.readInt();
		byte[] data = new byte[numBytes];
		dis.readFully(data);
		return new ByteStreamStateHandle(handleName, data);
	} else {
		throw new IOException("Unknown implementation of StreamStateHandle, code: " + type);
	}
}
 

private Collection<OperatorState> createOperatorState(OperatorID... operatorIds) {
	Random random = new Random();
	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,
			singleton(createNewInputChannelStateHandle(10, random)),
			singleton(createNewResultSubpartitionStateHandle(10, random)));
		operatorState.putState(0, subtaskState);
		operatorStates.add(operatorState);
	}

	return operatorStates;
}
 

private static CompletedCheckpointStorageLocation createSavepointWithOperatorSubtaskState(
		final long checkpointId,
		final OperatorID operatorId,
		final int parallelism) throws IOException {

	final Random rnd = new Random();

	final OperatorSubtaskState subtaskState = new OperatorSubtaskState(
		new OperatorStreamStateHandle(Collections.emptyMap(), new ByteStreamStateHandle("testHandler", new byte[0])),
		null,
		null,
		null,
		singleton(createNewInputChannelStateHandle(10, rnd)),
		singleton(createNewResultSubpartitionStateHandle(10, rnd)));

	final OperatorState state = new OperatorState(operatorId, parallelism, parallelism);
	state.putState(0, subtaskState);

	return createSavepointWithOperatorState(checkpointId, state);
}
 

/**
 * Tests that savepoint loading fails when there is non-restored coordinator state only,
 * and non-restored state is not allowed.
 */
@Test
public void testUnmatchedCoordinatorOnlyStateFails() throws Exception {
	final OperatorID operatorID = new OperatorID();
	final int maxParallelism = 1234;

	final OperatorState state = new OperatorState(operatorID, maxParallelism / 2, maxParallelism);
	state.setCoordinatorState(new ByteStreamStateHandle("coordinatorState", new byte[0]));

	final CompletedCheckpointStorageLocation testSavepoint = createSavepointWithOperatorState(42L, state);
	final Map<JobVertexID, ExecutionJobVertex> tasks = Collections.emptyMap();

	try {
		Checkpoints.loadAndValidateCheckpoint(new JobID(), tasks, testSavepoint, cl, false);
		fail("Did not throw expected Exception");
	} catch (IllegalStateException expected) {
		assertTrue(expected.getMessage().contains("allowNonRestoredState"));
	}
}
 

/**
 * Creates a new test {@link OperatorStreamStateHandle} with a given number of randomly created named states.
 */
public static OperatorStateHandle createNewOperatorStateHandle(int numNamedStates, Random random) {
	Map<String, OperatorStateHandle.StateMetaInfo> operatorStateMetaData = new HashMap<>(numNamedStates);
	byte[] streamData = new byte[numNamedStates * 4];
	random.nextBytes(streamData);
	long off = 0;
	for (int i = 0; i < numNamedStates; ++i) {
		long[] offsets = new long[4];
		for (int o = 0; o < offsets.length; ++o) {
			offsets[o] = off++;
		}
		OperatorStateHandle.StateMetaInfo metaInfo =
			new OperatorStateHandle.StateMetaInfo(offsets, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE);
		operatorStateMetaData.put(String.valueOf(UUID.randomUUID()), metaInfo);
	}
	ByteStreamStateHandle byteStreamStateHandle =
		new ByteStreamStateHandle(String.valueOf(UUID.randomUUID()), streamData);
	return new OperatorStreamStateHandle(operatorStateMetaData, byteStreamStateHandle);
}
 

/**
 * Creates a deep copy of the given {@link OperatorStreamStateHandle}.
 */
public static OperatorStateHandle deepDummyCopy(OperatorStateHandle original) {

	if (original == null) {
		return null;
	}

	ByteStreamStateHandle stateHandle = (ByteStreamStateHandle) original.getDelegateStateHandle();
	ByteStreamStateHandle stateHandleCopy = new ByteStreamStateHandle(
		String.valueOf(stateHandle.getHandleName()),
		stateHandle.getData().clone());
	Map<String, OperatorStateHandle.StateMetaInfo> offsets = original.getStateNameToPartitionOffsets();
	Map<String, OperatorStateHandle.StateMetaInfo> offsetsCopy = new HashMap<>(offsets.size());

	for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> entry : offsets.entrySet()) {
		OperatorStateHandle.StateMetaInfo metaInfo = entry.getValue();
		OperatorStateHandle.StateMetaInfo metaInfoCopy =
			new OperatorStateHandle.StateMetaInfo(metaInfo.getOffsets(), metaInfo.getDistributionMode());
		offsetsCopy.put(String.valueOf(entry.getKey()), metaInfoCopy);
	}
	return new OperatorStreamStateHandle(offsetsCopy, stateHandleCopy);
}
 

public static KeyGroupsStateHandle generateKeyGroupState(
		KeyGroupRange keyGroupRange,
		List<? extends Serializable> states) throws IOException {

	Preconditions.checkArgument(keyGroupRange.getNumberOfKeyGroups() == states.size());

	Tuple2<byte[], List<long[]>> serializedDataWithOffsets =
			serializeTogetherAndTrackOffsets(Collections.<List<? extends Serializable>>singletonList(states));

	KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(keyGroupRange, serializedDataWithOffsets.f1.get(0));

	ByteStreamStateHandle allSerializedStatesHandle = new ByteStreamStateHandle(
			String.valueOf(UUID.randomUUID()),
			serializedDataWithOffsets.f0);

	return new KeyGroupsStateHandle(keyGroupRangeOffsets, allSerializedStatesHandle);
}
 

@Test
public void testRepartitionSplitDistributeStates() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(1);
	metaInfoMap1.put("t-1", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[30]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
	metaInfoMap2.put("t-2", new OperatorStateHandle.StateMetaInfo(new long[]{0, 15}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[40]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

@Test
public void testRepartitionUnionState() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
	metaInfoMap1.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[]{0}, OperatorStateHandle.Mode.UNION));
	metaInfoMap1.put("t-4", new OperatorStateHandle.StateMetaInfo(new long[]{22, 44}, OperatorStateHandle.Mode.UNION));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[50]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
	metaInfoMap2.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[]{0}, OperatorStateHandle.Mode.UNION));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[20]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

@Test
public void testRepartitionBroadcastState() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
	metaInfoMap1.put("t-5", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.BROADCAST));
	metaInfoMap1.put("t-6", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.BROADCAST));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[60]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(2);
	metaInfoMap2.put("t-5", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.BROADCAST));
	metaInfoMap2.put("t-6", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.BROADCAST));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[60]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

@Test
public void testRepartitionBroadcastState() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
	metaInfoMap1.put("t-5", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.BROADCAST));
	metaInfoMap1.put("t-6", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.BROADCAST));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[60]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(2);
	metaInfoMap2.put("t-5", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10, 20}, OperatorStateHandle.Mode.BROADCAST));
	metaInfoMap2.put("t-6", new OperatorStateHandle.StateMetaInfo(new long[]{30, 40, 50}, OperatorStateHandle.Mode.BROADCAST));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[60]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

@Test
public void testRepartitionSplitDistributeStates() {
	OperatorID operatorID = new OperatorID();
	OperatorState operatorState = new OperatorState(operatorID, 2, 4);

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(1);
	metaInfoMap1.put("t-1", new OperatorStateHandle.StateMetaInfo(new long[]{0, 10}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
	OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[30]));
	operatorState.putState(0, new OperatorSubtaskState(osh1, null, null, null, null, null));

	Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
	metaInfoMap2.put("t-2", new OperatorStateHandle.StateMetaInfo(new long[]{0, 15}, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
	OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[40]));
	operatorState.putState(1, new OperatorSubtaskState(osh2, null, null, null, null, null));

	verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}
 

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

public static InputChannelStateHandle deepDummyCopy(InputChannelStateHandle original) {
	if (original == null) {
		return null;
	}
	return new InputChannelStateHandle(
		new InputChannelInfo(original.getInfo().getGateIdx(), original.getInfo().getInputChannelIdx()),
		cloneByteStreamStateHandle((ByteStreamStateHandle) original.getDelegate()),
		new ArrayList<>(original.getOffsets()));
}