Java Code Examples for org.apache.flink.runtime.state.FunctionInitializationContext#isRestored()

@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 {
    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 {

	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) 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 {
	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 {
	state = context.getOperatorStateStore().getUnionListState(
		new ListStateDescriptor<>("numbers", Types.INT)
	);

	if (context.isRestored()) {
		Set<Integer> expected = new HashSet<>();
		expected.add(1);
		expected.add(2);
		expected.add(3);

		for (Integer number : state.get()) {
			Assert.assertTrue("Duplicate state", expected.contains(number));
			expected.remove(number);
		}

		Assert.assertTrue("Failed to bootstrap all state elements: " + Arrays.toString(expected.toArray()), expected.isEmpty());
	}
}
 

@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) {
	if (!context.isRestored()) {
		int updatedValue = allowedInitializeCallsWithoutRestore.decrementAndGet();
		if (updatedValue < 0) {
			illegalRestores.getAndIncrement();
			throw new RuntimeException("We are not allowed any more restores.");
		}
	} else {
		if (!afterMessWithZooKeeper.get()) {
			illegalRestores.getAndIncrement();
		} else if (successfulRestores.getAndIncrement() > 0) {
			// already saw the one allowed successful restore
			illegalRestores.getAndIncrement();
		}
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	this.helper = new StreamingFileSinkHelper<>(
			bucketsBuilder.createBuckets(getRuntimeContext().getIndexOfThisSubtask()),
			context.isRestored(),
			context.getOperatorStateStore(),
			((StreamingRuntimeContext) getRuntimeContext()).getProcessingTimeService(),
			bucketCheckInterval);
}
 

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

	// 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.
	this.checkpointedState = context
		.getOperatorStateStore()
		.getListState(new ListStateDescriptor<>("message-acknowledging-source-state", new JavaSerializer<>()));

	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 {
	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 {
	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 {
	canFail = false;
	position = 0L;
	isRestored = context.isRestored();
	positionState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>("posState", Long.class));

	if (isRestored) {
		for (long value : positionState.get()) {
			position += value;
		}
	}
}
 

@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 {
	canFail = false;
	position = 0L;
	isRestored = context.isRestored();
	positionState = context.getOperatorStateStore().getListState(
		new ListStateDescriptor<>("posState", Long.class));

	if (isRestored) {
		for (long value : positionState.get()) {
			position += value;
		}
	}
}
 

@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 {
	TypeInformation<Tuple2<StreamShardMetadata, SequenceNumber>> shardsStateTypeInfo = new TupleTypeInfo<>(
		TypeInformation.of(StreamShardMetadata.class),
		TypeInformation.of(SequenceNumber.class));

	sequenceNumsStateForCheckpoint = context.getOperatorStateStore().getUnionListState(
		new ListStateDescriptor<>(sequenceNumsStateStoreName, shardsStateTypeInfo));

	if (context.isRestored()) {
		if (sequenceNumsToRestore == null) {
			sequenceNumsToRestore = new HashMap<>();
			for (Tuple2<StreamShardMetadata, SequenceNumber> kinesisSequenceNumber : sequenceNumsStateForCheckpoint.get()) {
				sequenceNumsToRestore.put(
					// we wrap the restored metadata inside an equivalence wrapper that checks only stream name and shard id,
					// so that if a shard had been closed (due to a Kinesis reshard operation, for example) since
					// the savepoint and has a different metadata than what we last stored,
					// we will still be able to match it in sequenceNumsToRestore. Please see FLINK-8484 for details.
					new StreamShardMetadata.EquivalenceWrapper(kinesisSequenceNumber.f0),
					kinesisSequenceNumber.f1);
			}

			LOG.info("Setting restore state in the FlinkKinesisConsumer. Using the following offsets: {}",
				sequenceNumsToRestore);
		}
	} else {
		LOG.info("No restore state for FlinkKinesisConsumer.");
	}
}
 

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	// when we are restoring state with pendingCommitTransactions, we don't really know whether the
	// transactions were already committed, or whether there was a failure between
	// completing the checkpoint on the master, and notifying the writer here.

	// (the common case is actually that is was already committed, the window
	// between the commit on the master and the notification here is very small)

	// it is possible to not have any transactions at all if there was a failure before
	// the first completed checkpoint, or in case of a scale-out event, where some of the
	// new task do not have and transactions assigned to check)

	// we can have more than one transaction to check in case of a scale-in event, or
	// for the reasons discussed in the 'notifyCheckpointComplete()' method.

	state = context.getOperatorStateStore().getListState(stateDescriptor);

	boolean recoveredUserContext = false;
	if (context.isRestored()) {
		LOG.info("{} - restoring state", name());
		for (State<TXN, CONTEXT> operatorState : state.get()) {
			userContext = operatorState.getContext();
			List<TransactionHolder<TXN>> recoveredTransactions = operatorState.getPendingCommitTransactions();
			List<TXN> handledTransactions = new ArrayList<>(recoveredTransactions.size() + 1);
			for (TransactionHolder<TXN> recoveredTransaction : recoveredTransactions) {
				// If this fails to succeed eventually, there is actually data loss
				recoverAndCommitInternal(recoveredTransaction);
				handledTransactions.add(recoveredTransaction.handle);
				LOG.info("{} committed recovered transaction {}", name(), recoveredTransaction);
			}

			{
				TXN transaction = operatorState.getPendingTransaction().handle;
				recoverAndAbort(transaction);
				handledTransactions.add(transaction);
				LOG.info("{} aborted recovered transaction {}", name(), operatorState.getPendingTransaction());
			}

			if (userContext.isPresent()) {
				finishRecoveringContext(handledTransactions);
				recoveredUserContext = true;
			}
		}
	}

	// if in restore we didn't get any userContext or we are initializing from scratch
	if (!recoveredUserContext) {
		LOG.info("{} - no state to restore", name());

		userContext = initializeUserContext();
	}
	this.pendingCommitTransactions.clear();

	currentTransactionHolder = beginTransactionInternal();
	LOG.debug("{} - started new transaction '{}'", name(), currentTransactionHolder);
}
 

/**
 * The actual paths in this depend on the binary checkpoint so it you update this the paths
 * here have to be updated as well.
 */
@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
	OperatorStateStore stateStore = context.getOperatorStateStore();

	// 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.
	ListState<State<T>> restoredBucketStates = stateStore.getListState(new ListStateDescriptor<>("bucket-states", new JavaSerializer<>()));

	if (context.isRestored()) {

		for (State<T> states : restoredBucketStates.get()) {
			for (String bucketPath : states.bucketStates.keySet()) {
				BucketState state = states.getBucketState(new Path(bucketPath));
				String current = state.currentFile;
				long validLength = state.currentFileValidLength;

				Assert.assertEquals(expectedBucketFilesPrefix + "4", current);
				Assert.assertEquals(6, validLength);

				List<String> pendingFiles = state.pendingFiles;
				assertTrue(pendingFiles.isEmpty());

				final Map<Long, List<String>> pendingFilesPerCheckpoint = state.pendingFilesPerCheckpoint;
				Assert.assertEquals(1, pendingFilesPerCheckpoint.size());

				for (Map.Entry<Long, List<String>> entry: pendingFilesPerCheckpoint.entrySet()) {
					long checkpoint = entry.getKey();
					List<String> files = entry.getValue();

					Assert.assertEquals(0L, checkpoint);
					Assert.assertEquals(4, files.size());

					for (int i = 0; i < 4; i++) {
						Assert.assertEquals(
								expectedBucketFilesPrefix + i,
								files.get(i));
					}
				}
			}
		}
	}

	initializeCalled = true;
	super.initializeState(context);
}
 

@Override
public void initializeState(FunctionInitializationContext functionInitializationContext) throws Exception {
	ValueStateDescriptor<String> stateDescriptor =
		new ValueStateDescriptor<>("state", String.class);
	valueState = functionInitializationContext.getKeyedStateStore().getState(stateDescriptor);

	ListStateDescriptor<MapperSchedulingAndFailureInfo> mapperInfoStateDescriptor =
		new ListStateDescriptor<>("mapperState", MapperSchedulingAndFailureInfo.class);
	schedulingAndFailureState =
		functionInitializationContext.getOperatorStateStore().getUnionListState(mapperInfoStateDescriptor);

	StreamingRuntimeContext runtimeContext = (StreamingRuntimeContext) getRuntimeContext();
	String allocationID = runtimeContext.getAllocationIDAsString();

	final int thisJvmPid = getJvmPid();
	final Set<Integer> killedJvmPids = new HashSet<>();

	// here we check if the sticky scheduling worked as expected
	if (functionInitializationContext.isRestored()) {
		Iterable<MapperSchedulingAndFailureInfo> iterable = schedulingAndFailureState.get();
		String taskNameWithSubtasks = runtimeContext.getTaskNameWithSubtasks();

		MapperSchedulingAndFailureInfo infoForThisTask = null;
		List<MapperSchedulingAndFailureInfo> completeInfo = new ArrayList<>();
		if (iterable != null) {
			for (MapperSchedulingAndFailureInfo testInfo : iterable) {

				completeInfo.add(testInfo);

				if (taskNameWithSubtasks.equals(testInfo.taskNameWithSubtask)) {
					infoForThisTask = testInfo;
				}

				if (testInfo.killedJvm) {
					killedJvmPids.add(testInfo.jvmPid);
				}
			}
		}

		Preconditions.checkNotNull(infoForThisTask, "Expected to find info here.");

		if (!isScheduledToCorrectAllocation(infoForThisTask, allocationID, killedJvmPids)) {
			allocationFailureMessage = String.format(
				"Sticky allocation test failed: Subtask %s in attempt %d was rescheduled from allocation %s " +
					"on JVM with PID %d to unexpected allocation %s on JVM with PID %d.\n" +
					"Complete information from before the crash: %s.",
				runtimeContext.getTaskNameWithSubtasks(),
				runtimeContext.getAttemptNumber(),
				infoForThisTask.allocationId,
				infoForThisTask.jvmPid,
				allocationID,
				thisJvmPid,
				completeInfo);
		}
	}

	// We determine which of the subtasks will produce the artificial failure
	boolean failingTask = shouldTaskFailForThisAttempt();

	// We take note of all the meta info that we require to check sticky scheduling in the next re-attempt
	this.currentSchedulingAndFailureInfo = new MapperSchedulingAndFailureInfo(
		failingTask,
		failingTask && killTaskOnFailure,
		thisJvmPid,
		runtimeContext.getTaskNameWithSubtasks(),
		allocationID);

	schedulingAndFailureState.clear();
	schedulingAndFailureState.add(currentSchedulingAndFailureInfo);
}