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

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {
	if (operatorIdentifier.contains(DECLINE_SINK_NAME)) {
		return new DeclineSinkFailingOperatorStateBackend(
			env.getExecutionConfig(),
			cancelStreamRegistry,
			new DeclineSinkFailingSnapshotStrategy());
	} else {
		return new DefaultOperatorStateBackendBuilder(
			env.getUserClassLoader(),
			env.getExecutionConfig(),
			false,
			stateHandles,
			cancelStreamRegistry).build();
	}
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	//the default for RocksDB; eventually there can be a operator state backend based on RocksDB, too.
	final boolean asyncSnapshots = true;
	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		asyncSnapshots,
		stateHandles,
		cancelStreamRegistry).build();
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	//the default for RocksDB; eventually there can be a operator state backend based on RocksDB, too.
	final boolean asyncSnapshots = true;
	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		asyncSnapshots,
		stateHandles,
		cancelStreamRegistry).build();
}
 

public static void snapshotFunctionState(
		StateSnapshotContext context,
		OperatorStateBackend backend,
		Function userFunction) throws Exception {

	Preconditions.checkNotNull(context);
	Preconditions.checkNotNull(backend);

	while (true) {

		if (trySnapshotFunctionState(context, backend, userFunction)) {
			break;
		}

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

@Override
public BufferedOutputManager<OutputT> create(
    Output<StreamRecord<WindowedValue<OutputT>>> output,
    Lock bufferLock,
    OperatorStateBackend operatorStateBackend)
    throws Exception {
  Preconditions.checkNotNull(output);
  Preconditions.checkNotNull(bufferLock);
  Preconditions.checkNotNull(operatorStateBackend);

  TaggedKvCoder taggedKvCoder = buildTaggedKvCoder();
  ListStateDescriptor<KV<Integer, WindowedValue<?>>> taggedOutputPushbackStateDescriptor =
      new ListStateDescriptor<>("bundle-buffer-tag", new CoderTypeSerializer<>(taggedKvCoder));
  ListState<KV<Integer, WindowedValue<?>>> listStateBuffer =
      operatorStateBackend.getListState(taggedOutputPushbackStateDescriptor);
  PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler =
      NonKeyedPushedBackElementsHandler.create(listStateBuffer);

  return new BufferedOutputManager<>(
      output, mainTag, tagsToOutputTags, tagsToIds, bufferLock, pushedBackElementsHandler);
}
 

@Override
public void open(OperatorStateInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(handles) -> createOperatorStateBackend(getRuntimeContext(), handles, registry),
			registry,
			operatorState.getOperatorID().toString()
		);

	try {
		restoredBackend = backendRestorer.createAndRestore(split.getPrioritizedManagedOperatorState());
	} catch (Exception exception) {
		throw new IOException("Failed to restore state backend", exception);
	}

	try {
		elements = getElements(restoredBackend).iterator();
	} catch (Exception e) {
		throw new IOException("Failed to read operator state from restored state backend", e);
	}
}
 

private static OperatorStateBackend createOperatorStateBackend(
	RuntimeContext runtimeContext,
	Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry){

	try {
		return new DefaultOperatorStateBackendBuilder(
			runtimeContext.getUserCodeClassLoader(),
			runtimeContext.getExecutionConfig(),
			false,
			stateHandles,
			cancelStreamRegistry
		).build();
	} catch (BackendBuildingException e) {
		throw new RuntimeException(e);
	}
}
 

public static <InputT, OutputT> BufferingDoFnRunner<InputT, OutputT> create(
    DoFnRunner<InputT, OutputT> doFnRunner,
    String stateName,
    org.apache.beam.sdk.coders.Coder windowedInputCoder,
    org.apache.beam.sdk.coders.Coder windowCoder,
    OperatorStateBackend operatorStateBackend,
    @Nullable KeyedStateBackend<Object> keyedStateBackend,
    int maxConcurrentCheckpoints)
    throws Exception {
  return new BufferingDoFnRunner<>(
      doFnRunner,
      stateName,
      windowedInputCoder,
      windowCoder,
      operatorStateBackend,
      keyedStateBackend,
      maxConcurrentCheckpoints);
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	//the default for RocksDB; eventually there can be a operator state backend based on RocksDB, too.
	final boolean asyncSnapshots = true;
	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		asyncSnapshots,
		stateHandles,
		cancelStreamRegistry).build();
}
 

FlinkCombiningStateWithContext(
    OperatorStateBackend flinkStateBackend,
    StateTag<CombiningState<InputT, AccumT, OutputT>> address,
    CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,
    StateNamespace namespace,
    Coder<AccumT> accumCoder,
    FlinkBroadcastStateInternals<K2> flinkStateInternals,
    CombineWithContext.Context context) {
  super(flinkStateBackend, address.getId(), namespace, accumCoder);

  this.namespace = namespace;
  this.address = address;
  this.combineFn = combineFn;
  this.flinkStateInternals = flinkStateInternals;
  this.context = context;
}
 

