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

private static <T extends StateObject> void performCheck(
	Future<SnapshotResult<T>> resultFuture,
	StateObjectCollection<T> jmState,
	StateObjectCollection<T> tmState) {

	SnapshotResult<T> snapshotResult;
	try {
		snapshotResult = resultFuture.get();
	} catch (Exception e) {
		throw new RuntimeException(e);
	}

	Assert.assertEquals(
		snapshotResult.getJobManagerOwnedSnapshot(),
		jmState.iterator().next());

	Assert.assertEquals(
		snapshotResult.getTaskLocalSnapshot(),
		tmState.iterator().next());
}
 

public static <T extends StateObject> Map<OperatorInstanceID, List<T>>  reDistributePartitionableStates(
		List<OperatorState> oldOperatorStates,
		int newParallelism,
		List<OperatorIDPair> newOperatorIDs,
		Function<OperatorSubtaskState, StateObjectCollection<T>> extractHandle,
		OperatorStateRepartitioner<T> stateRepartitioner) {

	//TODO: rewrite this method to only use OperatorID
	checkState(newOperatorIDs.size() == oldOperatorStates.size(),
		"This method still depends on the order of the new and old operators");

	// The nested list wraps as the level of operator -> subtask -> state object collection
	List<List<List<T>>> oldStates = splitManagedAndRawOperatorStates(oldOperatorStates, extractHandle);

	Map<OperatorInstanceID, List<T>> result = new HashMap<>();
	for (int operatorIndex = 0; operatorIndex < newOperatorIDs.size(); operatorIndex++) {
		result.putAll(applyRepartitioner(
			newOperatorIDs.get(operatorIndex).getGeneratedOperatorID(),
			stateRepartitioner,
			oldStates.get(operatorIndex),
			oldOperatorStates.get(operatorIndex).getParallelism(),
			newParallelism));
	}

	return result;
}
 

/**
 * Returns true iff, in iteration order, all objects in the first collection are equal by reference to their
 * corresponding object (by order) in the second collection and the size of the collections is equal.
 */
public boolean checkContainedObjectsReferentialEquality(StateObjectCollection<?> a, StateObjectCollection<?> b) {

	if (a == b) {
		return true;
	}

	if (a == null || b == null) {
		return false;
	}

	if (a.size() != b.size()) {
		return false;
	}

	Iterator<?> bIter = b.iterator();
	for (StateObject stateObject : a) {
		if (!bIter.hasNext() || bIter.next() != stateObject) {
			return false;
		}
	}
	return true;
}
 

@Override
public void discardState() {
	try {
		List<StateObject> toDispose =
			new ArrayList<>(
					managedOperatorState.size() +
					rawOperatorState.size() +
					managedKeyedState.size() +
					rawKeyedState.size());
		toDispose.addAll(managedOperatorState);
		toDispose.addAll(rawOperatorState);
		toDispose.addAll(managedKeyedState);
		toDispose.addAll(rawKeyedState);
		StateUtil.bestEffortDiscardAllStateObjects(toDispose);
	} catch (Exception e) {
		LOG.warn("Error while discarding operator states.", e);
	}
}
 

