org.apache.flink.core.fs.CloseableRegistry Java Examples

SubtaskCheckpointCoordinatorImpl(
		CheckpointStorageWorkerView checkpointStorage,
		String taskName,
		StreamTaskActionExecutor actionExecutor,
		CloseableRegistry closeableRegistry,
		ExecutorService executorService,
		Environment env,
		AsyncExceptionHandler asyncExceptionHandler,
		boolean unalignedCheckpointEnabled,
		BiFunctionWithException<ChannelStateWriter, Long, CompletableFuture<Void>, IOException> prepareInputSnapshot,
		int maxRecordAbortedCheckpoints) throws IOException {
	this(
		checkpointStorage,
		taskName,
		actionExecutor,
		closeableRegistry,
		executorService,
		env,
		asyncExceptionHandler,
		prepareInputSnapshot,
		maxRecordAbortedCheckpoints,
		unalignedCheckpointEnabled ? openChannelStateWriter(taskName, checkpointStorage) : ChannelStateWriter.NO_OP);
}
 

protected <K> AbstractKeyedStateBackend<K> restoreKeyedBackend(
		TypeSerializer<K> keySerializer,
		int numberOfKeyGroups,
		KeyGroupRange keyGroupRange,
		List<KeyedStateHandle> state,
		Environment env) throws Exception {

	AbstractKeyedStateBackend<K> backend = getStateBackend().createKeyedStateBackend(
		env,
		new JobID(),
		"test_op",
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		env.getTaskKvStateRegistry(),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		state,
		new CloseableRegistry());

	return backend;
}
 

private HeapKeyedStateBackend<Long> getLongHeapKeyedStateBackend(final long key) throws BackendBuildingException {
	final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 0);
	ExecutionConfig executionConfig = new ExecutionConfig();
	// objects for heap state list serialisation
	final HeapKeyedStateBackend<Long> longHeapKeyedStateBackend =
		new HeapKeyedStateBackendBuilder<>(
			mock(TaskKvStateRegistry.class),
			LongSerializer.INSTANCE,
			ClassLoader.getSystemClassLoader(),
			keyGroupRange.getNumberOfKeyGroups(),
			keyGroupRange,
			executionConfig,
			TtlTimeProvider.DEFAULT,
			Collections.emptyList(),
			AbstractStateBackend.getCompressionDecorator(executionConfig),
			TestLocalRecoveryConfig.disabled(),
			new HeapPriorityQueueSetFactory(keyGroupRange, keyGroupRange.getNumberOfKeyGroups(), 128),
			async,
			new CloseableRegistry()).build();
	longHeapKeyedStateBackend.setCurrentKey(key);
	return longHeapKeyedStateBackend;
}
 

private static AbstractKeyedStateBackend<Integer> createKeyedStateBackend() {
	MockStateBackend backend = new MockStateBackend();

	return backend.createKeyedStateBackend(
		new DummyEnvironment(),
		new JobID(),
		"mock-backend",
		IntSerializer.INSTANCE,
		129,
		KeyGroupRange.of(0, 128),
		null,
		TtlTimeProvider.DEFAULT,
		UnregisteredMetricGroups.createUnregisteredTaskMetricGroup(),
		Collections.emptyList(),
		new CloseableRegistry());
}
 

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

public StreamOperatorStateHandler(
		StreamOperatorStateContext context,
		ExecutionConfig executionConfig,
		CloseableRegistry closeableRegistry) {
	this.context = context;
	operatorStateBackend = context.operatorStateBackend();
	keyedStateBackend = context.keyedStateBackend();
	this.closeableRegistry = closeableRegistry;

	if (keyedStateBackend != null) {
		keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, executionConfig);
	}
	else {
		keyedStateStore = null;
	}
}
 

private <K> AbstractKeyedStateBackend<K> restoreKeyedBackend(
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	List<KeyedStateHandle> state,
	Environment env) throws Exception {
	AbstractKeyedStateBackend<K> backend = getStateBackend().createKeyedStateBackend(
		env,
		new JobID(),
		"test_op",
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		env.getTaskKvStateRegistry(),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		state,
		new CloseableRegistry());
	return backend;
}
 

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

