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

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkState(this.checkpointedState == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	this.checkpointedState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>(
			"from-elements-state",
			IntSerializer.INSTANCE
		)
	);

	if (context.isRestored()) {
		List<Integer> retrievedStates = new ArrayList<>();
		for (Integer entry : this.checkpointedState.get()) {
			retrievedStates.add(entry);
		}

		// given that the parallelism of the function is 1, we can only have 1 state
		Preconditions.checkArgument(retrievedStates.size() == 1,
			getClass().getSimpleName() + " retrieved invalid state.");

		this.numElementsToSkip = retrievedStates.get(0);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	ListState<String> unionListState = context.getOperatorStateStore().getUnionListState(
			CheckpointingNonParallelSourceWithListState.STATE_DESCRIPTOR);

	if (context.isRestored()) {
		assertThat(unionListState.get(),
				containsInAnyOrder(CheckpointingParallelSourceWithUnionListState.CHECKPOINTED_STRINGS));

		getRuntimeContext().addAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, new IntCounter());
		getRuntimeContext().getAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR).add(1);
	} else {
		throw new RuntimeException(
				"This source should always be restored because it's only used when restoring from a savepoint.");
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	if (broadcast) {
		this.counterPartitions = context
				.getOperatorStateStore()
				.getUnionListState(new ListStateDescriptor<>("counter_partitions", IntSerializer.INSTANCE));
	} else {
		this.counterPartitions = context
				.getOperatorStateStore()
				.getListState(new ListStateDescriptor<>("counter_partitions", IntSerializer.INSTANCE));
	}

	if (context.isRestored()) {
		for (int v : counterPartitions.get()) {
			counter += v;
		}
		checkCorrectRestore[getRuntimeContext().getIndexOfThisSubtask()] = counter;
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	ListState<String> unionListState = context.getOperatorStateStore().getListState(
			CheckpointingNonParallelSourceWithListState.STATE_DESCRIPTOR);

	if (context.isRestored()) {
		assertThat(unionListState.get(),
				containsInAnyOrder(
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING,
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING_1,
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING_2,
						CheckpointingNonParallelSourceWithListState.CHECKPOINTED_STRING_3));

		getRuntimeContext().addAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, new IntCounter());
		getRuntimeContext().getAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR).add(1);
	} else {
		throw new RuntimeException(
				"This source should always be restored because it's only used when restoring from a savepoint.");
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

@Test
public void testOverrideConnection() throws Exception {
	final Connection mockConnection = Mockito.mock(Connection.class);
	Channel channel = Mockito.mock(Channel.class);
	Mockito.when(mockConnection.createChannel()).thenReturn(channel);

	RMQMockedRuntimeTestSource source = new RMQMockedRuntimeTestSource() {
		@Override
		protected Connection setupConnection() throws Exception {
			return mockConnection;
		}
	};

	FunctionInitializationContext mockContext = getMockContext();
	source.initializeState(mockContext);
	source.open(new Configuration());

	Mockito.verify(mockConnection, Mockito.times(1)).createChannel();
}
 

@Override
public final void initializeState(FunctionInitializationContext context) throws Exception {

	OperatorStateStore stateStore = context.getOperatorStateStore();

	this.unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor<>(OFFSETS_STATE_NAME,
		createStateSerializer(getRuntimeContext().getExecutionConfig())));

	if (context.isRestored()) {
		restoredState = new TreeMap<>(new KafkaTopicPartition.Comparator());

		// populate actual holder for restored state
		for (Tuple2<KafkaTopicPartition, Long> kafkaOffset : unionOffsetStates.get()) {
			restoredState.put(kafkaOffset.f0, kafkaOffset.f1);
		}

		LOG.info("Consumer subtask {} restored state: {}.", getRuntimeContext().getIndexOfThisSubtask(), restoredState);
	} else {
		LOG.info("Consumer subtask {} has no restore state.", getRuntimeContext().getIndexOfThisSubtask());
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {

	readyToFail = false;

	if (context.isRestored()) {
		isRunning = false;
	} else {
		isRunning = true;

		OperatorStateStore operatorStateStore = context.getOperatorStateStore();
		for (int i = 0; i < numListStates; ++i) {

			ListStateDescriptor<String> listStateDescriptor =
				new ListStateDescriptor<>("test-list-state-" + i, String.class);

			ListState<String> unionListState =
				operatorStateStore.getUnionListState(listStateDescriptor);

			for (int j = 0; j < numPartitionsPerListState; ++j) {
				unionListState.add(String.valueOf(j));
			}
		}
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) {
	states = TtlStateVerifier.VERIFIERS.stream()
		.collect(Collectors.toMap(TtlStateVerifier::getId, v -> v.createState(context, ttlConfig)));
	prevUpdatesByVerifierId = TtlStateVerifier.VERIFIERS.stream()
		.collect(Collectors.toMap(TtlStateVerifier::getId, v -> {
			checkNotNull(v);
			final TypeSerializer<ValueWithTs<?>> typeSerializer = new ValueWithTs.Serializer(
				v.getUpdateSerializer(),
				LongSerializer.INSTANCE);

			ListStateDescriptor<ValueWithTs<?>> stateDesc = new ListStateDescriptor<>(
				"TtlPrevValueState_" + v.getId(), typeSerializer);
			KeyedStateStore store = context.getKeyedStateStore();
			return store.getListState(stateDesc);
		}));
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkState(this.checkpointedState == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	this.checkpointedState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>(
			"from-elements-state",
			IntSerializer.INSTANCE
		)
	);

	if (context.isRestored()) {
		List<Integer> retrievedStates = new ArrayList<>();
		for (Integer entry : this.checkpointedState.get()) {
			retrievedStates.add(entry);
		}

		// given that the parallelism of the function is 1, we can only have 1 state
		Preconditions.checkArgument(retrievedStates.size() == 1,
			getClass().getSimpleName() + " retrieved invalid state.");

		this.numElementsToSkip = retrievedStates.get(0);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	ListState<String> unionListState = context.getOperatorStateStore().getUnionListState(
			CheckpointingNonParallelSourceWithListState.STATE_DESCRIPTOR);

	if (context.isRestored()) {
		assertThat(unionListState.get(),
				containsInAnyOrder(CheckpointingParallelSourceWithUnionListState.CHECKPOINTED_STRINGS));

		getRuntimeContext().addAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR, new IntCounter());
		getRuntimeContext().getAccumulator(SUCCESSFUL_RESTORE_CHECK_ACCUMULATOR).add(1);
	} else {
		throw new RuntimeException(
				"This source should always be restored because it's only used when restoring from a savepoint.");
	}
}
 

@Before
public void beforeTest() throws Exception {

	OperatorStateStore mockStore = Mockito.mock(OperatorStateStore.class);
	FunctionInitializationContext mockContext = Mockito.mock(FunctionInitializationContext.class);
	Mockito.when(mockContext.getOperatorStateStore()).thenReturn(mockStore);
	Mockito.when(mockStore.getSerializableListState(any(String.class))).thenReturn(null);

	source = new RMQTestSource();
	source.initializeState(mockContext);
	source.open(config);

	messageId = 0;
	generateCorrelationIds = true;

	sourceThread = new Thread(new Runnable() {
		@Override
		public void run() {
			try {
				source.run(new DummySourceContext());
			} catch (Exception e) {
				exception = e;
			}
		}
	});
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
    OperatorStateStore stateStore = context.getOperatorStateStore();

    unionOffsetStates = stateStore.getUnionListState(
            new ListStateDescriptor<>(
                    OFFSETS_STATE_NAME,
                    TypeInformation.of(new TypeHint<Tuple2<String, MessageId>>() {
                    })));

    if (context.isRestored()) {
        restoredState = new TreeMap<>();
        unionOffsetStates.get().forEach(e -> restoredState.put(e.f0, e.f1));
        log.info("Source subtask {} restored state {}",
                taskIndex,
                StringUtils.join(restoredState.entrySet()));
    } else {
        log.info("Source subtask {} has no restore state", taskIndex);
    }
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkState(this.checkpointedState == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	this.checkpointedState = context
		.getOperatorStateStore()
		.getSerializableListState("message-acknowledging-source-state");

	this.idsForCurrentCheckpoint = new HashSet<>(64);
	this.pendingCheckpoints = new ArrayDeque<>();
	this.idsProcessedButNotAcknowledged = new HashSet<>();

	if (context.isRestored()) {
		LOG.info("Restoring state for the {}.", getClass().getSimpleName());

		List<SerializedCheckpointData[]> retrievedStates = new ArrayList<>();
		for (SerializedCheckpointData[] entry : this.checkpointedState.get()) {
			retrievedStates.add(entry);
		}

		// given that the parallelism of the function is 1, we can only have at most 1 state
		Preconditions.checkArgument(retrievedStates.size() == 1,
			getClass().getSimpleName() + " retrieved invalid state.");

		pendingCheckpoints = SerializedCheckpointData.toDeque(retrievedStates.get(0), idSerializer);
		// build a set which contains all processed ids. It may be used to check if we have
		// already processed an incoming message.
		for (Tuple2<Long, Set<UId>> checkpoint : pendingCheckpoints) {
			idsProcessedButNotAcknowledged.addAll(checkpoint.f1);
		}
	} else {
		LOG.info("No state to restore for the {}.", getClass().getSimpleName());
	}
}
 

@SuppressWarnings("unchecked")
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	pojoClass = getRuntimeContext().getUserCodeClassLoader().loadClass(POJO_NAME);

	fieldA = pojoClass.getDeclaredField("a");
	fieldA.setAccessible(true);

	if (hasBField) {
		fieldB = pojoClass.getDeclaredField("b");
		fieldB.setAccessible(true);
	}

	if (keyed) {
		keyedValueState = context.getKeyedStateStore().getState(
			new ValueStateDescriptor<>("keyedValueState", (Class<Object>) pojoClass));
		keyedListState = context.getKeyedStateStore().getListState(
			new ListStateDescriptor<>("keyedListState", (Class<Object>) pojoClass));

		ReduceFunction<Object> reduceFunction = new FirstValueReducer<>();
		keyedReducingState = context.getKeyedStateStore().getReducingState(
			new ReducingStateDescriptor<>("keyedReducingState", reduceFunction, (Class<Object>) pojoClass));
	} else {
		partitionableListState = context.getOperatorStateStore().getListState(
			new ListStateDescriptor<>("partitionableListState", (Class<Object>) pojoClass));
		unionListState = context.getOperatorStateStore().getUnionListState(
			new ListStateDescriptor<>("unionListState", (Class<Object>) pojoClass));
	}
}
 

@SuppressWarnings("unchecked")
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	pojoClass = getRuntimeContext().getUserCodeClassLoader().loadClass(POJO_NAME);

	fieldA = pojoClass.getDeclaredField("a");
	fieldA.setAccessible(true);

	if (hasBField) {
		fieldB = pojoClass.getDeclaredField("b");
		fieldB.setAccessible(true);
	}

	if (keyed) {
		keyedValueState = context.getKeyedStateStore().getState(
			new ValueStateDescriptor<>("keyedValueState", (Class<Object>) pojoClass));
		keyedListState = context.getKeyedStateStore().getListState(
			new ListStateDescriptor<>("keyedListState", (Class<Object>) pojoClass));

		ReduceFunction<Object> reduceFunction = new FirstValueReducer<>();
		keyedReducingState = context.getKeyedStateStore().getReducingState(
			new ReducingStateDescriptor<>("keyedReducingState", reduceFunction, (Class<Object>) pojoClass));
	} else {
		partitionableListState = context.getOperatorStateStore().getListState(
			new ListStateDescriptor<>("partitionableListState", (Class<Object>) pojoClass));
		unionListState = context.getOperatorStateStore().getUnionListState(
			new ListStateDescriptor<>("unionListState", (Class<Object>) pojoClass));
	}
}
 

private void migrateNextTransactionalIdHindState(FunctionInitializationContext context) throws Exception {
	ListState<NextTransactionalIdHint> oldNextTransactionalIdHintState = context.getOperatorStateStore().getUnionListState(
		NEXT_TRANSACTIONAL_ID_HINT_DESCRIPTOR);
	nextTransactionalIdHintState = context.getOperatorStateStore().getUnionListState(NEXT_TRANSACTIONAL_ID_HINT_DESCRIPTOR_V2);

	ArrayList<NextTransactionalIdHint> oldTransactionalIdHints = Lists.newArrayList(oldNextTransactionalIdHintState.get());
	if (!oldTransactionalIdHints.isEmpty()) {
		nextTransactionalIdHintState.addAll(oldTransactionalIdHints);
		//clear old state
		oldNextTransactionalIdHintState.clear();
	}
}
 

@Override
public void open(
		String name,
		FunctionInitializationContext context,
		RuntimeContext runtimeContext) throws Exception {
	this.random = new RandomDataGenerator();
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkState(this.checkpointedState == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	this.checkpointedState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>(
			"arrow-source-state",
			new TupleSerializer<>(
				(Class<Tuple2<Integer, Integer>>) (Class<?>) Tuple2.class,
				new TypeSerializer[]{IntSerializer.INSTANCE, IntSerializer.INSTANCE})
		)
	);

	this.indexesToEmit = new ArrayDeque<>();
	if (context.isRestored()) {
		// upon restoring
		for (Tuple2<Integer, Integer> v : this.checkpointedState.get()) {
			this.indexesToEmit.add(v);
		}
		LOG.info("Subtask {} restored state: {}.", getRuntimeContext().getIndexOfThisSubtask(), indexesToEmit);
	} else {
		// the first time the job is executed
		final int stepSize = getRuntimeContext().getNumberOfParallelSubtasks();
		final int taskIdx = getRuntimeContext().getIndexOfThisSubtask();

		for (int i = taskIdx; i < arrowData.length; i += stepSize) {
			this.indexesToEmit.add(Tuple2.of(i, 0));
		}
		LOG.info("Subtask {} has no restore state, initialized with {}.", taskIdx, indexesToEmit);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	state = context.getOperatorStateStore().getListState(
			new ListStateDescriptor<Integer>("state", IntSerializer.INSTANCE));

	for (Integer i : state.get()) {
		numRecordsEmitted += i;
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	state = context.getOperatorStateStore().getListState(
			new ListStateDescriptor<Long>("state", LongSerializer.INSTANCE));

	for (Long l : state.get()) {
		ms += l;
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");

	try {
		initFileSystem();
	} catch (IOException e) {
		LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
		throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
	}

	if (this.refTruncate == null) {
		this.refTruncate = reflectTruncate(fs);
	}

	// We are using JavaSerializer from the flink-runtime module here. This is very naughty and
	// we shouldn't be doing it because ideally nothing in the API modules/connector depends
	// directly on flink-runtime. We are doing it here because we need to maintain backwards
	// compatibility with old state and because we will have to rework/remove this code soon.
	OperatorStateStore stateStore = context.getOperatorStateStore();
	this.restoredBucketStates = stateStore.getListState(new ListStateDescriptor<>("bucket-states", new JavaSerializer<>()));

	int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);

		for (State<T> recoveredState : restoredBucketStates.get()) {
			handleRestoredBucketState(recoveredState);
			if (LOG.isDebugEnabled()) {
				LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
	}
}
 

private static <T> FunctionInitializationContext getInitializationContext(ListState<T> listState)
    throws Exception {
  FunctionInitializationContext mock = Mockito.mock(FunctionInitializationContext.class);
  OperatorStateStore mockOperatorState = getMockOperatorState(listState);
  when(mock.getOperatorStateStore()).thenReturn(mockOperatorState);
  return mock;
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	final int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
	this.buckets = bucketsBuilder.createBuckets(subtaskIndex);

	final OperatorStateStore stateStore = context.getOperatorStateStore();
	bucketStates = stateStore.getListState(BUCKET_STATE_DESC);
	maxPartCountersState = stateStore.getUnionListState(MAX_PART_COUNTER_STATE_DESC);

	if (context.isRestored()) {
		buckets.initializeState(bucketStates, maxPartCountersState);
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	serializableListState = context
			.getOperatorStateStore()
			.getListState(new ListStateDescriptor<>("test-state", IntSerializer.INSTANCE));
	if (context.isRestored()) {
		Iterator<Integer> integers = serializableListState.get().iterator();
		int act = integers.next();
		Assert.assertEquals(42, act);
		Assert.assertFalse(integers.hasNext());
		wasRestored = true;
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	Preconditions.checkState(this.checkpointedState == null,
		"The " + getClass().getSimpleName() + " has already been initialized.");

	this.checkpointedState = context
		.getOperatorStateStore()
		.getSerializableListState("message-acknowledging-source-state");

	this.idsForCurrentCheckpoint = new HashSet<>(64);
	this.pendingCheckpoints = new ArrayDeque<>();
	this.idsProcessedButNotAcknowledged = new HashSet<>();

	if (context.isRestored()) {
		LOG.info("Restoring state for the {}.", getClass().getSimpleName());

		List<SerializedCheckpointData[]> retrievedStates = new ArrayList<>();
		for (SerializedCheckpointData[] entry : this.checkpointedState.get()) {
			retrievedStates.add(entry);
		}

		// given that the parallelism of the function is 1, we can only have at most 1 state
		Preconditions.checkArgument(retrievedStates.size() == 1,
			getClass().getSimpleName() + " retrieved invalid state.");

		pendingCheckpoints = SerializedCheckpointData.toDeque(retrievedStates.get(0), idSerializer);
		// build a set which contains all processed ids. It may be used to check if we have
		// already processed an incoming message.
		for (Tuple2<Long, Set<UId>> checkpoint : pendingCheckpoints) {
			idsProcessedButNotAcknowledged.addAll(checkpoint.f1);
		}
	} else {
		LOG.info("No state to restore for the {}.", getClass().getSimpleName());
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	sumState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>("sumState", Integer.class));

	if (context.isRestored()) {
		for (int value : sumState.get()) {
			runningSum += value;
		}
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
	if (context.isRestored()) {
		restoredState = true;

		unionListState = context.getOperatorStateStore().getUnionListState(unionStateDescriptor);
		Set<Integer> actualIndices = StreamSupport.stream(unionListState.get().spliterator(), false).collect(Collectors.toSet());
		if (getRuntimeContext().getTaskName().contains(SINGLE_REGION_SOURCE_NAME)) {
			Assert.assertTrue(CollectionUtils.isEqualCollection(EXPECTED_INDICES_SINGLE_REGION, actualIndices));
		} else {
			Assert.assertTrue(CollectionUtils.isEqualCollection(EXPECTED_INDICES_MULTI_REGION, actualIndices));
		}

		if (indexOfThisSubtask == 0) {
			listState = context.getOperatorStateStore().getListState(stateDescriptor);
			Assert.assertTrue("list state should be empty for subtask-0",
				((List<Integer>) listState.get()).isEmpty());
		} else {
			listState = context.getOperatorStateStore().getListState(stateDescriptor);
			Assert.assertTrue("list state should not be empty for subtask-" + indexOfThisSubtask,
				((List<Integer>) listState.get()).size() > 0);

			if (indexOfThisSubtask == NUM_OF_REGIONS - 1) {
				index = listState.get().iterator().next();
				if (index != snapshotIndicesOfSubTask.get(lastCompletedCheckpointId.get())) {
					throw new RuntimeException("Test failed due to unexpected recovered index: " + index +
						", while last completed checkpoint record index: " + snapshotIndicesOfSubTask.get(lastCompletedCheckpointId.get()));
				}
			}
		}
	} else {
		unionListState = context.getOperatorStateStore().getUnionListState(unionStateDescriptor);

		if (indexOfThisSubtask != 0) {
			listState = context.getOperatorStateStore().getListState(stateDescriptor);
		}
	}

}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	state = context.getOperatorStateStore().getListState(
			new ListStateDescriptor<Integer>("state", IntSerializer.INSTANCE));

	for (Integer i : state.get()) {
		saveRecordCnt += i;
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	sumState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>("sumState", Integer.class));

	if (context.isRestored()) {
		for (int value : sumState.get()) {
			runningSum += value;
		}
	}
}