@Test
public void ensureStateHandlesHaveSerialVersionUID() {
	try {
		Reflections reflections = new Reflections("org.apache.flink");

		// check all state handles

		@SuppressWarnings("unchecked")
		Set<Class<?>> stateHandleImplementations = (Set<Class<?>>) (Set<?>)
				reflections.getSubTypesOf(StateObject.class);

		for (Class<?> clazz : stateHandleImplementations) {
			validataSerialVersionUID(clazz);
		}
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

private static <T extends StateObject> void performCheck(
	Future<SnapshotResult<T>> resultFuture,
	StateObjectCollection<T> jmState,
	StateObjectCollection<T> tmState) {

	SnapshotResult<T> snapshotResult;
	try {
		snapshotResult = resultFuture.get();
	} catch (Exception e) {
		throw new RuntimeException(e);
	}

	Assert.assertEquals(
		snapshotResult.getJobManagerOwnedSnapshot(),
		jmState.iterator().next());

	Assert.assertEquals(
		snapshotResult.getTaskLocalSnapshot(),
		tmState.iterator().next());
}
 

private static <T extends StateObject> List<List<List<T>>> splitManagedAndRawOperatorStates(
		List<OperatorState> operatorStates,
		Function<OperatorSubtaskState, StateObjectCollection<T>> extracthandle) {
	List<List<List<T>>> result = new ArrayList<>();

	for (OperatorState operatorState : operatorStates) {
		List<List<T>> statePerSubtask = new ArrayList<>(operatorState.getParallelism());

		for (int subTaskIndex = 0; subTaskIndex < operatorState.getParallelism(); subTaskIndex++) {
			OperatorSubtaskState subtaskState = operatorState.getState(subTaskIndex);
			statePerSubtask.add(subtaskState == null ? emptyList() : extracthandle.apply(subtaskState).asList());
		}
		result.add(statePerSubtask);
	}
	return result;
}
 

@Test
public void ensureStateHandlesHaveSerialVersionUID() {
	try {
		Reflections reflections = new Reflections("org.apache.flink");

		// check all state handles

		@SuppressWarnings("unchecked")
		Set<Class<?>> stateHandleImplementations = (Set<Class<?>>) (Set<?>)
				reflections.getSubTypesOf(StateObject.class);

		for (Class<?> clazz : stateHandleImplementations) {
			validataSerialVersionUID(clazz);
		}
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

public static <T extends StateObject> Map<OperatorInstanceID, List<T>> applyRepartitioner(
		OperatorID operatorID,
		OperatorStateRepartitioner<T> opStateRepartitioner,
		List<List<T>> chainOpParallelStates,
		int oldParallelism,
		int newParallelism) {

	List<List<T>> states = applyRepartitioner(
		opStateRepartitioner,
		chainOpParallelStates,
		oldParallelism,
		newParallelism);

	Map<OperatorInstanceID, List<T>> result = new HashMap<>(states.size());

	for (int subtaskIndex = 0; subtaskIndex < states.size(); subtaskIndex++) {
		checkNotNull(states.get(subtaskIndex) != null, "states.get(subtaskIndex) is null");
		result.put(OperatorInstanceID.of(subtaskIndex, operatorID), states.get(subtaskIndex));
	}

	return result;
}
 

/**
 * Creates a deep copy of the first state object in the given collection, or null if the collection is empy.
 */
private <T extends StateObject> T deepCopyFirstElement(StateObjectCollection<T> original) {
	if (original.isEmpty()) {
		return null;
	}

	T stateObject = original.iterator().next();
	StateObject result;
	if (stateObject instanceof OperatorStreamStateHandle) {
		result = deepDummyCopy((OperatorStateHandle) stateObject);
	} else if (stateObject instanceof KeyedStateHandle) {
		result = deepDummyCopy((KeyedStateHandle) stateObject);
	} else {
		throw new IllegalStateException();
	}
	return (T) result;
}
 

/**
 * Returns true iff, in iteration order, all objects in the first collection are equal by reference to their
 * corresponding object (by order) in the second collection and the size of the collections is equal.
 */
public boolean checkContainedObjectsReferentialEquality(StateObjectCollection<?> a, StateObjectCollection<?> b) {

	if (a == b) {
		return true;
	}

	if(a == null || b == null) {
		return false;
	}

	if (a.size() != b.size()) {
		return false;
	}

	Iterator<?> bIter = b.iterator();
	for (StateObject stateObject : a) {
		if (!bIter.hasNext() || bIter.next() != stateObject) {
			return false;
		}
	}
	return true;
}
 

/**
 * Creates a deep copy of the first state object in the given collection, or null if the collection is empy.
 */
private <T extends StateObject> T deepCopyFirstElement(StateObjectCollection<T> original) {
	if (original.isEmpty()) {
		return null;
	}

	T stateObject = original.iterator().next();
	StateObject result;
	if (stateObject instanceof OperatorStreamStateHandle) {
		result = deepDummyCopy((OperatorStateHandle) stateObject);
	} else if (stateObject instanceof KeyedStateHandle) {
		result = deepDummyCopy((KeyedStateHandle) stateObject);
	} else {
		throw new IllegalStateException();
	}
	return (T) result;
}
 

public void cancel() throws Exception {
	List<Tuple2<Future<? extends StateObject>, String>> pairs = new ArrayList<>();
	pairs.add(new Tuple2<>(getKeyedStateManagedFuture(), "managed keyed"));
	pairs.add(new Tuple2<>(getKeyedStateRawFuture(), "managed operator"));
	pairs.add(new Tuple2<>(getOperatorStateManagedFuture(), "raw keyed"));
	pairs.add(new Tuple2<>(getOperatorStateRawFuture(), "raw operator"));
	pairs.add(new Tuple2<>(getInputChannelStateFuture(), "input channel"));
	pairs.add(new Tuple2<>(getResultSubpartitionStateFuture(), "result subpartition"));
	try (Closer closer = Closer.create()) {
		for (Tuple2<Future<? extends StateObject>, String> pair : pairs) {
			closer.register(() -> {
				try {
					discardStateFuture(pair.f0);
				} catch (Exception e) {
					throw new RuntimeException(String.format("Could not properly cancel %s state future", pair.f1), e);
				}
			});
		}
	}
}
 

@Test
public void ensureStateHandlesHaveSerialVersionUID() {
	try {
		Reflections reflections = new Reflections("org.apache.flink");

		// check all state handles

		@SuppressWarnings("unchecked")
		Set<Class<?>> stateHandleImplementations = (Set<Class<?>>) (Set<?>)
				reflections.getSubTypesOf(StateObject.class);

		for (Class<?> clazz : stateHandleImplementations) {
			validataSerialVersionUID(clazz);
		}
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Override
public void discardState() {
	try {
		List<StateObject> toDispose =
			new ArrayList<>(
					managedOperatorState.size() +
					rawOperatorState.size() +
					managedKeyedState.size() +
					rawKeyedState.size());
		toDispose.addAll(managedOperatorState);
		toDispose.addAll(rawOperatorState);
		toDispose.addAll(managedKeyedState);
		toDispose.addAll(rawKeyedState);
		StateUtil.bestEffortDiscardAllStateObjects(toDispose);
	} catch (Exception e) {
		LOG.warn("Error while discarding operator states.", e);
	}
}
 

/**
 * Returns true iff, in iteration order, all objects in the first collection are equal by reference to their
 * corresponding object (by order) in the second collection and the size of the collections is equal.
 */
public boolean checkContainedObjectsReferentialEquality(StateObjectCollection<?> a, StateObjectCollection<?> b) {

	if (a == b) {
		return true;
	}

	if(a == null || b == null) {
		return false;
	}

	if (a.size() != b.size()) {
		return false;
	}

	Iterator<?> bIter = b.iterator();
	for (StateObject stateObject : a) {
		if (!bIter.hasNext() || bIter.next() != stateObject) {
			return false;
		}
	}
	return true;
}
 

private static <T extends StateObject> void performCheck(
	Future<SnapshotResult<T>> resultFuture,
	StateObjectCollection<T> jmState,
	StateObjectCollection<T> tmState) {

	SnapshotResult<T> snapshotResult;
	try {
		snapshotResult = resultFuture.get();
	} catch (Exception e) {
		throw new RuntimeException(e);
	}

	Assert.assertEquals(
		snapshotResult.getJobManagerOwnedSnapshot(),
		jmState.iterator().next());

	Assert.assertEquals(
		snapshotResult.getTaskLocalSnapshot(),
		tmState.iterator().next());
}
 

@Override
public void discardState() {
	try {
		List<StateObject> toDispose =
			new ArrayList<>(
					managedOperatorState.size() +
					rawOperatorState.size() +
					managedKeyedState.size() +
					rawKeyedState.size() +
					inputChannelState.size() +
					resultSubpartitionState.size());
		toDispose.addAll(managedOperatorState);
		toDispose.addAll(rawOperatorState);
		toDispose.addAll(managedKeyedState);
		toDispose.addAll(rawKeyedState);
		toDispose.addAll(collectUniqueDelegates(inputChannelState, resultSubpartitionState));
		StateUtil.bestEffortDiscardAllStateObjects(toDispose);
	} catch (Exception e) {
		LOG.warn("Error while discarding operator states.", e);
	}
}
 

/**
 * Returns true if this contains at least one {@link StateObject}.
 */
public boolean hasState() {
	for (StateObject state : stateObjects) {
		if (state != null) {
			return true;
		}
	}
	return false;
}
 

static <T extends StateObject> StateObjectCollection<T> deserializeCollection(
	DataInputStream dis,
	DeserializationContext context,
	BiFunctionWithException<DataInputStream, DeserializationContext, T, IOException> s) throws IOException {

	int size = dis.readInt();
	List<T> result = new ArrayList<>();
	for (int i = 0; i < size; i++) {
		result.add(s.apply(dis, context));
	}
	return new StateObjectCollection<>(result);
}
 

/**
 * Returns true if this contains at least one {@link StateObject}.
 */
public boolean hasState() {
	for (StateObject state : stateObjects) {
		if (state != null) {
			return true;
		}
	}
	return false;
}
 

@Test
public void testHasState() {
	StateObjectCollection<StateObject> stateObjects = new StateObjectCollection<>(new ArrayList<>());
	Assert.assertFalse(stateObjects.hasState());

	stateObjects = new StateObjectCollection<>(Collections.singletonList(null));
	Assert.assertFalse(stateObjects.hasState());

	stateObjects = new StateObjectCollection<>(Collections.singletonList(mock(StateObject.class)));
	Assert.assertTrue(stateObjects.hasState());
}
 

private static long sumAllSizes(Collection<? extends StateObject> stateObject) {
	long size = 0L;
	for (StateObject object : stateObject) {
		size += getSizeNullSafe(object);
	}

	return size;
}
 

/**
 * Creates a deep copy of the first state object in the given collection, or null if the collection is empy.
 */
private static  <T extends StateObject> StateObjectCollection<T> deepCopy(StateObjectCollection<T> original) {
	if (original == null || original.isEmpty()) {
		return StateObjectCollection.empty();
	}
	//noinspection unchecked
	return new StateObjectCollection<>((List<T>) original.stream().map(PrioritizedOperatorSubtaskStateTest::deepCopy).collect(Collectors.toList()));
}
 

private static long sumAllSizes(Collection<? extends StateObject> stateObject) {
	long size = 0L;
	for (StateObject object : stateObject) {
		size += getSizeNullSafe(object);
	}

	return size;
}
 

private <T extends StateObject> boolean checkResult(T expected, StateObjectCollection<T> actual) {
	if (expected == null) {
		return actual.isEmpty();
	}

	return actual.size() == 1 && expected == actual.iterator().next();
}
 

@Test
public void testHasState() {
	StateObjectCollection<StateObject> stateObjects = new StateObjectCollection<>(new ArrayList<>());
	Assert.assertFalse(stateObjects.hasState());

	stateObjects = new StateObjectCollection<>(Collections.singletonList(null));
	Assert.assertFalse(stateObjects.hasState());

	stateObjects = new StateObjectCollection<>(Collections.singletonList(mock(StateObject.class)));
	Assert.assertTrue(stateObjects.hasState());
}
 

private static long sumAllSizes(Collection<? extends StateObject> stateObject) {
	long size = 0L;
	for (StateObject object : stateObject) {
		size += getSizeNullSafe(object);
	}

	return size;
}
 

/**
 * Returns true if this contains at least one {@link StateObject}.
 */
public boolean hasState() {
	for (StateObject state : stateObjects) {
		if (state != null) {
			return true;
		}
	}
	return false;
}
 

private <T extends StateObject> boolean checkResultAsExpected(
	Function<OperatorSubtaskState, StateObjectCollection<T>> extractor,
	Function<PrioritizedOperatorSubtaskState, List<StateObjectCollection<T>>> extractor2,
	PrioritizedOperatorSubtaskState prioritizedResult,
	OperatorSubtaskState... expectedOrdered) {

	List<StateObjectCollection<T>> collector = new ArrayList<>(expectedOrdered.length);
	for (OperatorSubtaskState operatorSubtaskState : expectedOrdered) {
		collector.add(extractor.apply(operatorSubtaskState));
	}

	return checkRepresentSameOrder(
		extractor2.apply(prioritizedResult).iterator(),
		collector.toArray(new StateObjectCollection[collector.size()]));
}