Java Code Examples for org.apache.flink.runtime.state.CheckpointStreamFactory#CheckpointStateOutputStream

/**
 * Test that the exception arose in the thread pool will rethrow to the main thread.
 */
@Test
public void testMultiThreadUploadThreadPoolExceptionRethrow() throws IOException {
	SpecifiedException expectedException = new SpecifiedException("throw exception while multi thread upload states.");

	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = createFailingCheckpointStateOutputStream(expectedException);
	CheckpointStreamFactory checkpointStreamFactory = (CheckpointedStateScope scope) -> outputStream;

	File file = temporaryFolder.newFile(String.valueOf(UUID.randomUUID()));
	generateRandomFileContent(file.getPath(), 20);

	Map<StateHandleID, Path> filePaths = new HashMap<>(1);
	filePaths.put(new StateHandleID("mockHandleID"), new Path(file.getPath()));
	try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
		rocksDBStateUploader.uploadFilesToCheckpointFs(filePaths, checkpointStreamFactory, new CloseableRegistry());
		fail();
	} catch (Exception e) {
		assertEquals(expectedException, e);
	}
}
 

/**
 * Test that the exception arose in the thread pool will rethrow to the main thread.
 */
@Test
public void testMultiThreadUploadThreadPoolExceptionRethrow() throws IOException {
	SpecifiedException expectedException = new SpecifiedException("throw exception while multi thread upload states.");

	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = createFailingCheckpointStateOutputStream(expectedException);
	CheckpointStreamFactory checkpointStreamFactory = (CheckpointedStateScope scope) -> outputStream;

	File file = temporaryFolder.newFile(String.valueOf(UUID.randomUUID()));
	generateRandomFileContent(file.getPath(), 20);

	Map<StateHandleID, Path> filePaths = new HashMap<>(1);
	filePaths.put(new StateHandleID("mockHandleID"), file.toPath());
	try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
		rocksDBStateUploader.uploadFilesToCheckpointFs(filePaths, checkpointStreamFactory, new CloseableRegistry());
		fail();
	} catch (Exception e) {
		assertEquals(expectedException, e);
	}
}
 

/**
 * Tests that the underlying stream file is deleted upon calling close.
 */
@Test
public void testCleanupWhenClosingStream() throws IOException {

	final FileSystem fs = mock(FileSystem.class);
	final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);

	final ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);

	when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);

	CheckpointStreamFactory.CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
		Path.fromLocalFile(tempDir.newFolder()),
		fs,
		4,
		0);

	// this should create the underlying file stream
	stream.write(new byte[] {1, 2, 3, 4, 5});

	verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));

	stream.close();

	verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}
 

@Nullable
@Override
public StreamStateHandle closeAndGetHandle() throws IOException {

	if (!closed.compareAndSet(false, true)) {
		throw new IOException("Stream was already closed!");
	}

	if (delegate instanceof CheckpointStreamFactory.CheckpointStateOutputStream) {
		StreamStateHandle streamStateHandle =
			((CheckpointStreamFactory.CheckpointStateOutputStream) delegate).closeAndGetHandle();
		unblockAll();
		return streamStateHandle;
	} else {
		unblockAll();
		throw new IOException("Delegate is not a CheckpointStateOutputStream!");
	}
}
 

@Nullable
@Override
public StreamStateHandle closeAndGetHandle() throws IOException {

	if (!closed.compareAndSet(false, true)) {
		throw new IOException("Stream was already closed!");
	}

	if (delegate instanceof CheckpointStreamFactory.CheckpointStateOutputStream) {
		StreamStateHandle streamStateHandle =
			((CheckpointStreamFactory.CheckpointStateOutputStream) delegate).closeAndGetHandle();
		unblockAll();
		return streamStateHandle;
	} else {
		unblockAll();
		throw new IOException("Delegate is not a CheckpointStateOutputStream!");
	}
}
 

/**
 * Tests that the underlying stream file is deleted upon calling close.
 */
