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

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	counterState = context
		.getKeyedStateStore()
		.getState(new ValueStateDescriptor<>(prefix + "counter-state", LongSerializer.INSTANCE));

	// set key manually to make RocksDBListState get the serialized key.
	setCurrentKey("10");

	if (context.isRestored()) {
		counter =  counterState.value();
		assertEquals(snapshotOutData, counter);
		counterState.clear();
	}
}
 

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

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

public void initializeOperatorState(CheckpointedStreamOperator streamOperator) throws Exception {
	CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = context.rawKeyedStateInputs();
	CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = context.rawOperatorStateInputs();

	try {
		StateInitializationContext initializationContext = new StateInitializationContextImpl(
			context.isRestored(), // information whether we restore or start for the first time
			operatorStateBackend, // access to operator state backend
			keyedStateStore, // access to keyed state backend
			keyedStateInputs, // access to keyed state stream
			operatorStateInputs); // access to operator state stream

		streamOperator.initializeState(initializationContext);
	} finally {
		closeFromRegistry(operatorStateInputs, closeableRegistry);
		closeFromRegistry(keyedStateInputs, closeableRegistry);
	}
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);
	currentWatermark = Long.MIN_VALUE;
	this.trigger = PartitionCommitTrigger.create(
			context.isRestored(),
			context.getOperatorStateStore(),
			conf,
			getUserCodeClassloader(),
			partitionKeys,
			getProcessingTimeService());
	this.policies = PartitionCommitPolicy.createPolicyChain(
			getUserCodeClassloader(),
			conf.get(SINK_PARTITION_COMMIT_POLICY_KIND),
			conf.get(SINK_PARTITION_COMMIT_POLICY_CLASS),
			conf.get(SINK_PARTITION_COMMIT_SUCCESS_FILE_NAME),
			() -> {
				try {
					return fsFactory.create(locationPath.toUri());
				} catch (IOException e) {
					throw new RuntimeException(e);
				}
			});
}
 

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

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

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	this.leftBuffer = context.getKeyedStateStore().getMapState(new MapStateDescriptor<>(
		LEFT_BUFFER,
		LongSerializer.INSTANCE,
		new ListSerializer<>(new BufferEntrySerializer<>(leftTypeSerializer))
	));

	this.rightBuffer = context.getKeyedStateStore().getMapState(new MapStateDescriptor<>(
		RIGHT_BUFFER,
		LongSerializer.INSTANCE,
		new ListSerializer<>(new BufferEntrySerializer<>(rightTypeSerializer))
	));
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	// initializeState through the provided context
	computationStates = context.getKeyedStateStore().getState(
		new ValueStateDescriptor<>(
			NFA_STATE_NAME,
			new NFAStateSerializer()));

	partialMatches = new SharedBuffer<>(context.getKeyedStateStore(), inputSerializer);

	elementQueueState = context.getKeyedStateStore().getMapState(
			new MapStateDescriptor<>(
					EVENT_QUEUE_STATE_NAME,
					LongSerializer.INSTANCE,
					new ListSerializer<>(inputSerializer)));

	migrateOldState();
}
 

private static <T> SourceOperator<T, MockSourceSplit> createTestOperator(
		SourceReader<T, MockSourceSplit> reader,
		WatermarkStrategy<T> watermarkStrategy,
		ProcessingTimeService timeService) throws Exception {

	final OperatorStateStore operatorStateStore =
			new MemoryStateBackend().createOperatorStateBackend(
					new MockEnvironmentBuilder().build(),
					"test-operator",
					Collections.emptyList(),
					new CloseableRegistry());

	final StateInitializationContext stateContext = new StateInitializationContextImpl(
		false, operatorStateStore, null, null, null);

	final SourceOperator<T, MockSourceSplit> sourceOperator =
			new TestingSourceOperator<>(reader, watermarkStrategy, timeService);
	sourceOperator.initializeState(stateContext);
	sourceOperator.open();

	return sourceOperator;
}
 

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

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	counterState = context
		.getKeyedStateStore()
		.getState(new ValueStateDescriptor<>(prefix + "counter-state", LongSerializer.INSTANCE));

	// set key manually to make RocksDBListState get the serialized key.
	setCurrentKey("10");

	if (context.isRestored()) {
		counter =  counterState.value();
		assertEquals(snapshotOutData, counter);
		counterState.clear();
	}
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	// initializeState through the provided context
	computationStates = context.getKeyedStateStore().getState(
		new ValueStateDescriptor<>(
			NFA_STATE_NAME,
			new NFAStateSerializer()));

	partialMatches = new SharedBuffer<>(context.getKeyedStateStore(), inputSerializer);

	elementQueueState = context.getKeyedStateStore().getMapState(
			new MapStateDescriptor<>(
					EVENT_QUEUE_STATE_NAME,
					LongSerializer.INSTANCE,
					new ListSerializer<>(inputSerializer)));

	migrateOldState();
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	counterState = context
		.getKeyedStateStore()
		.getState(new ValueStateDescriptor<>(prefix + "counter-state", LongSerializer.INSTANCE));

	// set key manually to make RocksDBListState get the serialized key.
	setCurrentKey("10");

	if (context.isRestored()) {
		counter =  counterState.value();
		assertEquals(snapshotOutData, counter);
		counterState.clear();
	}
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	// initializeState through the provided context
	computationStates = context.getKeyedStateStore().getState(
		new ValueStateDescriptor<>(
			NFA_STATE_NAME,
			new NFAStateSerializer()));

	partialMatches = new SharedBuffer<>(context.getKeyedStateStore(), inputSerializer);

	elementQueueState = context.getKeyedStateStore().getMapState(
			new MapStateDescriptor<>(
					EVENT_QUEUE_STATE_NAME,
					LongSerializer.INSTANCE,
					new ListSerializer<>(inputSerializer)));

	migrateOldState();
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	this.leftBuffer = context.getKeyedStateStore().getMapState(new MapStateDescriptor<>(
		LEFT_BUFFER,
		LongSerializer.INSTANCE,
		new ListSerializer<>(new BufferEntrySerializer<>(leftTypeSerializer))
	));

	this.rightBuffer = context.getKeyedStateStore().getMapState(new MapStateDescriptor<>(
		RIGHT_BUFFER,
		LongSerializer.INSTANCE,
		new ListSerializer<>(new BufferEntrySerializer<>(rightTypeSerializer))
	));
}
 

