org.apache.flink.util.IOUtils Java Examples

@Override
public TransportClient createClient(Map<String, String> clientConfig) {
	Settings settings = Settings.settingsBuilder().put(clientConfig).build();

	TransportClient transportClient = TransportClient.builder().settings(settings).build();
	for (TransportAddress address : ElasticsearchUtils.convertInetSocketAddresses(transportAddresses)) {
		transportClient.addTransportAddress(address);
	}

	// verify that we actually are connected to a cluster
	if (transportClient.connectedNodes().isEmpty()) {

		// close the transportClient here
		IOUtils.closeQuietly(transportClient);

		throw new RuntimeException("Elasticsearch client is not connected to any Elasticsearch nodes!");
	}

	if (LOG.isInfoEnabled()) {
		LOG.info("Created Elasticsearch TransportClient with connected nodes {}", transportClient.connectedNodes());
	}

	return transportClient;
}
 

@Test(expected = IOException.class)
public void testResumeAfterCommit() throws Exception {
	final Path testDir = getBasePathForTest();

	final RecoverableWriter writer = getNewFileSystemWriter();
	final Path path = new Path(testDir, "part-0");

	RecoverableWriter.ResumeRecoverable recoverable;
	RecoverableFsDataOutputStream stream = null;
	try {
		stream = writer.open(path);
		stream.write(testData1.getBytes(StandardCharsets.UTF_8));

		recoverable = stream.persist();
		stream.write(testData2.getBytes(StandardCharsets.UTF_8));

		stream.closeForCommit().commit();
	} finally {
		IOUtils.closeQuietly(stream);
	}

	// this should throw an exception as the file is already committed
	writer.recover(recoverable);
	fail();
}
 

@Test
public void testTaskManagerRpcServiceShouldBindToIpAddressDeterminedByConnectingToResourceManager() throws Exception {
	final ServerSocket testJobManagerSocket = openServerSocket();
	final Configuration config = createFlinkConfigWithJobManagerPort(testJobManagerSocket.getLocalPort());
	final HighAvailabilityServices highAvailabilityServices = createHighAvailabilityServices(config);

	RpcService taskManagerRpcService = null;
	try {
		taskManagerRpcService = TaskManagerRunner.createRpcService(config, highAvailabilityServices);
		assertThat(taskManagerRpcService.getAddress(), is(ipAddress()));
	} finally {
		maybeCloseRpcService(taskManagerRpcService);
		highAvailabilityServices.closeAndCleanupAllData();
		IOUtils.closeQuietly(testJobManagerSocket);
	}
}
 

private void writeSnapshotToOutputStream(
	@Nonnull CheckpointStreamWithResultProvider checkpointStreamWithResultProvider,
	@Nonnull KeyGroupRangeOffsets keyGroupRangeOffsets) throws IOException, InterruptedException {

	final List<Tuple2<RocksIteratorWrapper, Integer>> kvStateIterators =
		new ArrayList<>(metaData.size());
	final DataOutputView outputView =
		new DataOutputViewStreamWrapper(checkpointStreamWithResultProvider.getCheckpointOutputStream());
	final ReadOptions readOptions = new ReadOptions();
	try {
		readOptions.setSnapshot(snapshot);
		writeKVStateMetaData(kvStateIterators, readOptions, outputView);
		writeKVStateData(kvStateIterators, checkpointStreamWithResultProvider, keyGroupRangeOffsets);
	} finally {

		for (Tuple2<RocksIteratorWrapper, Integer> kvStateIterator : kvStateIterators) {
			IOUtils.closeQuietly(kvStateIterator.f0);
		}

		IOUtils.closeQuietly(readOptions);
	}
}
 

private void writeSnapshotToOutputStream(
	@Nonnull CheckpointStreamWithResultProvider checkpointStreamWithResultProvider,
	@Nonnull KeyGroupRangeOffsets keyGroupRangeOffsets) throws IOException, InterruptedException {

	final List<Tuple2<RocksIteratorWrapper, Integer>> kvStateIterators =
		new ArrayList<>(metaData.size());
	final DataOutputView outputView =
		new DataOutputViewStreamWrapper(checkpointStreamWithResultProvider.getCheckpointOutputStream());
	final ReadOptions readOptions = new ReadOptions();
	try {
		readOptions.setSnapshot(snapshot);
		writeKVStateMetaData(kvStateIterators, readOptions, outputView);
		writeKVStateData(kvStateIterators, checkpointStreamWithResultProvider, keyGroupRangeOffsets);
	} finally {

		for (Tuple2<RocksIteratorWrapper, Integer> kvStateIterator : kvStateIterators) {
			IOUtils.closeQuietly(kvStateIterator.f0);
		}

		IOUtils.closeQuietly(readOptions);
	}
}
 

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