@Test
public void testCleanupWhenClosingStream() throws IOException {

	final FileSystem fs = mock(FileSystem.class);
	final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);

	final ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);

	when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);

	CheckpointStreamFactory.CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
		Path.fromLocalFile(tempDir.newFolder()),
		fs,
		4,
		0);

	// this should create the underlying file stream
	stream.write(new byte[] {1, 2, 3, 4, 5});

	verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));

	stream.close();

	verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}
 

/**
 * Test that the exception arose in the thread pool will rethrow to the main thread.
 */
@Test
public void testMultiThreadUploadThreadPoolExceptionRethrow() throws IOException {
	SpecifiedException expectedException = new SpecifiedException("throw exception while multi thread upload states.");

	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = createFailingCheckpointStateOutputStream(expectedException);
	CheckpointStreamFactory checkpointStreamFactory = (CheckpointedStateScope scope) -> outputStream;

	File file = temporaryFolder.newFile(String.valueOf(UUID.randomUUID()));
	generateRandomFileContent(file.getPath(), 20);

	Map<StateHandleID, Path> filePaths = new HashMap<>(1);
	filePaths.put(new StateHandleID("mockHandleID"), new Path(file.getPath()));
	try (RocksDBStateUploader rocksDBStateUploader = new RocksDBStateUploader(5)) {
		rocksDBStateUploader.uploadFilesToCheckpointFs(filePaths, checkpointStreamFactory, new CloseableRegistry());
		fail();
	} catch (Exception e) {
		assertEquals(expectedException, e);
	}
}
 

private CheckpointStreamFactory.CheckpointStateOutputStream createFailingCheckpointStateOutputStream(
	IOException failureException) {
	return new CheckpointStreamFactory.CheckpointStateOutputStream() {
		@Nullable
		@Override
		public StreamStateHandle closeAndGetHandle() {
			return new ByteStreamStateHandle("testHandle", "testHandle".getBytes());
		}

		@Override
		public void close() {
		}

		@Override
		public long getPos() {
			return 0;
		}

		@Override
		public void flush() {
		}

		@Override
		public void sync() {
		}

		@Override
		public void write(int b) throws IOException {
			throw failureException;
		}
	};
}
 

/**
 * Tests that the underlying stream file is deleted if the closeAndGetHandle method fails.
 */
@Test
public void testCleanupWhenFailingCloseAndGetHandle() throws IOException {
	final FileSystem fs = mock(FileSystem.class);
	final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);

	final ArgumentCaptor<Path>  pathCaptor = ArgumentCaptor.forClass(Path.class);

	when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);
	doThrow(new IOException("Test IOException.")).when(outputStream).close();

	CheckpointStreamFactory.CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
		Path.fromLocalFile(tempDir.newFolder()),
		fs,
		4,
		0);

	// this should create the underlying file stream
	stream.write(new byte[] {1, 2, 3, 4, 5});

	verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));

	try {
		stream.closeAndGetHandle();
		fail("Expected IOException");
	} catch (IOException ioE) {
		// expected exception
	}

	verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}
 

private CheckpointStreamFactory.CheckpointStateOutputStream createFailingCheckpointStateOutputStream(
	IOException failureException) {
	return new CheckpointStreamFactory.CheckpointStateOutputStream() {
		@Nullable
		@Override
		public StreamStateHandle closeAndGetHandle() {
			return new ByteStreamStateHandle("testHandle", "testHandle".getBytes());
		}

		@Override
		public void close() {
		}

		@Override
		public long getPos() {
			return 0;
		}

		@Override
		public void flush() {
		}

		@Override
		public void sync() {
		}

		@Override
		public void write(int b) throws IOException {
			throw failureException;
		}
	};
}
 

/**
 * Tests that the underlying stream file is deleted upon calling close.
 */
@Test
public void testCleanupWhenClosingStream() throws IOException {

	final FileSystem fs = mock(FileSystem.class);
	final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);

	final ArgumentCaptor<Path> pathCaptor = ArgumentCaptor.forClass(Path.class);

	when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);

	CheckpointStreamFactory.CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
		Path.fromLocalFile(tempDir.newFolder()),
		fs,
		4,
		0,
		relativePaths);

	// this should create the underlying file stream
	stream.write(new byte[] {1, 2, 3, 4, 5});

	verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));

	stream.close();

	verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}
 

