Java Code Examples for org.apache.flink.util.IOUtils#closeQuietly()

/**
 * Tests that the TaskManagerRunner startup fails if the network stack cannot be initialized.
 */
@Test
public void testStartupWhenNetworkStackFailsToInitialize() throws Exception {
	final ServerSocket blocker = new ServerSocket(0, 50, InetAddress.getByName(LOCAL_HOST));

	try {
		final Configuration cfg = new Configuration();
		cfg.setInteger(TaskManagerOptions.DATA_PORT, blocker.getLocalPort());

		startTaskManager(
			cfg,
			rpcService,
			highAvailabilityServices);

		fail("Should throw IOException when the network stack cannot be initialized.");
	} catch (IOException e) {
		// ignored
	} finally {
		IOUtils.closeQuietly(blocker);
	}
}
 

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

private PriorityQueue<RocksSingleStateIterator> buildIteratorHeap(
	List<Tuple2<RocksIteratorWrapper, Integer>> kvStateIterators) {

	Comparator<RocksSingleStateIterator> iteratorComparator = COMPARATORS.get(keyGroupPrefixByteCount - 1);

	PriorityQueue<RocksSingleStateIterator> iteratorPriorityQueue =
		new PriorityQueue<>(kvStateIterators.size(), iteratorComparator);

	for (Tuple2<RocksIteratorWrapper, Integer> rocksIteratorWithKVStateId : kvStateIterators) {
		final RocksIteratorWrapper rocksIterator = rocksIteratorWithKVStateId.f0;
		rocksIterator.seekToFirst();
		if (rocksIterator.isValid()) {
			iteratorPriorityQueue.offer(
				new RocksSingleStateIterator(rocksIterator, rocksIteratorWithKVStateId.f1));
		} else {
			IOUtils.closeQuietly(rocksIterator);
		}
	}
	return iteratorPriorityQueue;
}
 

@Test
public void testNonConcurrentSnapshotTransformerAccess() throws Exception {
	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
	AbstractKeyedStateBackend<Integer> backend = null;
	try {
		backend = createKeyedBackend(IntSerializer.INSTANCE);
		new StateSnapshotTransformerTest(backend, streamFactory)
			.testNonConcurrentSnapshotTransformerAccess();
	} finally {
		if (backend != null) {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

/**
 * This test checks the behavior for basic setting of local DB directories.
 */
@Test
public void testSetDbPath() throws Exception {
	final RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());

	final String testDir1 = tempFolder.newFolder().getAbsolutePath();
	final String testDir2 = tempFolder.newFolder().getAbsolutePath();

	assertNull(rocksDbBackend.getDbStoragePaths());

	rocksDbBackend.setDbStoragePath(testDir1);
	assertArrayEquals(new String[] { testDir1 }, rocksDbBackend.getDbStoragePaths());

	rocksDbBackend.setDbStoragePath(null);
	assertNull(rocksDbBackend.getDbStoragePaths());

	rocksDbBackend.setDbStoragePaths(testDir1, testDir2);
	assertArrayEquals(new String[] { testDir1, testDir2 }, rocksDbBackend.getDbStoragePaths());

	final Environment env = getMockEnvironment(tempFolder.newFolder());
	final RocksDBKeyedStateBackend<Integer> keyedBackend = createKeyedStateBackend(rocksDbBackend, env);

	try {
		File instanceBasePath = keyedBackend.getInstanceBasePath();
		assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(testDir1), startsWith(testDir2)));

		//noinspection NullArgumentToVariableArgMethod
		rocksDbBackend.setDbStoragePaths(null);
		assertNull(rocksDbBackend.getDbStoragePaths());
	}
	finally {
		IOUtils.closeQuietly(keyedBackend);
		keyedBackend.dispose();
	}
}
 

@Override
protected void after() {
	// destruct in reversed order of creation.
	IOUtils.closeQuietly(this.batchWrapper);
	for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
		IOUtils.closeQuietly(columnFamilyHandle);
	}
	IOUtils.closeQuietly(this.rocksDB);
	IOUtils.closeQuietly(this.readOptions);
	IOUtils.closeQuietly(this.writeOptions);
	IOUtils.closeQuietly(this.columnFamilyOptions);
	IOUtils.closeQuietly(this.dbOptions);
	temporaryFolder.delete();
}
 