@Override
public void releaseAllResources() throws IOException {
	// it is not a problem if this method executes multiple times
	isReleased = true;

	IOUtils.closeQuietly(dataReader);

	// nulling these fields means the read method and will fail fast
	nextBuffer = null;
	dataReader = null;

	// Notify the parent that this one is released. This allows the parent to
	// eventually release all resources (when all readers are done and the
	// parent is disposed).
	parent.releaseReaderReference(this);
}
 

@Test
public void testRegisterProvidedLocalResourcesWithDuplication() throws IOException {
	final File flinkLibDir1 = temporaryFolder.newFolder();
	final File flinkLibDir2 = temporaryFolder.newFolder();

	generateFilesInDirectory(flinkLibDir1, getLibJars());
	generateFilesInDirectory(flinkLibDir2, getLibJars());

	final FileSystem fileSystem = FileSystem.get(new YarnConfiguration());
	try {
		assertThrows(
			"Two files with the same filename exist in the shared libs",
			RuntimeException.class,
			() -> YarnApplicationFileUploader.from(
					fileSystem,
					new Path(temporaryFolder.getRoot().toURI()),
					Arrays.asList(new Path(flinkLibDir1.toURI()), new Path(flinkLibDir2.toURI())),
					ApplicationId.newInstance(0, 0),
					DFSConfigKeys.DFS_REPLICATION_DEFAULT));
	} finally {
		IOUtils.closeQuietly(fileSystem);
	}
}
 

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

@After
public void cleanupRocksDB() {
	if (keyedStateBackend != null) {
		IOUtils.closeQuietly(keyedStateBackend);
		keyedStateBackend.dispose();
	}
	IOUtils.closeQuietly(defaultCFHandle);
	IOUtils.closeQuietly(db);
	IOUtils.closeQuietly(optionsContainer);

	if (allCreatedCloseables != null) {
		for (RocksObject rocksCloseable : allCreatedCloseables) {
			verify(rocksCloseable, times(1)).close();
		}
		allCreatedCloseables = null;
	}
}
 