@Test
public void testStreamClosingExceptionally() throws Exception {
	long checkpointId = 42L;
	long checkpointTimestamp = 1L;

	CheckpointStreamFactory.CheckpointStateOutputStream outputStream1 = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream2 = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);

	CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
	when(streamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE)).thenReturn(outputStream1, outputStream2);

	InsightCloseableRegistry closableRegistry = new InsightCloseableRegistry();

	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);

	StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl(
		checkpointId,
		checkpointTimestamp,
		streamFactory,
		keyGroupRange,
		closableRegistry);

	// creating the output streams
	context.getRawKeyedOperatorStateOutput();
	context.getRawOperatorStateOutput();

	verify(streamFactory, times(2)).createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);

	assertEquals(2, closableRegistry.size());
	assertTrue(closableRegistry.contains(outputStream1));
	assertTrue(closableRegistry.contains(outputStream2));

	context.closeExceptionally();

	verify(outputStream1).close();
	verify(outputStream2).close();

	assertEquals(0, closableRegistry.size());
}
 

/**
 * Tests that the underlying stream file is deleted if the closeAndGetHandle method fails.
 */
@Test
public void testCleanupWhenFailingCloseAndGetHandle() throws IOException {
	final FileSystem fs = mock(FileSystem.class);
	final FSDataOutputStream outputStream = mock(FSDataOutputStream.class);

	final ArgumentCaptor<Path>  pathCaptor = ArgumentCaptor.forClass(Path.class);

	when(fs.create(pathCaptor.capture(), any(FileSystem.WriteMode.class))).thenReturn(outputStream);
	doThrow(new IOException("Test IOException.")).when(outputStream).close();

	CheckpointStreamFactory.CheckpointStateOutputStream stream = new FsCheckpointStreamFactory.FsCheckpointStateOutputStream(
		Path.fromLocalFile(tempDir.newFolder()),
		fs,
		4,
		0);

	// this should create the underlying file stream
	stream.write(new byte[] {1, 2, 3, 4, 5});

	verify(fs).create(any(Path.class), any(FileSystem.WriteMode.class));

	try {
		stream.closeAndGetHandle();
		fail("Expected IOException");
	} catch (IOException ioE) {
		// expected exception
	}

	verify(fs).delete(eq(pathCaptor.getValue()), anyBoolean());
}
 