@Test
public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
	final File targetDir1 = tempFolder.newFolder();
	final File targetDir2 = tempFolder.newFolder();
	Assume.assumeTrue("Cannot mark directory non-writable", targetDir1.setWritable(false, false));

	String checkpointPath = tempFolder.newFolder().toURI().toString();
	RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

	try (MockEnvironment env = getMockEnvironment(tempFolder.newFolder())) {
		rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());

		try {
			AbstractKeyedStateBackend<Integer> keyedStateBackend = rocksDbBackend.createKeyedStateBackend(
				env,
				env.getJobID(),
				"foobar",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());

			IOUtils.closeQuietly(keyedStateBackend);
			keyedStateBackend.dispose();
		}
		catch (Exception e) {
			e.printStackTrace();
			fail("Backend initialization failed even though some paths were available");
		}
	} finally {
		//noinspection ResultOfMethodCallIgnored
		targetDir1.setWritable(true, false);
	}
}
 

/**
 * Closes the safety net for a thread. This closes all remaining unclosed streams that were opened
 * by safety-net-guarded file systems. After this method was called, no streams can be opened any more
 * from any FileSystem instance that was obtained while the thread was guarded by the safety net.
 *
 * <p>This method should be called at the very end of a guarded thread.
 */
@Internal
public static void closeSafetyNetAndGuardedResourcesForThread() {
	SafetyNetCloseableRegistry registry = REGISTRIES.get();
	if (null != registry) {
		REGISTRIES.remove();
		IOUtils.closeQuietly(registry);
	}
}
 

@Override
public void close() {
	IOUtils.closeQuietly(currentSubIterator);
	currentSubIterator = null;

	IOUtils.closeAllQuietly(heap);
	heap.clear();
}
 

private void updateCurrentIterator() {
	IOUtils.closeQuietly(currentIterator);
	if (iteratorQueue.isEmpty()) {
		currentIterator = null;
		currentName = null;
		return;
	} else {
		Entry<String, RocksIteratorWrapper> e = iteratorQueue.pop();
		currentName = e.getKey();
		currentIterator = e.getValue();
	}
}
 

/**
 * This tests whether the RocksDB backends uses the temp directories that are provided
 * from the {@link Environment} when no db storage path is set.
 */
@Test
public void testUseTempDirectories() throws Exception {
	String checkpointPath = tempFolder.newFolder().toURI().toString();
	RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

	File dir1 = tempFolder.newFolder();
	File dir2 = tempFolder.newFolder();

	assertNull(rocksDbBackend.getDbStoragePaths());

	Environment env = getMockEnvironment(dir1, dir2);
	RocksDBKeyedStateBackend<Integer> keyedBackend = (RocksDBKeyedStateBackend<Integer>) rocksDbBackend.
		createKeyedStateBackend(
			env,
			env.getJobID(),
			"test_op",
			IntSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			env.getTaskKvStateRegistry(),
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			new CloseableRegistry());

	try {
		File instanceBasePath = keyedBackend.getInstanceBasePath();
		assertThat(instanceBasePath.getAbsolutePath(), anyOf(startsWith(dir1.getAbsolutePath()), startsWith(dir2.getAbsolutePath())));
	} finally {
		IOUtils.closeQuietly(keyedBackend);
		keyedBackend.dispose();
	}
}
 