HeapSnapshotStrategy(
	SnapshotStrategySynchronicityBehavior<K> snapshotStrategySynchronicityTrait,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	StreamCompressionDecorator keyGroupCompressionDecorator,
	LocalRecoveryConfig localRecoveryConfig,
	KeyGroupRange keyGroupRange,
	CloseableRegistry cancelStreamRegistry,
	StateSerializerProvider<K> keySerializerProvider) {
	super("Heap backend snapshot");
	this.snapshotStrategySynchronicityTrait = snapshotStrategySynchronicityTrait;
	this.registeredKVStates = registeredKVStates;
	this.registeredPQStates = registeredPQStates;
	this.keyGroupCompressionDecorator = keyGroupCompressionDecorator;
	this.localRecoveryConfig = localRecoveryConfig;
	this.keyGroupRange = keyGroupRange;
	this.cancelStreamRegistry = cancelStreamRegistry;
	this.keySerializerProvider = keySerializerProvider;
}
 

private HeapSnapshotStrategy<K> initSnapshotStrategy(
	boolean asynchronousSnapshots,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	CloseableRegistry cancelStreamRegistry) {
	SnapshotStrategySynchronicityBehavior<K> synchronicityTrait = asynchronousSnapshots ?
		new AsyncSnapshotStrategySynchronicityBehavior<>() :
		new SyncSnapshotStrategySynchronicityBehavior<>();
	return new HeapSnapshotStrategy<>(
		synchronicityTrait,
		registeredKVStates,
		registeredPQStates,
		keyGroupCompressionDecorator,
		localRecoveryConfig,
		keyGroupRange,
		cancelStreamRegistry,
		keySerializerProvider);
}
 

public AbstractKeyedStateBackend(
	TaskKvStateRegistry kvStateRegistry,
	TypeSerializer<K> keySerializer,
	ClassLoader userCodeClassLoader,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	CloseableRegistry cancelStreamRegistry,
	StreamCompressionDecorator keyGroupCompressionDecorator,
	InternalKeyContext<K> keyContext) {
	this.keyContext = Preconditions.checkNotNull(keyContext);
	this.numberOfKeyGroups = keyContext.getNumberOfKeyGroups();
	this.keyGroupRange = Preconditions.checkNotNull(keyContext.getKeyGroupRange());
	Preconditions.checkArgument(numberOfKeyGroups >= 1, "NumberOfKeyGroups must be a positive number");
	Preconditions.checkArgument(numberOfKeyGroups >= keyGroupRange.getNumberOfKeyGroups(), "The total number of key groups must be at least the number in the key group range assigned to this backend");

	this.kvStateRegistry = kvStateRegistry;
	this.keySerializer = keySerializer;
	this.userCodeClassLoader = Preconditions.checkNotNull(userCodeClassLoader);
	this.cancelStreamRegistry = cancelStreamRegistry;
	this.keyValueStatesByName = new HashMap<>();
	this.executionConfig = executionConfig;
	this.keyGroupCompressionDecorator = keyGroupCompressionDecorator;
	this.ttlTimeProvider = Preconditions.checkNotNull(ttlTimeProvider);
	this.keySelectionListeners = new ArrayList<>(1);
}
 

private HeapSnapshotStrategy<K> initSnapshotStrategy(
	boolean asynchronousSnapshots,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	CloseableRegistry cancelStreamRegistry) {
	SnapshotStrategySynchronicityBehavior<K> synchronicityTrait = asynchronousSnapshots ?
		new AsyncSnapshotStrategySynchronicityBehavior<>() :
		new SyncSnapshotStrategySynchronicityBehavior<>();
	return new HeapSnapshotStrategy<>(
		synchronicityTrait,
		registeredKVStates,
		registeredPQStates,
		keyGroupCompressionDecorator,
		localRecoveryConfig,
		keyGroupRange,
		cancelStreamRegistry,
		keySerializerProvider);
}
 

