Java StateInitializationContext类代码示例

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
  if (getKeyedStateBackend() != null) {
    KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();

    for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
      DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(streamProvider.getStream());

      int keyGroupIdx = streamProvider.getKeyGroupId();
      checkArgument(localKeyGroupRange.contains(keyGroupIdx),
          "Key Group " + keyGroupIdx + " does not belong to the local range.");

      // if (this instanceof KeyGroupRestoringOperator)
      restoreKeyGroupState(keyGroupIdx, div);

    }
  }
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	if (nfaOperatorState == null) {
		nfaOperatorState = getRuntimeContext().getState(
			new ValueStateDescriptor<>(
					NFA_OPERATOR_STATE_NAME,
					new NFA.NFASerializer<>(inputSerializer)));
	}

	if (elementQueueState == null) {
		elementQueueState = getRuntimeContext().getMapState(
				new MapStateDescriptor<>(
						EVENT_QUEUE_STATE_NAME,
						LongSerializer.INSTANCE,
						new ListSerializer<>(inputSerializer)
				)
		);
	}
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
/**
 * Stream operators with state which can be restored need to override this hook method.
 *
 * @param context context that allows to register different states.
 */
public void initializeState(StateInitializationContext context) throws Exception {
	if (getKeyedStateBackend() != null) {
		KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();

		// and then initialize the timer services
		for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
			int keyGroupIdx = streamProvider.getKeyGroupId();

			checkArgument(localKeyGroupRange.contains(keyGroupIdx),
				"Key Group " + keyGroupIdx + " does not belong to the local range.");

			timeServiceManager.restoreStateForKeyGroup(
				new DataInputViewStreamWrapper(streamProvider.getStream()),
				keyGroupIdx, getUserCodeClassloader());
		}
	}
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
public static void restoreFunctionState(
		StateInitializationContext context,
		Function userFunction) throws Exception {

	Preconditions.checkNotNull(context);

	while (true) {

		if (tryRestoreFunction(context, userFunction)) {
			break;
		}

		// inspect if the user function is wrapped, then unwrap and try again if we can restore the inner function
		if (userFunction instanceof WrappingFunction) {
			userFunction = ((WrappingFunction<?>) userFunction).getWrappedFunction();
		} else {
			break;
		}
	}
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	Preconditions.checkState(this.checkpointedState == null,
		"The reader state has already been initialized.");

	checkpointedState = context.getOperatorStateStore()
		.getSerializableListState("pending-checkpoints");

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

		for (PendingCheckpoint pendingCheckpoint : checkpointedState.get()) {
			this.pendingCheckpoints.add(pendingCheckpoint);
		}

		if (LOG.isDebugEnabled()) {
			LOG.debug("GenericWriteAheadSink idx {} restored {}.", subtaskIdx, this.pendingCheckpoints);
		}
	} else {
		LOG.info("No state to restore for the GenericWriteAheadSink (taskIdx={}).", subtaskIdx);
	}
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
/**
 * Tests that no checkpoints happen when the fetcher is not running.
 */