private StreamStateHandle uploadLocalFileToCheckpointFs(
	Path filePath,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws IOException {
	FSDataInputStream inputStream = null;
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;

	try {
		final byte[] buffer = new byte[READ_BUFFER_SIZE];

		FileSystem backupFileSystem = filePath.getFileSystem();
		inputStream = backupFileSystem.open(filePath);
		closeableRegistry.registerCloseable(inputStream);

		outputStream = checkpointStreamFactory
			.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
		closeableRegistry.registerCloseable(outputStream);

		while (true) {
			int numBytes = inputStream.read(buffer);

			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}

		StreamStateHandle result = null;
		if (closeableRegistry.unregisterCloseable(outputStream)) {
			result = outputStream.closeAndGetHandle();
			outputStream = null;
		}
		return result;

	} finally {

		if (closeableRegistry.unregisterCloseable(inputStream)) {
			IOUtils.closeQuietly(inputStream);
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			IOUtils.closeQuietly(outputStream);
		}
	}
}
 

@Override
public CheckpointStreamFactory.CheckpointStateOutputStream createTaskOwnedStateStream() throws IOException {
	return delegate.createTaskOwnedStateStream();
}
 

@Test
public void testMixedBelowAndAboveThreshold() throws Exception {
	final byte[] state1 = new byte[1274673];
	final byte[] state2 = new byte[1];
	final byte[] state3 = new byte[0];
	final byte[] state4 = new byte[177];

	final Random rnd = new Random();
	rnd.nextBytes(state1);
	rnd.nextBytes(state2);
	rnd.nextBytes(state3);
	rnd.nextBytes(state4);

	final File directory = tempDir.newFolder();
	final Path basePath = Path.fromLocalFile(directory);

	final Supplier<CheckpointStateOutputStream> factory = () ->
			new FsCheckpointStateOutputStream(basePath, FileSystem.getLocalFileSystem(), 1024, 15);

	CheckpointStateOutputStream stream1 = factory.get();
	CheckpointStateOutputStream stream2 = factory.get();
	CheckpointStateOutputStream stream3 = factory.get();

	stream1.write(state1);
	stream2.write(state2);
	stream3.write(state3);

	FileStateHandle handle1 = (FileStateHandle) stream1.closeAndGetHandle();
	ByteStreamStateHandle handle2 = (ByteStreamStateHandle) stream2.closeAndGetHandle();
	ByteStreamStateHandle handle3 = (ByteStreamStateHandle) stream3.closeAndGetHandle();

	// use with try-with-resources
	StreamStateHandle handle4;
	try (CheckpointStreamFactory.CheckpointStateOutputStream stream4 = factory.get()) {
		stream4.write(state4);
		handle4 = stream4.closeAndGetHandle();
	}

	// close before accessing handle
	CheckpointStreamFactory.CheckpointStateOutputStream stream5 = factory.get();
	stream5.write(state4);
	stream5.close();
	try {
		stream5.closeAndGetHandle();
		fail();
	} catch (IOException e) {
		// uh-huh
	}

	validateBytesInStream(handle1.openInputStream(), state1);
	handle1.discardState();
	assertFalse(isDirectoryEmpty(directory));
	ensureLocalFileDeleted(handle1.getFilePath());

	validateBytesInStream(handle2.openInputStream(), state2);
	handle2.discardState();
	assertFalse(isDirectoryEmpty(directory));

	// nothing was written to the stream, so it will return nothing
	assertNull(handle3);
	assertFalse(isDirectoryEmpty(directory));

	validateBytesInStream(handle4.openInputStream(), state4);
	handle4.discardState();
	assertTrue(isDirectoryEmpty(directory));
}
 

/**
 * Tests that closing the StateSnapshotContextSynchronousImpl will also close the associated
 * output streams.
 */
@Test
public void testStreamClosingWhenClosing() throws Exception {
	long checkpointId = 42L;
	long checkpointTimestamp = 1L;

	CheckpointStreamFactory.CheckpointStateOutputStream outputStream1 = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream2 = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);

	CheckpointStreamFactory streamFactory = mock(CheckpointStreamFactory.class);
	when(streamFactory.createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE)).thenReturn(outputStream1, outputStream2);

	InsightCloseableRegistry closableRegistry = new InsightCloseableRegistry();

	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);

	StateSnapshotContextSynchronousImpl context = new StateSnapshotContextSynchronousImpl(
		checkpointId,
		checkpointTimestamp,
		streamFactory,
		keyGroupRange,
		closableRegistry);

	// creating the output streams
	context.getRawKeyedOperatorStateOutput();
	context.getRawOperatorStateOutput();

	verify(streamFactory, times(2)).createCheckpointStateOutputStream(CheckpointedStateScope.EXCLUSIVE);

	assertEquals(2, closableRegistry.size());
	assertTrue(closableRegistry.contains(outputStream1));
	assertTrue(closableRegistry.contains(outputStream2));

	context.getKeyedStateStreamFuture().run();
	context.getOperatorStateStreamFuture().run();

	verify(outputStream1).closeAndGetHandle();
	verify(outputStream2).closeAndGetHandle();

	assertEquals(0, closableRegistry.size());
}
 