public AbstractKeyedStateBackend(
	TaskKvStateRegistry kvStateRegistry,
	TypeSerializer<K> keySerializer,
	ClassLoader userCodeClassLoader,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	CloseableRegistry cancelStreamRegistry,
	InternalKeyContext<K> keyContext) {
	this(
		kvStateRegistry,
		keySerializer,
		userCodeClassLoader,
		executionConfig,
		ttlTimeProvider,
		cancelStreamRegistry,
		determineStreamCompression(executionConfig),
		keyContext
	);
}
 

public <K> HeapKeyedStateBackend<K> createKeyedBackend(
	TypeSerializer<K> keySerializer,
	Collection<KeyedStateHandle> stateHandles) throws Exception {
	final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 15);
	final int numKeyGroups = keyGroupRange.getNumberOfKeyGroups();
	ExecutionConfig executionConfig = new ExecutionConfig();

	return new HeapKeyedStateBackendBuilder<>(
		mock(TaskKvStateRegistry.class),
		keySerializer,
		HeapStateBackendTestBase.class.getClassLoader(),
		numKeyGroups,
		keyGroupRange,
		executionConfig,
		TtlTimeProvider.DEFAULT,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		TestLocalRecoveryConfig.disabled(),
		new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128),
		async,
		new CloseableRegistry()).build();
}
 

public RocksFullSnapshotStrategy(
	@Nonnull RocksDB db,
	@Nonnull ResourceGuard rocksDBResourceGuard,
	@Nonnull TypeSerializer<K> keySerializer,
	@Nonnull LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	@Nonnull KeyGroupRange keyGroupRange,
	@Nonnegative int keyGroupPrefixBytes,
	@Nonnull LocalRecoveryConfig localRecoveryConfig,
	@Nonnull CloseableRegistry cancelStreamRegistry,
	@Nonnull StreamCompressionDecorator keyGroupCompressionDecorator) {
	super(
		DESCRIPTION,
		db,
		rocksDBResourceGuard,
		keySerializer,
		kvStateInformation,
		keyGroupRange,
		keyGroupPrefixBytes,
		localRecoveryConfig,
		cancelStreamRegistry);

	this.keyGroupCompressionDecorator = keyGroupCompressionDecorator;
}
 

protected <K> AbstractKeyedStateBackend<K> createKeyedBackend(
		TypeSerializer<K> keySerializer,
		int numberOfKeyGroups,
		KeyGroupRange keyGroupRange,
		Environment env) throws Exception {

	AbstractKeyedStateBackend<K> backend = getStateBackend().createKeyedStateBackend(
		env,
		new JobID(),
		"test_op",
		keySerializer,
		numberOfKeyGroups,
		keyGroupRange,
		env.getTaskKvStateRegistry(),
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());

	return backend;
}
 

@Override
public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
	Environment env,
	JobID jobID,
	String operatorIdentifier,
	TypeSerializer<K> keySerializer,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	TaskKvStateRegistry kvStateRegistry,
	TtlTimeProvider ttlTimeProvider,
	MetricGroup metricGroup,
	@Nonnull Collection<KeyedStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) {
	return new MockKeyedStateBackendBuilder<>(
		new KvStateRegistry().createTaskRegistry(jobID, new JobVertexID()),
		keySerializer,
		env.getUserClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		env.getExecutionConfig(),
		ttlTimeProvider,
		stateHandles,
		AbstractStateBackend.getCompressionDecorator(env.getExecutionConfig()),
		cancelStreamRegistry).build();
}
 

