org.apache.flink.util.IOUtils Java Examples

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

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

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

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

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

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

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

@Override
public void cancel() {
	isRunning = false;

	// we need to close the socket as well, because the Thread.interrupt() function will
	// not wake the thread in the socketStream.read() method when blocked.
	Socket theSocket = this.currentSocket;
	if (theSocket != null) {
		IOUtils.closeSocket(theSocket);
	}
}
 

/**
 * Restore one key groups state handle.
 */
private void restoreKeyGroupsInStateHandle()
	throws IOException, StateMigrationException, RocksDBException {
	try {
		currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
		cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
		currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
		restoreKVStateMetaData();
		restoreKVStateData();
	} finally {
		if (cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
			IOUtils.closeQuietly(currentStateHandleInStream);
		}
	}
}
 

@Override
public void close() throws IOException {
	List<AsyncCheckpointRunnable> asyncCheckpointRunnables = null;
	synchronized (lock) {
		if (!closed) {
			closed = true;
			asyncCheckpointRunnables = new ArrayList<>(checkpoints.values());
			checkpoints.clear();
		}
	}
	IOUtils.closeAllQuietly(asyncCheckpointRunnables);
	channelStateWriter.close();
}
 

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

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

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

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

@Override
public void close() {
	if (closed.compareAndSet(false, true)) {
		if (delegate != null) {
			IOUtils.closeQuietly(delegate);
		}
		// trigger all the latches, essentially all blocking ops on the stream should resume after close.
		unblockAll();
	}
}
 

/**
 * Remove all the temp directories.
 */
@Override
public void close() throws Exception {
	IOUtils.closeAll(Arrays.stream(paths)
		.filter(File::exists)
		.map(FileChannelManagerImpl::getFileCloser)
		.collect(Collectors.toList()));
}
 

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

/**
 * Closes the state backend, releasing all internal resources, but does not delete any persistent
 * checkpoint data.
 *
 */
@Override
public void dispose() {

	IOUtils.closeQuietly(cancelStreamRegistry);

	if (kvStateRegistry != null) {
		kvStateRegistry.unregisterAll();
	}

	lastName = null;
	lastState = null;
	keyValueStatesByName.clear();
}
 

@Override
public boolean hasNext() {
	boolean currentHasNext = current.hasNext();

	// find the iterator of the next partition that has elements.
	while (!currentHasNext && index < keyGroupLists.length) {
		IOUtils.closeQuietly(current);
		current = keyGroupLists[index++].iterator();
		currentHasNext = current.hasNext();
	}
	return currentHasNext;
}
 

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