private StateInitializationContext getStateContext() throws Exception {
	// Create a mock split.
	byte[] serializedSplitWithVersion = SimpleVersionedSerialization
			.writeVersionAndSerialize(new MockSourceSplitSerializer(), MOCK_SPLIT);

	// Crate the state context.
	OperatorStateStore operatorStateStore = createOperatorStateStore();
	StateInitializationContext stateContext = new StateInitializationContextImpl(
			false,
			operatorStateStore,
			null,
			null,
			null);

	// Update the context.
	stateContext.getOperatorStateStore()
				.getListState(SourceOperator.SPLITS_STATE_DESC)
				.update(Collections.singletonList(serializedSplitWithVersion));

	return stateContext;
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);

	counterState = context
		.getOperatorStateStore()
		.getListState(new ListStateDescriptor<>("counter-state", LongSerializer.INSTANCE));

	if (context.isRestored()) {
		for (Long value : counterState.get()) {
			counter += value;
		}
		counterState.clear();
	}
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
  super.initializeState(context);

  timerService =
      getInternalTimerService("dedup-cleanup-timer", VoidNamespaceSerializer.INSTANCE, this);
}
 

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

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);
	TupleTypeInfo<Tuple2<RowData, Long>> tupleType = new TupleTypeInfo<>(inputRowType, Types.LONG);
	this.bufferState = context.getOperatorStateStore()
			.getListState(new ListStateDescriptor<>("localBufferState", tupleType));
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	keyedStateBackend = (AbstractKeyedStateBackend<?>) getKeyedStateBackend();
	operatorStateBackend = getOperatorStateBackend();
	rawOperatorStateInputs =
		(CloseableIterable<StatePartitionStreamProvider>) context.getRawOperatorStateInputs();
	rawKeyedStateInputs =
		(CloseableIterable<KeyGroupStatePartitionStreamProvider>) context.getRawKeyedStateInputs();
	super.initializeState(context);
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
  super.initializeState(context);

  final IOManager ioManager = getContainingTask().getEnvironment().getIOManager();
  final int maxParallelism = getRuntimeContext().getMaxNumberOfParallelSubtasks();

  this.reusable = new StreamRecord<>(null);

  //
  // Initialize the unbounded feedback logger
  //
  @SuppressWarnings("unchecked")
  UnboundedFeedbackLoggerFactory<T> feedbackLoggerFactory =
      (UnboundedFeedbackLoggerFactory<T>)
          Loggers.unboundedSpillableLoggerFactory(
              ioManager,
              maxParallelism,
              totalMemoryUsedForFeedbackCheckpointing,
              elementSerializer,
              keySelector);

  this.checkpoints = new Checkpoints<>(feedbackLoggerFactory::create);

  //
  // we first must reply previously check-pointed envelopes before we start
  // processing any new envelopes.
  //
  UnboundedFeedbackLogger<T> logger = feedbackLoggerFactory.create();
  for (KeyGroupStatePartitionStreamProvider keyedStateInput : context.getRawKeyedStateInputs()) {
    logger.replyLoggedEnvelops(keyedStateInput.getStream(), this);
  }
  //
  // now we can start processing new messages. We do so by registering ourselves as a
  // FeedbackConsumer
  //
  registerFeedbackConsumer(new MailboxExecutorFacade(mailboxExecutor, "Feedback Consumer"));
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
  super.initializeState(context);

  final State stateAccessor = createStateAccessor(getRuntimeContext(), getKeyedStateBackend());
  this.stateBootstrapper = new StateBootstrapper(stateBootstrapFunctionRegistry, stateAccessor);
}
 

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

@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();
    }
}
 

@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();
    }
}
 

@Override
public void initializeState(StateInitializationContext context) throws Exception {
  super.initializeState(context);

  final State stateAccessor =
      createStateAccessor(getRuntimeContext(), getKeyedStateBackend(), disableMultiplexState);
  this.stateBootstrapper = new StateBootstrapper(stateBootstrapFunctionRegistry, stateAccessor);
}
 

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

@Override
public void initializeState(StateInitializationContext context) throws Exception {
	super.initializeState(context);
	recoveredStreamElements = context
		.getOperatorStateStore()
		.getListState(new ListStateDescriptor<>(STATE_NAME, inStreamElementSerializer));

}