HeapRestoreOperation(
	@Nonnull Collection<KeyedStateHandle> restoreStateHandles,
	StateSerializerProvider<K> keySerializerProvider,
	ClassLoader userCodeClassLoader,
	Map<String, StateTable<K, ?, ?>> registeredKVStates,
	Map<String, HeapPriorityQueueSnapshotRestoreWrapper> registeredPQStates,
	CloseableRegistry cancelStreamRegistry,
	HeapPriorityQueueSetFactory priorityQueueSetFactory,
	@Nonnull KeyGroupRange keyGroupRange,
	int numberOfKeyGroups,
	HeapSnapshotStrategy<K> snapshotStrategy,
	HeapKeyedStateBackend<K> backend) {
	this.restoreStateHandles = restoreStateHandles;
	this.keySerializerProvider = keySerializerProvider;
	this.userCodeClassLoader = userCodeClassLoader;
	this.registeredKVStates = registeredKVStates;
	this.registeredPQStates = registeredPQStates;
	this.cancelStreamRegistry = cancelStreamRegistry;
	this.priorityQueueSetFactory = priorityQueueSetFactory;
	this.keyGroupRange = keyGroupRange;
	this.numberOfKeyGroups = numberOfKeyGroups;
	this.snapshotStrategy = snapshotStrategy;
	this.backend = backend;
}
 

private static HeapKeyedStateBackend<Long> createHeapKeyedStateBackend(File rootDir) throws IOException {
	File recoveryBaseDir = prepareDirectory(recoveryDirName, rootDir);
	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
	int numberOfKeyGroups = keyGroupRange.getNumberOfKeyGroups();
	ExecutionConfig executionConfig = new ExecutionConfig();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	HeapKeyedStateBackendBuilder<Long> backendBuilder = new HeapKeyedStateBackendBuilder<>(
		null,
		new LongSerializer(),
		Thread.currentThread().getContextClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		executionConfig,
		TtlTimeProvider.DEFAULT,
		Collections.emptyList(),
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		new LocalRecoveryConfig(false, new LocalRecoveryDirectoryProviderImpl(recoveryBaseDir, new JobID(), new JobVertexID(), 0)),
		priorityQueueSetFactory,
		false,
		new CloseableRegistry()
	);
	return backendBuilder.build();
}
 

private void validateSupportForAsyncSnapshots(StateBackend backend) throws Exception {

		AbstractKeyedStateBackend<Integer> keyedStateBackend = backend.createKeyedStateBackend(
			new DummyEnvironment("Test", 1, 0),
			new JobID(),
			"testOperator",
			IntSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			null,
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			new CloseableRegistry()
		);

		assertTrue(keyedStateBackend.supportsAsynchronousSnapshots());

		IOUtils.closeQuietly(keyedStateBackend);
		keyedStateBackend.dispose();
	}
 

@Before
public void setup() throws IOException {
	ownerRegistry = new CloseableRegistry();
	testProvidedResource = new TestBlockingCloseable();
	testBlocker = new TestBlockingCloseable();
	testAsyncSnapshotCallable = new TestAsyncSnapshotCallable(testProvidedResource, testBlocker);
	task = testAsyncSnapshotCallable.toAsyncSnapshotFutureTask(ownerRegistry);
	Assert.assertEquals(1, ownerRegistry.getNumberOfRegisteredCloseables());
}
 

void createAndRestoreKeyedStateBackend(int numberOfKeyGroups, KeyedStateHandle snapshot) {
	Collection<KeyedStateHandle> stateHandles;
	if (snapshot == null) {
		stateHandles = Collections.emptyList();
	} else {
		stateHandles = new ArrayList<>(1);
		stateHandles.add(snapshot);
	}
	Environment env = new DummyEnvironment();
	try {
		disposeKeyedStateBackend();
		keyedStateBackend = stateBackend.createKeyedStateBackend(
			env,
			new JobID(),
			"test",
			StringSerializer.INSTANCE,
			numberOfKeyGroups,
			new KeyGroupRange(0, numberOfKeyGroups - 1),
			env.getTaskKvStateRegistry(),
			timeProvider,
			new UnregisteredMetricsGroup(),
			stateHandles,
			new CloseableRegistry());
	} catch (Exception e) {
		throw new RuntimeException("unexpected", e);
	}
}
 