AbstractBroadcastState(
    OperatorStateBackend flinkStateBackend,
    String name,
    StateNamespace namespace,
    Coder<T> coder) {
  this.name = name;

  this.namespace = namespace;
  this.flinkStateBackend = flinkStateBackend;

  CoderTypeInformation<Map<String, T>> typeInfo =
      new CoderTypeInformation<>(MapCoder.of(StringUtf8Coder.of(), coder));

  flinkStateDescriptor =
      new ListStateDescriptor<>(name, typeInfo.createSerializer(new ExecutionConfig()));
}
 

@Override
public void open(OperatorStateInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(handles) -> createOperatorStateBackend(getRuntimeContext(), handles, registry),
			registry,
			operatorState.getOperatorID().toString()
		);

	try {
		restoredBackend = backendRestorer.createAndRestore(split.getPrioritizedManagedOperatorState());
	} catch (Exception exception) {
		throw new IOException("Failed to restore state backend", exception);
	}

	try {
		elements = getElements(restoredBackend).iterator();
	} catch (Exception e) {
		throw new IOException("Failed to read operator state from restored state backend", e);
	}
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		true,
		stateHandles,
		cancelStreamRegistry) {
		@Override
		@SuppressWarnings("unchecked")
		public DefaultOperatorStateBackend build() {
			return new DefaultOperatorStateBackend(
				executionConfig,
				cancelStreamRegistry,
				new HashMap<>(),
				new HashMap<>(),
				new HashMap<>(),
				new HashMap<>(),
				mock(AbstractSnapshotStrategy.class)
			) {
				@Nonnull
				@Override
				public RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot(
					long checkpointId,
					long timestamp,
					@Nonnull CheckpointStreamFactory streamFactory,
					@Nonnull CheckpointOptions checkpointOptions) throws Exception {

					return new FutureTask<>(() -> {
						throw new Exception("Async part snapshot exception.");
					});
				}
			};
		}
	}.build();
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		true,
		stateHandles,
		cancelStreamRegistry) {
		@Override
		@SuppressWarnings("unchecked")
		public DefaultOperatorStateBackend build() {
			return new DefaultOperatorStateBackend(
				executionConfig,
				cancelStreamRegistry,
				new HashMap<>(),
				new HashMap<>(),
				new HashMap<>(),
				new HashMap<>(),
				mock(AbstractSnapshotStrategy.class)
			) {
				@Nonnull
				@Override
				public RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot(
					long checkpointId,
					long timestamp,
					@Nonnull CheckpointStreamFactory streamFactory,
					@Nonnull CheckpointOptions checkpointOptions) throws Exception {

					throw new Exception("Sync part snapshot exception.");
				}
			};
		}
	}.build();
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	return backend.createOperatorStateBackend(env, operatorIdentifier, stateHandles, cancelStreamRegistry);
}
 

private StateObjectCollection<OperatorStateHandle> transformSubtaskOpState(Path outDir, Integer subtaskId,
		StateObjectCollection<OperatorStateHandle> baseState) {

	if (transformer == null) {
		return baseState;
	}

	StateObjectCollection<OperatorStateHandle> opHandle = baseState;
	try (OperatorStateBackend opBackend = OperatorStateReader
			.restoreOperatorStateBackend(opHandle)) {

		transformer.accept(subtaskId, opBackend);

		OperatorStateHandle newSnapshot = opBackend
				.snapshot(checkpointId, System.currentTimeMillis(), new CheckpointStreamFactory() {
					@Override
					public CheckpointStateOutputStream createCheckpointStateOutputStream(
							CheckpointedStateScope scope)
							throws IOException {
						return new FileBasedStateOutputStream(outDir.getFileSystem(),
								new Path(outDir, String.valueOf(UUID.randomUUID())));
					}
				}, null).get().getJobManagerOwnedSnapshot();
		return new StateObjectCollection<>(Lists.newArrayList(newSnapshot));
	} catch (Exception e) {
		throw new RuntimeException(e);
	}
}
 

/**
 * Restores the OperatorStateBackends corresponding to the different
 * subtasks. The backends are completely restored in-memory.
 */
public Map<Integer, OperatorStateBackend> createOperatorStateBackendsFromSnapshot() throws Exception {
	return Maps.transformValues(opState.getSubtaskStates(),
			sst -> {
				try {
					return restoreOperatorStateBackend(sst.getManagedOperatorState());
				} catch (Exception e) {
					throw new RuntimeException(e);
				}
			});
}
 

/**
 * Read the serializableListState stored in the checkpoint for the given
 * operator subtask
 */
public List<Serializable> getSerializableListState(int subtask) throws Exception {
	OperatorStateBackend backend = createOperatorStateBackendFromSnapshot(subtask);
	@SuppressWarnings("deprecation")
	ListState<Serializable> listState = backend
			.getSerializableListState(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME);

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

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

	return list;
}
 