@Override
public void close() throws FlinkKafka011Exception {
	// First close the producer for current transaction.
	try {
		final KafkaTransactionState currentTransaction = currentTransaction();
		if (currentTransaction != null) {
			// to avoid exceptions on aborting transactions with some pending records
			flush(currentTransaction);

			// normal abort for AT_LEAST_ONCE and NONE do not clean up resources because of producer reusing, thus
			// we need to close it manually
			switch (semantic) {
				case EXACTLY_ONCE:
					break;
				case AT_LEAST_ONCE:
				case NONE:
					currentTransaction.producer.close();
					break;
			}
		}
		super.close();
	} catch (Exception e) {
		asyncException = ExceptionUtils.firstOrSuppressed(e, asyncException);
	} finally {
		// We may have to close producer of the current transaction in case some exception was thrown before
		// the normal close routine finishes.
		if (currentTransaction() != null) {
			IOUtils.closeQuietly(currentTransaction().producer);
		}
		// Make sure all the producers for pending transactions are closed.
		pendingTransactions().forEach(transaction -> IOUtils.closeQuietly(transaction.getValue().producer)
		);
		// make sure we propagate pending errors
		checkErroneous();
	}
}
 

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
	if (enableIncrementalCheckpointing) {
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
		try {
			ValueStateDescriptor<String> kvId =
				new ValueStateDescriptor<>("id", String.class, null);

			kvId.initializeSerializerUnlessSet(new ExecutionConfig());

			ValueState<String> state =
				backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

			Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
			SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistry());
			for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {

				reset(sharedStateRegistry);

				backend.setCurrentKey(checkpointId);
				state.update("Hello-" + checkpointId);

				RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(
					checkpointId,
					checkpointId,
					createStreamFactory(),
					CheckpointOptions.forCheckpointWithDefaultLocation());

				snapshot.run();

				SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();

				IncrementalRemoteKeyedStateHandle stateHandle =
					(IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();

				Map<StateHandleID, StreamStateHandle> sharedState =
					new HashMap<>(stateHandle.getSharedState());

				stateHandle.registerSharedStates(sharedStateRegistry);

				for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
					verify(sharedStateRegistry).registerReference(
						stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()),
						e.getValue());
				}

				previousStateHandles.add(stateHandle);
				backend.notifyCheckpointComplete(checkpointId);

				//-----------------------------------------------------------------

				if (previousStateHandles.size() > 1) {
					checkRemove(previousStateHandles.remove(), sharedStateRegistry);
				}
			}

			while (!previousStateHandles.isEmpty()) {

				reset(sharedStateRegistry);

				checkRemove(previousStateHandles.remove(), sharedStateRegistry);
			}
		} finally {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

@Test
public void testContinueOnSomeDbDirectoriesMissing() throws Exception {
	File targetDir1 = tempFolder.newFolder();
	File targetDir2 = tempFolder.newFolder();

	String checkpointPath = tempFolder.newFolder().toURI().toString();
	RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(checkpointPath);

	try {

		if (!targetDir1.setWritable(false, false)) {
			System.err.println("Cannot execute 'testContinueOnSomeDbDirectoriesMissing' because cannot mark directory non-writable");
			return;
		}

		rocksDbBackend.setDbStoragePaths(targetDir1.getAbsolutePath(), targetDir2.getAbsolutePath());

		try {
			Environment env = getMockEnvironment(tempFolder.newFolder());
			AbstractKeyedStateBackend<Integer> keyedStateBackend = rocksDbBackend.createKeyedStateBackend(
				env,
				env.getJobID(),
				"foobar",
				IntSerializer.INSTANCE,
				1,
				new KeyGroupRange(0, 0),
				new KvStateRegistry().createTaskRegistry(env.getJobID(), new JobVertexID()),
				TtlTimeProvider.DEFAULT,
				new UnregisteredMetricsGroup(),
				Collections.emptyList(),
				new CloseableRegistry());

			IOUtils.closeQuietly(keyedStateBackend);
			keyedStateBackend.dispose();
		}
		catch (Exception e) {
			e.printStackTrace();
			fail("Backend initialization failed even though some paths were available");
		}
	} finally {
		//noinspection ResultOfMethodCallIgnored
		targetDir1.setWritable(true, false);
	}
}
 

@Override
public void dispose() {
	// we can suppress exceptions here, because we do not rely on close() to
	// flush or persist any data
	IOUtils.closeQuietly(currentPartStream);
}
 

/**
 * Should only be called by one thread, and only after all accesses to the DB happened.
 */
@Override
public void dispose() {
	if (this.disposed) {
		return;
	}
	super.dispose();

	// This call will block until all clients that still acquire access to the RocksDB instance have released it,
	// so that we cannot release the native resources while clients are still working with it in parallel.
	rocksDBResourceGuard.close();

	// IMPORTANT: null reference to signal potential async checkpoint workers that the db was disposed, as
	// working on the disposed object results in SEGFAULTS.
	if (db != null) {

		IOUtils.closeQuietly(writeBatchWrapper);

		// Metric collection occurs on a background thread. When this method returns
		// it is guaranteed that thr RocksDB reference has been invalidated
		// and no more metric collection will be attempted against the database.
		if (nativeMetricMonitor != null) {
			nativeMetricMonitor.close();
		}

		List<ColumnFamilyOptions> columnFamilyOptions = new ArrayList<>(kvStateInformation.values().size());

		// RocksDB's native memory management requires that *all* CFs (including default) are closed before the
		// DB is closed. See:
		// https://github.com/facebook/rocksdb/wiki/RocksJava-Basics#opening-a-database-with-column-families
		// Start with default CF ...
		RocksDBOperationUtils.addColumnFamilyOptionsToCloseLater(columnFamilyOptions, defaultColumnFamily);
		IOUtils.closeQuietly(defaultColumnFamily);

		// ... continue with the ones created by Flink...
		for (RocksDbKvStateInfo kvStateInfo : kvStateInformation.values()) {
			RocksDBOperationUtils.addColumnFamilyOptionsToCloseLater(columnFamilyOptions, kvStateInfo.columnFamilyHandle);
			IOUtils.closeQuietly(kvStateInfo.columnFamilyHandle);
		}

		// ... and finally close the DB instance ...
		IOUtils.closeQuietly(db);

		columnFamilyOptions.forEach(IOUtils::closeQuietly);

		IOUtils.closeQuietly(dbOptions);
		IOUtils.closeQuietly(writeOptions);

		ttlCompactFiltersManager.disposeAndClearRegisteredCompactionFactories();

		kvStateInformation.clear();

		cleanInstanceBasePath();
	}
	this.disposed = true;
}
 

void dispose() {
	// we can suppress exceptions here, because we do not rely on close() to
	// flush or persist any data
	IOUtils.closeQuietly(currentPartStream);
}
 

public void disposeAndClearRegisteredCompactionFactories() {
	for (FlinkCompactionFilterFactory factory : compactionFilterFactories.values()) {
		IOUtils.closeQuietly(factory);
	}
	compactionFilterFactories.clear();
}
 

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
	if (enableIncrementalCheckpointing) {
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
		try {
			ValueStateDescriptor<String> kvId =
				new ValueStateDescriptor<>("id", String.class, null);

			kvId.initializeSerializerUnlessSet(new ExecutionConfig());

			ValueState<String> state =
				backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

			Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
			SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistry());
			for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {

				reset(sharedStateRegistry);

				backend.setCurrentKey(checkpointId);
				state.update("Hello-" + checkpointId);

				RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(
					checkpointId,
					checkpointId,
					createStreamFactory(),
					CheckpointOptions.forCheckpointWithDefaultLocation());

				snapshot.run();

				SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();

				IncrementalRemoteKeyedStateHandle stateHandle =
					(IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();

				Map<StateHandleID, StreamStateHandle> sharedState =
					new HashMap<>(stateHandle.getSharedState());

				stateHandle.registerSharedStates(sharedStateRegistry);

				for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
					verify(sharedStateRegistry).registerReference(
						stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()),
						e.getValue());
				}

				previousStateHandles.add(stateHandle);
				backend.notifyCheckpointComplete(checkpointId);

				//-----------------------------------------------------------------

				if (previousStateHandles.size() > 1) {
					checkRemove(previousStateHandles.remove(), sharedStateRegistry);
				}
			}

			while (!previousStateHandles.isEmpty()) {

				reset(sharedStateRegistry);

				checkRemove(previousStateHandles.remove(), sharedStateRegistry);
			}
		} finally {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

@Test
public void testCorrectMergeOperatorSet() throws Exception {
	prepareRocksDB();
	final ColumnFamilyOptions columnFamilyOptions = spy(new ColumnFamilyOptions());
	RocksDBKeyedStateBackend<Integer> test = null;
	try (DBOptions options = new DBOptions().setCreateIfMissing(true)) {
		ExecutionConfig executionConfig = new ExecutionConfig();
		test = new RocksDBKeyedStateBackendBuilder<>(
			"test",
			Thread.currentThread().getContextClassLoader(),
			tempFolder.newFolder(),
			options,
			stateName -> columnFamilyOptions,
			mock(TaskKvStateRegistry.class),
			IntSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			executionConfig,
			mock(LocalRecoveryConfig.class),
			RocksDBStateBackend.PriorityQueueStateType.HEAP,
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			AbstractStateBackend.getCompressionDecorator(executionConfig),
			db,
			defaultCFHandle,
			new CloseableRegistry())
			.setEnableIncrementalCheckpointing(enableIncrementalCheckpointing)
			.build();
		ValueStateDescriptor<String> stubState1 =
			new ValueStateDescriptor<>("StubState-1", StringSerializer.INSTANCE);
		test.createInternalState(StringSerializer.INSTANCE, stubState1);
		ValueStateDescriptor<String> stubState2 =
			new ValueStateDescriptor<>("StubState-2", StringSerializer.INSTANCE);
		test.createInternalState(StringSerializer.INSTANCE, stubState2);

		// The default CF is pre-created so sum up to 2 times (once for each stub state)
		verify(columnFamilyOptions, Mockito.times(2))
			.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);
	} finally {
		if (test != null) {
			IOUtils.closeQuietly(test);
			test.dispose();
		}
		columnFamilyOptions.close();
	}
}