private static OperatorStateBackend recreateOperatorStateBackend(
	OperatorStateBackend oldOperatorStateBackend,
	AbstractStateBackend abstractStateBackend,
	Collection<OperatorStateHandle> toRestore
) throws Exception {
	oldOperatorStateBackend.close();
	oldOperatorStateBackend.dispose();
	return abstractStateBackend.createOperatorStateBackend(
		createMockEnvironment(),
		"testOperator",
		toRestore,
		new CloseableRegistry());
}
 

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

	throw new SuccessException();
}
 

/**
 * Upload all the files to checkpoint fileSystem using specified number of threads.
 *
 * @param files The files will be uploaded to checkpoint filesystem.
 * @param checkpointStreamFactory The checkpoint streamFactory used to create outputstream.
 *
 * @throws Exception Thrown if can not upload all the files.
 */
public Map<StateHandleID, StreamStateHandle> uploadFilesToCheckpointFs(
	@Nonnull Map<StateHandleID, Path> files,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws Exception {

	Map<StateHandleID, StreamStateHandle> handles = new HashMap<>();

	Map<StateHandleID, CompletableFuture<StreamStateHandle>> futures =
		createUploadFutures(files, checkpointStreamFactory, closeableRegistry);

	try {
		FutureUtils.waitForAll(futures.values()).get();

		for (Map.Entry<StateHandleID, CompletableFuture<StreamStateHandle>> entry : futures.entrySet()) {
			handles.put(entry.getKey(), entry.getValue().get());
		}
	} catch (ExecutionException e) {
		Throwable throwable = ExceptionUtils.stripExecutionException(e);
		throwable = ExceptionUtils.stripException(throwable, RuntimeException.class);
		if (throwable instanceof IOException) {
			throw (IOException) throwable;
		} else {
			throw new FlinkRuntimeException("Failed to upload data for state handles.", e);
		}
	}

	return handles;
}
 

public DeclineSinkFailingOperatorStateBackend(
	ExecutionConfig executionConfig,
	CloseableRegistry closeStreamOnCancelRegistry,
	AbstractSnapshotStrategy<OperatorStateHandle> snapshotStrategy) {
	super(executionConfig,
		closeStreamOnCancelRegistry,
		new HashMap<>(),
		new HashMap<>(),
		new HashMap<>(),
		new HashMap<>(),
		snapshotStrategy);
}
 

@VisibleForTesting
public StateSnapshotContextSynchronousImpl(long checkpointId, long checkpointTimestamp) {
	this.checkpointId = checkpointId;
	this.checkpointTimestamp = checkpointTimestamp;
	this.streamFactory = null;
	this.keyGroupRange = KeyGroupRange.EMPTY_KEY_GROUP_RANGE;
	this.closableRegistry = new CloseableRegistry();
}
 

@Override
public OperatorStateBackend createOperatorStateBackend(
	Environment env,
	String operatorIdentifier,
	@Nonnull Collection<OperatorStateHandle> stateHandles,
	CloseableRegistry cancelStreamRegistry) throws Exception {
	OperatorStateBackend operatorStateBackend = mock(OperatorStateBackend.class);
	when(operatorStateBackend.snapshot(anyLong(), anyLong(), any(CheckpointStreamFactory.class), any(CheckpointOptions.class)))
		.thenReturn(new FutureTask<>(new BlockingCallable()));

	return operatorStateBackend;
}
 

public OperatorStateRestoreOperation(
	CloseableRegistry closeStreamOnCancelRegistry,
	ClassLoader userClassloader,
	Map<String, PartitionableListState<?>> registeredOperatorStates,
	Map<String, BackendWritableBroadcastState<?, ?>> registeredBroadcastStates,
	@Nonnull Collection<OperatorStateHandle> stateHandles) {
	this.closeStreamOnCancelRegistry = closeStreamOnCancelRegistry;
	this.userClassloader = userClassloader;
	this.registeredOperatorStates = registeredOperatorStates;
	this.registeredBroadcastStates = registeredBroadcastStates;
	this.stateHandles = stateHandles;
}
 

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