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

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
  if (checkpointCoder == null) {
    // no checkpoint coder available in this source
    return;
  }

  OperatorStateStore stateStore = context.getOperatorStateStore();
  @SuppressWarnings("unchecked")
  CoderTypeInformation<KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>>
      typeInformation = (CoderTypeInformation) new CoderTypeInformation<>(checkpointCoder);
  stateForCheckpoint =
      stateStore.getOperatorState(
          new ListStateDescriptor<>(
              DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
              typeInformation.createSerializer(new ExecutionConfig())));

  if (context.isRestored()) {
    isRestored = true;
    LOG.info("Restoring state in the UnboundedSourceWrapper.");
  } else {
    LOG.info("No restore state for UnboundedSourceWrapper.");
  }
}
 

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

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

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;

		// 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.
		ListStateDescriptor<Serializable> listStateDescriptor =
			new ListStateDescriptor<>(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, new JavaSerializer<>());
		ListState<Serializable> listState = context.getOperatorStateStore().getListState(listStateDescriptor);

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

protected CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs(
	Iterator<StateObjectCollection<OperatorStateHandle>> restoreStateAlternatives) {

	if (restoreStateAlternatives.hasNext()) {

		Collection<OperatorStateHandle> rawOperatorState = restoreStateAlternatives.next();
		// TODO currently this does not support local state recovery, so we expect there is only one handle.
		Preconditions.checkState(
			!restoreStateAlternatives.hasNext(),
			"Local recovery is currently not implemented for raw operator state, but found state alternative.");

		if (rawOperatorState != null) {
			return new CloseableIterable<StatePartitionStreamProvider>() {

				final CloseableRegistry closeableRegistry = new CloseableRegistry();

				@Override
				public void close() throws IOException {
					closeableRegistry.close();
				}

				@Nonnull
				@Override
				public Iterator<StatePartitionStreamProvider> iterator() {
					return new OperatorStateStreamIterator(
						DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
						rawOperatorState.iterator(), closeableRegistry);
				}
			};
		}
	}

	return CloseableIterable.empty();
}
 

public static OperatorStateBackend restoreOperatorStateBackend(
		StateObjectCollection<OperatorStateHandle> managedOpState)
		throws Exception {

	DefaultOperatorStateBackend stateBackend = new DefaultOperatorStateBackend(
			OperatorStateReader.class.getClassLoader(), new ExecutionConfig(), false);

	stateBackend.restore(managedOpState);
	return stateBackend;
}
 

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

protected CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs(
	Iterator<StateObjectCollection<OperatorStateHandle>> restoreStateAlternatives) {

	if (restoreStateAlternatives.hasNext()) {

		Collection<OperatorStateHandle> rawOperatorState = restoreStateAlternatives.next();
		// TODO currently this does not support local state recovery, so we expect there is only one handle.
		Preconditions.checkState(
			!restoreStateAlternatives.hasNext(),
			"Local recovery is currently not implemented for raw operator state, but found state alternative.");

		if (rawOperatorState != null) {
			return new CloseableIterable<StatePartitionStreamProvider>() {

				final CloseableRegistry closeableRegistry = new CloseableRegistry();

				@Override
				public void close() throws IOException {
					closeableRegistry.close();
				}

				@Nonnull
				@Override
				public Iterator<StatePartitionStreamProvider> iterator() {
					return new OperatorStateStreamIterator(
						DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
						rawOperatorState.iterator(), closeableRegistry);
				}
			};
		}
	}

	return CloseableIterable.empty();
}
 

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

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

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

protected CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs(
	Iterator<StateObjectCollection<OperatorStateHandle>> restoreStateAlternatives) {

	if (restoreStateAlternatives.hasNext()) {

		Collection<OperatorStateHandle> rawOperatorState = restoreStateAlternatives.next();
		// TODO currently this does not support local state recovery, so we expect there is only one handle.
		Preconditions.checkState(
			!restoreStateAlternatives.hasNext(),
			"Local recovery is currently not implemented for raw operator state, but found state alternative.");

		if (rawOperatorState != null) {
			return new CloseableIterable<StatePartitionStreamProvider>() {

				final CloseableRegistry closeableRegistry = new CloseableRegistry();

				@Override
				public void close() throws IOException {
					closeableRegistry.close();
				}

				@Nonnull
				@Override
				public Iterator<StatePartitionStreamProvider> iterator() {
					return new OperatorStateStreamIterator(
						DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME,
						rawOperatorState.iterator(), closeableRegistry);
				}
			};
		}
	}

	return CloseableIterable.empty();
}
 

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

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

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

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

		// 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.
		ListStateDescriptor<Serializable> listStateDescriptor =
			new ListStateDescriptor<>(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, new JavaSerializer<>());
		ListState<Serializable> listState = backend.getListState(listStateDescriptor);

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