/**
 * 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 = createFlinkConfiguration();
		cfg.setInteger(NettyShuffleEnvironmentOptions.DATA_PORT, blocker.getLocalPort());
		cfg.setString(TaskManagerOptions.BIND_HOST, LOCAL_HOST);

		startTaskManager(
			cfg,
			rpcService,
			highAvailabilityServices);

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

@Test
public void testTaskManagerRpcServiceShouldBindToIpAddressDeterminedByConnectingToResourceManager() throws Exception {
	final ServerSocket testJobManagerSocket = openServerSocket();
	final Configuration config = createFlinkConfigWithJobManagerPort(testJobManagerSocket.getLocalPort());
	final HighAvailabilityServices highAvailabilityServices = createHighAvailabilityServices(config);

	RpcService taskManagerRpcService = null;
	try {
		taskManagerRpcService = TaskManagerRunner.createRpcService(config, highAvailabilityServices);
		assertThat(taskManagerRpcService.getAddress(), is(ipAddress()));
	} finally {
		maybeCloseRpcService(taskManagerRpcService);
		highAvailabilityServices.closeAndCleanupAllData();
		IOUtils.closeQuietly(testJobManagerSocket);
	}
}
 

@Test
public void testApplyToAllKeysLambdaFunction() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {
		ListStateDescriptor<String> listStateDescriptor =
			new ListStateDescriptor<>("foo", StringSerializer.INSTANCE);

		ListState<String> listState =
			backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor);

		for (int i = 0; i < 100; ++i) {
			backend.setCurrentKey(i);
			listState.add("Hello" + i);
		}

		// valid state value via applyToAllKeys().
		backend.applyToAllKeys(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor,
			(Integer key, ListState<String> state) -> assertEquals("Hello" + key, state.get().iterator().next())
		);
	}
	finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
}
 

@After
public void cleanupRocksDB() {
	if (keyedStateBackend != null) {
		IOUtils.closeQuietly(keyedStateBackend);
		keyedStateBackend.dispose();
	}
	IOUtils.closeQuietly(defaultCFHandle);
	IOUtils.closeQuietly(db);
	IOUtils.closeQuietly(columnOptions);
	IOUtils.closeQuietly(dbOptions);

	if (allCreatedCloseables != null) {
		for (RocksObject rocksCloseable : allCreatedCloseables) {
			verify(rocksCloseable, times(1)).close();
		}
		allCreatedCloseables = null;
	}
	try {
		org.apache.flink.util.FileUtils.deleteDirectory(instanceBasePath);
	} catch (Exception ex) {
		// ignored
	}
}
 

@Test
public void testCommitAfterNormalClose() throws Exception {
	final RecoverableWriter writer = getNewFileSystemWriter();

	final Path testDir = getBasePathForTest();

	final Path path = new Path(testDir, "part-0");

	RecoverableFsDataOutputStream stream = null;
	try {
		stream = writer.open(path);
		stream.write(testData1.getBytes(StandardCharsets.UTF_8));
		stream.closeForCommit().commit();

		for (Map.Entry<Path, String> fileContents : getFileContentByPath(testDir).entrySet()) {
			Assert.assertEquals("part-0", fileContents.getKey().getName());
			Assert.assertEquals(testData1, fileContents.getValue());
		}
	} finally {
		IOUtils.closeQuietly(stream);
	}
}
 

private void testLocalDbPaths(String configuredPath, File expectedPath) throws Exception {
	final RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI().toString());
	rocksDbBackend.setDbStoragePath(configuredPath);

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

	try {
		File instanceBasePath = keyedBackend.getInstanceBasePath();
		assertThat(instanceBasePath.getAbsolutePath(), startsWith(expectedPath.getAbsolutePath()));

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

@Test(expected = IOException.class)
public void testExceptionWritingAfterCloseForCommit() throws Exception {
	final Path testDir = getBasePathForTest();

	final RecoverableWriter writer = getNewFileSystemWriter();
	final Path path = new Path(testDir, "part-0");

	RecoverableFsDataOutputStream stream = null;
	try {
		stream = writer.open(path);
		stream.write(testData1.getBytes(StandardCharsets.UTF_8));

		stream.closeForCommit().getRecoverable();
		stream.write(testData2.getBytes(StandardCharsets.UTF_8));
		fail();
	} finally {
		IOUtils.closeQuietly(stream);
	}
}
 

/**
 * Reads the given archive file and returns a {@link Collection} of contained {@link ArchivedJson}.
 *
 * @param file archive to extract
 * @return collection of archived jsons
 * @throws IOException if the file can't be opened, read or doesn't contain valid json
 */
public static Collection<ArchivedJson> getArchivedJsons(Path file) throws IOException {
	try (FSDataInputStream input = file.getFileSystem().open(file);
		ByteArrayOutputStream output = new ByteArrayOutputStream()) {
		IOUtils.copyBytes(input, output);

		try {
			JsonNode archive = mapper.readTree(output.toByteArray());

			Collection<ArchivedJson> archives = new ArrayList<>();
			for (JsonNode archivePart : archive.get(ARCHIVE)) {
				String path = archivePart.get(PATH).asText();
				String json = archivePart.get(JSON).asText();
				archives.add(new ArchivedJson(path, json));
			}
			return archives;
		} catch (NullPointerException npe) {
			// occurs if the archive is empty or any of the expected fields are not present
			throw new IOException("Job archive (" + file.getPath() + ") did not conform to expected format.");
		}
	}
}
 

@Override
public TransportClient createClient(Map<String, String> clientConfig) {
	Settings settings = Settings.builder().put(clientConfig)
		.put(NetworkModule.HTTP_TYPE_KEY, Netty3Plugin.NETTY_HTTP_TRANSPORT_NAME)
		.put(NetworkModule.TRANSPORT_TYPE_KEY, Netty3Plugin.NETTY_TRANSPORT_NAME)
		.build();

	TransportClient transportClient = new PreBuiltTransportClient(settings);
	for (TransportAddress transport : ElasticsearchUtils.convertInetSocketAddresses(transportAddresses)) {
		transportClient.addTransportAddress(transport);
	}

	// verify that we actually are connected to a cluster
	if (transportClient.connectedNodes().isEmpty()) {

		// close the transportClient here
		IOUtils.closeQuietly(transportClient);

		throw new RuntimeException("Elasticsearch client is not connected to any Elasticsearch nodes!");
	}

	if (LOG.isInfoEnabled()) {
		LOG.info("Created Elasticsearch TransportClient with connected nodes {}", transportClient.connectedNodes());
	}

	return transportClient;
}
 