private void validateSecondSavepoint(Savepoint savepoint) throws Exception {
	// Validate the contents of the broadcast state
	OperatorStateReader reader = new OperatorStateReader(ExecutionEnvironment.createLocalEnvironment(), savepoint,
			"stateful");
	OperatorStateBackend backend = reader.createOperatorStateBackendFromSnapshot(1);
	BroadcastState<Boolean, List<Integer>> state = backend.getBroadcastState(bcstate);
	assertEquals(Lists.newArrayList(2, 3), state.get(true));
}
 

private static boolean trySnapshotFunctionState(
		StateSnapshotContext context,
		OperatorStateBackend backend,
		Function userFunction) throws Exception {

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

		return true;
	}

	if (userFunction instanceof ListCheckpointed) {
		@SuppressWarnings("unchecked")
		List<Serializable> partitionableState = ((ListCheckpointed<Serializable>) userFunction).
				snapshotState(context.getCheckpointId(), context.getCheckpointTimestamp());

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

		listState.clear();

		if (null != partitionableState) {
			try {
				for (Serializable statePartition : partitionableState) {
					listState.add(statePartition);
				}
			} catch (Exception e) {
				listState.clear();

				throw new Exception("Could not write partitionable state to operator " +
					"state backend.", e);
			}
		}

		return true;
	}

	return false;
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		true,
		stateHandles,
		cancelStreamRegistry) {
		@Override
		@SuppressWarnings("unchecked")
		public DefaultOperatorStateBackend build() {
			return new DefaultOperatorStateBackend(
				executionConfig,
				cancelStreamRegistry,
				new HashMap<>(),
				new HashMap<>(),
				new HashMap<>(),
				new HashMap<>(),
				mock(AbstractSnapshotStrategy.class)
			) {
				@Nonnull
				@Override
				public RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot(
					long checkpointId,
					long timestamp,
					@Nonnull CheckpointStreamFactory streamFactory,
					@Nonnull CheckpointOptions checkpointOptions) throws Exception {

					throw new Exception("Sync part snapshot exception.");
				}
			};
		}
	}.build();
}
 

FlinkBroadcastBagState(
    OperatorStateBackend flinkStateBackend,
    StateTag<BagState<T>> address,
    StateNamespace namespace,
    Coder<T> coder) {
  super(flinkStateBackend, address.getId(), namespace, ListCoder.of(coder));

  this.namespace = namespace;
  this.address = address;
}
 

FlinkCombiningState(
    OperatorStateBackend flinkStateBackend,
    StateTag<CombiningState<InputT, AccumT, OutputT>> address,
    Combine.CombineFn<InputT, AccumT, OutputT> combineFn,
    StateNamespace namespace,
    Coder<AccumT> accumCoder) {
  super(flinkStateBackend, address.getId(), namespace, accumCoder);

  this.namespace = namespace;
  this.address = address;
  this.combineFn = combineFn;
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	throw new SuccessException();
}
 

StreamOperatorStateContextImpl(
	boolean restored,
	OperatorStateBackend operatorStateBackend,
	AbstractKeyedStateBackend<?> keyedStateBackend,
	InternalTimeServiceManager<?> internalTimeServiceManager,
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs,
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs) {

	this.restored = restored;
	this.operatorStateBackend = operatorStateBackend;
	this.keyedStateBackend = keyedStateBackend;
	this.internalTimeServiceManager = internalTimeServiceManager;
	this.rawOperatorStateInputs = rawOperatorStateInputs;
	this.rawKeyedStateInputs = rawKeyedStateInputs;
}
 

protected OperatorStateBackend operatorStateBackend(
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry) throws Exception {

	String logDescription = "operator state backend for " + operatorIdentifierText;

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createOperatorStateBackend(
				environment,
				operatorIdentifierText,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedOperatorState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

FlinkKeyedCombiningState(
    OperatorStateBackend flinkStateBackend,
    StateTag<CombiningState<InputT, AccumT, OutputT>> address,
    Combine.CombineFn<InputT, AccumT, OutputT> combineFn,
    StateNamespace namespace,
    Coder<AccumT> accumCoder,
    FlinkBroadcastStateInternals<K2> flinkStateInternals) {
  super(flinkStateBackend, address.getId(), namespace, accumCoder);

  this.namespace = namespace;
  this.address = address;
  this.combineFn = combineFn;
  this.flinkStateInternals = flinkStateInternals;
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {

	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		isUsingAsynchronousSnapshots(),
		stateHandles,
		cancelStreamRegistry).build();
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	throw new UnsupportedOperationException();
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws BackendBuildingException {

	return new DefaultOperatorStateBackendBuilder(
		env.getUserClassLoader(),
		env.getExecutionConfig(),
		isUsingAsynchronousSnapshots(),
		stateHandles,
		cancelStreamRegistry).build();
}