@Test
public void checkRestoredNullCheckpointWhenFetcherNotReady() throws Exception {
	FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	TestingListState<Serializable> listState = new TestingListState<>();
	when(operatorStateStore.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);

	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
	when(initializationContext.isRestored()).thenReturn(false);

	consumer.initializeState(initializationContext);

	consumer.snapshotState(new StateSnapshotContextSynchronousImpl(17, 17));

	assertFalse(listState.get().iterator().hasNext());
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
/**
 * Tests that on snapshots, states and offsets to commit to Kafka are correct
 */
@Test
public void checkUseFetcherWhenNoCheckpoint() throws Exception {

	FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);
	List<KafkaTopicPartition> partitionList = new ArrayList<>(1);
	partitionList.add(new KafkaTopicPartition("test", 0));
	consumer.setSubscribedPartitions(partitionList);

	OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
	TestingListState<Serializable> listState = new TestingListState<>();
	when(operatorStateStore.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);

	StateInitializationContext initializationContext = mock(StateInitializationContext.class);

	when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);

	// make the context signal that there is no restored state, then validate that
	when(initializationContext.isRestored()).thenReturn(false);
	consumer.initializeState(initializationContext);
	consumer.run(mock(SourceFunction.SourceContext.class));
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
    super.initializeState(context);

    if (this.engineState == null) {
        this.engineState = getRuntimeContext().getState(new ValueStateDescriptor<>(ESPER_SERVICE_PROVIDER_STATE, new EsperEngineSerializer()));
    }
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
    super.initializeState(context);

    // create a map-based queue to buffer input elements
    if (elementQueueState == null) {
        elementQueueState = getRuntimeContext().getMapState(
                new MapStateDescriptor<>(
                        EVENT_QUEUE_STATE_NAME,
                        LongSerializer.INSTANCE,
                        new ListSerializer<>(inputSerializer)
                )
        );
    }
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);
	if (siddhiRuntimeState == null) {
		siddhiRuntimeState = context.getOperatorStateStore().getUnionListState(new ListStateDescriptor<>(SIDDHI_RUNTIME_STATE_NAME,
				new BytePrimitiveArraySerializer()));
	}
	if (queuedRecordsState == null) {
		queuedRecordsState = context.getOperatorStateStore().getListState(
			new ListStateDescriptor<>(QUEUED_RECORDS_STATE_NAME, new BytePrimitiveArraySerializer()));
	}
	if (context.isRestored()) {
		restoreState();
	}
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
    super.initializeState(context);
    if (siddhiRuntimeState == null) {
        siddhiRuntimeState = context.getOperatorStateStore().getUnionListState(new ListStateDescriptor<>(SIDDHI_RUNTIME_STATE_NAME,
                new BytePrimitiveArraySerializer()));
    }
    if (queuedRecordsState == null) {
        queuedRecordsState = context.getOperatorStateStore().getListState(
            new ListStateDescriptor<>(QUEUED_RECORDS_STATE_NAME, new BytePrimitiveArraySerializer()));
    }
    if (context.isRestored()) {
        restoreState();
    }
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
  if (getKeyedStateBackend() != null) {
    int totalKeyGroups = getKeyedStateBackend().getNumberOfKeyGroups();
    KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();

    for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
      DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(streamProvider.getStream());

      int keyGroupIdx = streamProvider.getKeyGroupId();
      checkArgument(localKeyGroupRange.contains(keyGroupIdx),
          "Key Group " + keyGroupIdx + " does not belong to the local range.");

      // if (this instanceof KeyGroupRestoringOperator)
      restoreKeyGroupState(keyGroupIdx, div);

      // We just initialize our timerService
      if (keyCoder != null) {
        if (timerService == null) {
          timerService = new HeapInternalTimerService<>(
              totalKeyGroups,
              localKeyGroupRange,
              this,
              getRuntimeContext().getProcessingTimeService());
        }
        timerService.restoreTimersForKeyGroup(div, keyGroupIdx, getUserCodeClassloader());
      }
    }
  }
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);
	recoveredStreamElements = context
		.getOperatorStateStore()
		.getListState(new ListStateDescriptor<>(STATE_NAME, inStreamElementSerializer));

}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	checkState(checkpointedState == null,	"The reader state has already been initialized.");

	checkpointedState = context.getOperatorStateStore().getSerializableListState("splits");

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

		// this may not be null in case we migrate from a previous Flink version.
		if (restoredReaderState == null) {
			restoredReaderState = new ArrayList<>();
			for (TimestampedFileInputSplit split : checkpointedState.get()) {
				restoredReaderState.add(split);
			}

			if (LOG.isDebugEnabled()) {
				LOG.debug("{} (taskIdx={}) restored {}.", getClass().getSimpleName(), subtaskIdx, restoredReaderState);
			}
		}
	} else {
		LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIdx);
	}
}
 

import org.apache.flink.runtime.state.StateInitializationContext; //导入依赖的package包/类
private static boolean tryRestoreFunction(
		StateInitializationContext context,
		Function userFunction) throws Exception {

	if (userFunction instanceof CheckpointedFunction) {
		((CheckpointedFunction) userFunction).initializeState(context);

		return true;
	}

	if (context.isRestored() && userFunction instanceof ListCheckpointed) {
		@SuppressWarnings("unchecked")
		ListCheckpointed<Serializable> listCheckpointedFun = (ListCheckpointed<Serializable>) userFunction;

		ListState<Serializable> listState = context.getOperatorStateStore().
				getSerializableListState(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME);

		List<Serializable> list = new ArrayList<>();

		for (Serializable serializable : listState.get()) {
			list.add(serializable);
		}

		try {
			listCheckpointedFun.restoreState(list);
		} catch (Exception e) {

			throw new Exception("Failed to restore state to function: " + e.getMessage(), e);
		}

		return true;
	}

	return false;
}