Skip to content

[FLINK-29775][tests][JUnit5 migration] Module: flink-statebackend-rocksdb #28943

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
spuru9 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[FLINK-29775][tests][JUnit5 migration] Module: flink-statebackend-rocksdb #28943

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions ...kend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/HeapTimersSnapshottingTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,32 +31,32 @@
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.util.Collector;

import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import static org.hamcrest.CoreMatchers.equalTo;
import static org.junit.Assert.assertThat;
import java.io.File;

import static org.assertj.core.api.Assertions.assertThat;

/**
* The tests verify that {@link PriorityQueueStateType#HEAP heap timers} are not serialized into raw
* keyed operator state when taking a savepoint, but they are serialized for checkpoints. The heap
* timers still need to be serialized into the raw operator state because of RocksDB incremental
* checkpoints.
*/
public class HeapTimersSnapshottingTest {
class HeapTimersSnapshottingTest {

@Rule public TemporaryFolder temporaryFolder = new TemporaryFolder();
@TempDir private File temporaryFolder;

@Test
public void testNotSerializingTimersInRawStateForSavepoints() throws Exception {
void testNotSerializingTimersInRawStateForSavepoints() throws Exception {
try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
getTestHarness()) {
EmbeddedRocksDBStateBackend backend = new EmbeddedRocksDBStateBackend();
backend.setPriorityQueueStateType(PriorityQueueStateType.HEAP);
testHarness.setStateBackend(backend);
testHarness.setCheckpointStorage(
new FileSystemCheckpointStorage(temporaryFolder.newFolder().toURI()));
new FileSystemCheckpointStorage(temporaryFolder.toURI()));
testHarness.open();
testHarness.processElement(0, 0L);

Expand All @@ -65,27 +65,27 @@ public void testNotSerializingTimersInRawStateForSavepoints() throws Exception {
.snapshotWithLocalState(
0L, 1L, SavepointType.savepoint(SavepointFormatType.CANONICAL))
.getJobManagerOwnedState();
assertThat(state.getRawKeyedState().isEmpty(), equalTo(true));
assertThat(state.getRawKeyedState()).isEmpty();
}
}

@Test
public void testSerializingTimersInRawStateForCheckpoints() throws Exception {
void testSerializingTimersInRawStateForCheckpoints() throws Exception {
try (KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness =
getTestHarness()) {
EmbeddedRocksDBStateBackend backend = new EmbeddedRocksDBStateBackend();
backend.setPriorityQueueStateType(PriorityQueueStateType.HEAP);
testHarness.setStateBackend(backend);
testHarness.setCheckpointStorage(
new FileSystemCheckpointStorage(temporaryFolder.newFolder().toURI()));
new FileSystemCheckpointStorage(temporaryFolder.toURI()));
testHarness.open();
testHarness.processElement(0, 0L);

OperatorSubtaskState state =
testHarness
.snapshotWithLocalState(0L, 1L, CheckpointType.CHECKPOINT)
.getJobManagerOwnedState();
assertThat(state.getRawKeyedState().isEmpty(), equalTo(false));
assertThat(state.getRawKeyedState()).isNotEmpty();
}
}

Expand Down
90 changes: 31 additions & 59 deletions ...ackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/RocksDBAsyncSnapshotTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,20 +71,15 @@
import org.apache.flink.streaming.runtime.tasks.OneInputStreamTaskTestHarness;
import org.apache.flink.streaming.runtime.tasks.StreamMockEnvironment;
import org.apache.flink.util.IOUtils;
import org.apache.flink.util.TestLogger;
import org.apache.flink.util.concurrent.FutureUtils;

import org.assertj.core.api.Assertions;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import javax.annotation.Nullable;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
Expand All @@ -101,15 +96,12 @@
import static org.apache.flink.runtime.state.FullSnapshotUtil.clearMetaDataFollowsFlag;
import static org.apache.flink.runtime.state.FullSnapshotUtil.hasMetaDataFollowsFlag;
import static org.apache.flink.runtime.state.FullSnapshotUtil.setMetaDataFollowsFlagInKey;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

/** Tests for asynchronous RocksDB Key/Value state checkpoints. */
@SuppressWarnings("serial")
public class RocksDBAsyncSnapshotTest extends TestLogger {

/** Temporary fold for test. */
@Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
class RocksDBAsyncSnapshotTest {

/**
* This ensures that asynchronous state handles are actually materialized asynchronously.
Expand All @@ -119,7 +111,7 @@ public class RocksDBAsyncSnapshotTest extends TestLogger {
* simply lock forever.
*/
@Test
public void testFullyAsyncSnapshot() throws Exception {
void testFullyAsyncSnapshot(@TempDir File dbDir) throws Exception {

final OneInputStreamTaskTestHarness<String, String> testHarness =
new OneInputStreamTaskTestHarness<>(
Expand All @@ -139,8 +131,6 @@ public String getKey(String value) throws Exception {

StreamConfig streamConfig = testHarness.getStreamConfig();

File dbDir = temporaryFolder.newFolder();

EmbeddedRocksDBStateBackend backend = new EmbeddedRocksDBStateBackend();
backend.setDbStoragePath(dbDir.getAbsolutePath());

Expand Down Expand Up @@ -181,7 +171,7 @@ public void acknowledgeCheckpoint(
}

// should be one k/v state
assertTrue(hasManagedKeyedState);
assertThat(hasManagedKeyedState).isTrue();

// we now know that the checkpoint went through
ensureCheckpointLatch.trigger();
Expand Down Expand Up @@ -254,20 +244,18 @@ public void reportInitializationMetrics(

ExecutorService threadPool = task.getAsyncOperationsThreadPool();
threadPool.shutdown();
Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
assertThat(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS)).isTrue();

testHarness.waitForTaskCompletion();
if (errorRef.get() != null) {
fail("Unexpected exception during execution.");
}
assertThat(errorRef.get()).isNull();
}

/**
* This tests ensures that canceling of asynchronous snapshots works as expected and does not
* block.
*/
@Test
public void testCancelFullyAsyncCheckpoints() throws Exception {
void testCancelFullyAsyncCheckpoints(@TempDir File dbDir) throws Exception {
final OneInputStreamTaskTestHarness<String, String> testHarness =
new OneInputStreamTaskTestHarness<>(
OneInputStreamTask::new,
Expand All @@ -280,8 +268,6 @@ public void testCancelFullyAsyncCheckpoints() throws Exception {

StreamConfig streamConfig = testHarness.getStreamConfig();

File dbDir = temporaryFolder.newFolder();

final EmbeddedRocksDBStateBackend.PriorityQueueStateType timerServicePriorityQueueType =
RocksDBOptions.TIMER_SERVICE_FACTORY.defaultValue();

Expand Down Expand Up @@ -371,38 +357,26 @@ public CheckpointStateOutputStream createCheckpointStateOutputStream(

ExecutorService threadPool = task.getAsyncOperationsThreadPool();
threadPool.shutdown();
Assert.assertTrue(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS));
assertThat(threadPool.awaitTermination(60_000, TimeUnit.MILLISECONDS)).isTrue();

Set<BlockingCheckpointOutputStream> createdStreams =
blockerCheckpointStreamFactory.getAllCreatedStreams();

for (BlockingCheckpointOutputStream stream : createdStreams) {
Assert.assertTrue(
"Not all of the "
+ createdStreams.size()
+ " created streams have been closed.",
stream.isClosed());
}
assertThat(createdStreams)
.as("Not all of the %d created streams have been closed.", createdStreams.size())
.allMatch(BlockingCheckpointOutputStream::isClosed);

try {
testHarness.waitForTaskCompletion();
fail("Operation completed. Cancel failed.");
} catch (Exception expected) {

Throwable cause = expected.getCause();

if (!(cause instanceof CancelTaskException)) {
fail("Unexpected exception: " + expected);
}
}
assertThatThrownBy(
testHarness::waitForTaskCompletion, "Operation completed. Cancel failed.")
.hasCauseInstanceOf(CancelTaskException.class);
}

/**
* Test that the snapshot files are cleaned up in case of a failure during the snapshot
* procedure.
*/
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
void testCleanupOfSnapshotsInFailureCase(@TempDir File dbDir) throws Exception {
long checkpointId = 1L;
long timestamp = 42L;

Expand All @@ -413,7 +387,7 @@ public void testCleanupOfSnapshotsInFailureCase() throws Exception {

EmbeddedRocksDBStateBackend backend = new EmbeddedRocksDBStateBackend();

backend.setDbStoragePath(temporaryFolder.newFolder().toURI().toString());
backend.setDbStoragePath(dbDir.toURI().toString());

CheckpointableKeyedStateBackend<Void> keyedStateBackend =
backend.createKeyedStateBackend(
Expand Down Expand Up @@ -446,14 +420,12 @@ public void testCleanupOfSnapshotsInFailureCase() throws Exception {
new TestCheckpointStreamFactory(() -> outputStream),
CheckpointOptions.forCheckpointWithDefaultLocation());

try {
FutureUtils.runIfNotDoneAndGet(snapshotFuture);
fail("Expected an exception to be thrown here.");
} catch (ExecutionException e) {
Assert.assertEquals(testException, e.getCause());
}
assertThatThrownBy(() -> FutureUtils.runIfNotDoneAndGet(snapshotFuture))
.isInstanceOf(ExecutionException.class)
.cause()
.isSameAs(testException);

Assertions.assertThat(outputStream.isCloseCalled()).isEqualTo(true);
assertThat(outputStream.isCloseCalled()).isTrue();
} finally {
IOUtils.closeQuietly(keyedStateBackend);
keyedStateBackend.dispose();
Expand All @@ -462,23 +434,23 @@ public void testCleanupOfSnapshotsInFailureCase() throws Exception {
}

@Test
public void testConsistentSnapshotSerializationFlagsAndMasks() {
void testConsistentSnapshotSerializationFlagsAndMasks() {

Assert.assertEquals(0xFFFF, END_OF_KEY_GROUP_MARK);
Assert.assertEquals(0x80, FIRST_BIT_IN_BYTE_MASK);
assertThat(END_OF_KEY_GROUP_MARK).isEqualTo(0xFFFF);
assertThat(FIRST_BIT_IN_BYTE_MASK).isEqualTo(0x80);

byte[] expectedKey = new byte[] {42, 42};
byte[] modKey = expectedKey.clone();

Assert.assertFalse(hasMetaDataFollowsFlag(modKey));
assertThat(hasMetaDataFollowsFlag(modKey)).isFalse();

setMetaDataFollowsFlagInKey(modKey);
Assert.assertTrue(hasMetaDataFollowsFlag(modKey));
assertThat(hasMetaDataFollowsFlag(modKey)).isTrue();

clearMetaDataFollowsFlag(modKey);
Assert.assertFalse(hasMetaDataFollowsFlag(modKey));
assertThat(hasMetaDataFollowsFlag(modKey)).isFalse();

Assert.assertTrue(Arrays.equals(expectedKey, modKey));
assertThat(modKey).isEqualTo(expectedKey);
}

// ------------------------------------------------------------------------
Expand Down
16 changes: 8 additions & 8 deletions ...k-statebackend-rocksdb/src/test/java/org/apache/flink/state/rocksdb/RocksDBExtension.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,13 @@

package org.apache.flink.state.rocksdb;

import org.apache.flink.util.FileUtils;
import org.apache.flink.util.FlinkRuntimeException;
import org.apache.flink.util.IOUtils;

import org.junit.jupiter.api.extension.AfterEachCallback;
import org.junit.jupiter.api.extension.BeforeEachCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.junit.rules.TemporaryFolder;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
Expand All @@ -39,7 +39,8 @@

import javax.annotation.Nonnull;

import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
Expand All @@ -55,7 +56,7 @@ public class RocksDBExtension implements BeforeEachCallback, AfterEachCallback {
private final boolean enableStatistics;

/** Temporary folder that provides the working directory for the RocksDB instance. */
private TemporaryFolder temporaryFolder;
private Path temporaryFolder;

/** The options for the RocksDB instance. */
private DBOptions dbOptions;
Expand Down Expand Up @@ -169,9 +170,8 @@ public ColumnFamilyHandle createNewColumnFamily(String name) {
}

public void before() throws Exception {
this.temporaryFolder = new TemporaryFolder();
this.temporaryFolder.create();
final File rocksFolder = temporaryFolder.newFolder();
this.temporaryFolder = Files.createTempDirectory("rocksdb-extension");
final Path rocksFolder = Files.createDirectory(temporaryFolder.resolve("db"));
this.dbOptions =
optionsFactory
.createDBOptions(
Expand All @@ -195,7 +195,7 @@ public void before() throws Exception {
this.rocksDB =
RocksDB.open(
dbOptions,
rocksFolder.getAbsolutePath(),
rocksFolder.toAbsolutePath().toString(),
Collections.singletonList(
new ColumnFamilyDescriptor(
"default".getBytes(), columnFamilyOptions)),
Expand All @@ -215,7 +215,7 @@ public void after() throws Exception {
IOUtils.closeQuietly(this.columnFamilyOptions);
IOUtils.closeQuietly(this.dbOptions);
handlesToClose.forEach(IOUtils::closeQuietly);
temporaryFolder.delete();
FileUtils.deleteDirectoryQuietly(temporaryFolder.toFile());
}

@Override
Expand Down
Loading