private StreamStateHandle uploadLocalFileToCheckpointFs(
	Path filePath,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws IOException {

	InputStream inputStream = null;
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;

	try {
		final byte[] buffer = new byte[READ_BUFFER_SIZE];

		inputStream = Files.newInputStream(filePath);
		closeableRegistry.registerCloseable(inputStream);

		outputStream = checkpointStreamFactory
			.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
		closeableRegistry.registerCloseable(outputStream);

		while (true) {
			int numBytes = inputStream.read(buffer);

			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}

		StreamStateHandle result = null;
		if (closeableRegistry.unregisterCloseable(outputStream)) {
			result = outputStream.closeAndGetHandle();
			outputStream = null;
		}
		return result;

	} finally {

		if (closeableRegistry.unregisterCloseable(inputStream)) {
			IOUtils.closeQuietly(inputStream);
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			IOUtils.closeQuietly(outputStream);
		}
	}
}
 

@Test
public void testMixedBelowAndAboveThreshold() throws Exception {
	final byte[] state1 = new byte[1274673];
	final byte[] state2 = new byte[1];
	final byte[] state3 = new byte[0];
	final byte[] state4 = new byte[177];

	final Random rnd = new Random();
	rnd.nextBytes(state1);
	rnd.nextBytes(state2);
	rnd.nextBytes(state3);
	rnd.nextBytes(state4);

	final File directory = tempDir.newFolder();
	final Path basePath = Path.fromLocalFile(directory);

	final Supplier<CheckpointStateOutputStream> factory = () ->
			new FsCheckpointStateOutputStream(basePath, FileSystem.getLocalFileSystem(), 1024, 15, relativePaths);

	CheckpointStateOutputStream stream1 = factory.get();
	CheckpointStateOutputStream stream2 = factory.get();
	CheckpointStateOutputStream stream3 = factory.get();

	stream1.write(state1);
	stream2.write(state2);
	stream3.write(state3);

	FileStateHandle handle1 = (FileStateHandle) stream1.closeAndGetHandle();
	ByteStreamStateHandle handle2 = (ByteStreamStateHandle) stream2.closeAndGetHandle();
	ByteStreamStateHandle handle3 = (ByteStreamStateHandle) stream3.closeAndGetHandle();

	// use with try-with-resources
	StreamStateHandle handle4;
	try (CheckpointStreamFactory.CheckpointStateOutputStream stream4 = factory.get()) {
		stream4.write(state4);
		handle4 = stream4.closeAndGetHandle();
	}

	// close before accessing handle
	CheckpointStreamFactory.CheckpointStateOutputStream stream5 = factory.get();
	stream5.write(state4);
	stream5.close();
	try {
		stream5.closeAndGetHandle();
		fail();
	} catch (IOException e) {
		// uh-huh
	}

	validateBytesInStream(handle1.openInputStream(), state1);
	handle1.discardState();
	assertFalse(isDirectoryEmpty(directory));
	ensureLocalFileDeleted(handle1.getFilePath());

	validateBytesInStream(handle2.openInputStream(), state2);
	handle2.discardState();
	assertFalse(isDirectoryEmpty(directory));

	// nothing was written to the stream, so it will return nothing
	assertNull(handle3);
	assertFalse(isDirectoryEmpty(directory));

	validateBytesInStream(handle4.openInputStream(), state4);
	handle4.discardState();
	assertTrue(isDirectoryEmpty(directory));
}
 

private StreamStateHandle uploadLocalFileToCheckpointFs(
	Path filePath,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws IOException {
	FSDataInputStream inputStream = null;
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;

	try {
		final byte[] buffer = new byte[READ_BUFFER_SIZE];

		FileSystem backupFileSystem = filePath.getFileSystem();
		inputStream = backupFileSystem.open(filePath);
		closeableRegistry.registerCloseable(inputStream);

		outputStream = checkpointStreamFactory
			.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
		closeableRegistry.registerCloseable(outputStream);

		while (true) {
			int numBytes = inputStream.read(buffer);

			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}

		StreamStateHandle result = null;
		if (closeableRegistry.unregisterCloseable(outputStream)) {
			result = outputStream.closeAndGetHandle();
			outputStream = null;
		}
		return result;

	} finally {

		if (closeableRegistry.unregisterCloseable(inputStream)) {
			IOUtils.closeQuietly(inputStream);
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			IOUtils.closeQuietly(outputStream);
		}
	}
}