@Override
public DefaultOperatorStateBackend build() throws BackendBuildingException {
	Map<String, PartitionableListState<?>> registeredOperatorStates = new HashMap<>();
	Map<String, BackendWritableBroadcastState<?, ?>> registeredBroadcastStates = new HashMap<>();
	CloseableRegistry cancelStreamRegistryForBackend = new CloseableRegistry();
	AbstractSnapshotStrategy<OperatorStateHandle> snapshotStrategy =
		new DefaultOperatorStateBackendSnapshotStrategy(
			userClassloader,
			asynchronousSnapshots,
			registeredOperatorStates,
			registeredBroadcastStates,
			cancelStreamRegistryForBackend);
	OperatorStateRestoreOperation restoreOperation = new OperatorStateRestoreOperation(
		cancelStreamRegistry,
		userClassloader,
		registeredOperatorStates,
		registeredBroadcastStates,
		restoreStateHandles
	);
	try {
		restoreOperation.restore();
	} catch (Exception e) {
		IOUtils.closeQuietly(cancelStreamRegistryForBackend);
		throw new BackendBuildingException("Failed when trying to restore operator state backend", e);
	}
	return new DefaultOperatorStateBackend(
		executionConfig,
		cancelStreamRegistryForBackend,
		registeredOperatorStates,
		registeredBroadcastStates,
		new HashMap<>(),
		new HashMap<>(),
		snapshotStrategy
	);
}
 

@Override
public void commit() {
  try {
    flushToKeyedStateOutputStream();
  } catch (IOException e) {
    throw new RuntimeException(e);
  } finally {
    keyGroupStreams.clear();
    IOUtils.closeQuietly(snapshotLease);
    snapshotLease = null;
    keyedStateOutputStream = null;
  }
}
 

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

private static void unzipPythonLibrary(Path targetDir) throws IOException {
	FileSystem targetFs = targetDir.getFileSystem();
	ClassLoader classLoader = PythonPlanBinder.class.getClassLoader();
	try (ZipInputStream zis = new ZipInputStream(classLoader.getResourceAsStream("python-source.zip"))) {
		ZipEntry entry = zis.getNextEntry();
		while (entry != null) {
			String fileName = entry.getName();
			Path newFile = new Path(targetDir, fileName);
			if (entry.isDirectory()) {
				targetFs.mkdirs(newFile);
			} else {
				try {
					LOG.debug("Unzipping to {}.", newFile);
					FSDataOutputStream fsDataOutputStream = targetFs.create(newFile, FileSystem.WriteMode.NO_OVERWRITE);
					IOUtils.copyBytes(zis, fsDataOutputStream, false);
				} catch (Exception e) {
					zis.closeEntry();
					throw new IOException("Failed to unzip flink python library.", e);
				}
			}

			zis.closeEntry();
			entry = zis.getNextEntry();
		}
		zis.closeEntry();
	}
}
 

@Override
public void close() {
	List<ColumnFamilyOptions> columnFamilyOptions = new ArrayList<>(columnFamilyDescriptors.size() + 1);
	columnFamilyDescriptors.forEach((cfd) -> columnFamilyOptions.add(cfd.getOptions()));
	RocksDBOperationUtils.addColumnFamilyOptionsToCloseLater(columnFamilyOptions, defaultColumnFamilyHandle);
	IOUtils.closeQuietly(defaultColumnFamilyHandle);
	IOUtils.closeAllQuietly(columnFamilyHandles);
	IOUtils.closeQuietly(db);
	IOUtils.closeAllQuietly(columnFamilyOptions);
}
 

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

private void assertStateContentEqual(Path stateFilePath, FSDataInputStream inputStream) throws IOException {
	byte[] excepted = Files.readAllBytes(stateFilePath);
	byte[] actual = new byte[excepted.length];
	IOUtils.readFully(inputStream, actual, 0, actual.length);
	assertEquals(-1, inputStream.read());
	assertArrayEquals(excepted, actual);
}
 

private void startAllOperatorCoordinators() {
	final Collection<OperatorCoordinatorHolder> coordinators = getAllCoordinators();
	try {
		for (OperatorCoordinatorHolder coordinator : coordinators) {
			coordinator.start();
		}
	}
	catch (Throwable t) {
		ExceptionUtils.rethrowIfFatalErrorOrOOM(t);
		coordinators.forEach(IOUtils::closeQuietly);
		throw new FlinkRuntimeException("Failed to start the operator coordinators", t);
	}
}
 

@Override
public void close() throws RocksDBException {
	if (batch.count() != 0) {
		flush();
	}
	IOUtils.closeQuietly(batch);
}
 

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

	final MockEnvironment 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();
		env.close();
	}
}
 

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

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