Examples with MockListState org.apache.flink.streaming.api.functions.sink.filesystem.TestUtils.MockListState used on opensource projects

Example 1 with MockListState

use of org.apache.flink.streaming.api.functions.sink.filesystem.TestUtils.MockListState in project flink by apache.

the class BucketsTest method testBucketLifeCycleListenerOnRestoring.

@Test
public void testBucketLifeCycleListenerOnRestoring() throws Exception {
    File outDir = TEMP_FOLDER.newFolder();
    Path path = new Path(outDir.toURI());
    OnProcessingTimePolicy<String, String> rollOnProcessingTimeCountingPolicy = new OnProcessingTimePolicy<>(2L);
    RecordBucketLifeCycleListener bucketLifeCycleListener = new RecordBucketLifeCycleListener();
    Buckets<String, String> buckets = createBuckets(path, rollOnProcessingTimeCountingPolicy, bucketLifeCycleListener, null, 0, OutputFileConfig.builder().build());
    ListState<byte[]> bucketStateContainer = new MockListState<>();
    ListState<Long> partCounterContainer = new MockListState<>();
    buckets.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    buckets.onElement("test2", new TestUtils.MockSinkContext(null, 1L, 3L));
    // Will close the part file writer of the bucket "test1". Now bucket "test1" have only
    // one pending file while bucket "test2" has an on-writing in-progress file.
    buckets.onProcessingTime(4);
    buckets.snapshotState(0, bucketStateContainer, partCounterContainer);
    // On restoring the bucket "test1" will commit its pending file and become inactive.
    buckets = restoreBuckets(path, rollOnProcessingTimeCountingPolicy, bucketLifeCycleListener, null, 0, bucketStateContainer, partCounterContainer, OutputFileConfig.builder().build());
    Assert.assertEquals(new HashSet<>(Collections.singletonList("test2")), buckets.getActiveBuckets().keySet());
    List<Tuple2<RecordBucketLifeCycleListener.EventType, String>> expectedEvents = Arrays.asList(new Tuple2<>(RecordBucketLifeCycleListener.EventType.CREATED, "test1"), new Tuple2<>(RecordBucketLifeCycleListener.EventType.CREATED, "test2"), new Tuple2<>(RecordBucketLifeCycleListener.EventType.INACTIVE, "test1"));
    Assert.assertEquals(expectedEvents, bucketLifeCycleListener.getEvents());
}

Example 2 with MockListState

use of org.apache.flink.streaming.api.functions.sink.filesystem.TestUtils.MockListState in project flink by apache.

the class BucketsTest method testMergeAtScaleInAndMaxCounterAtRecovery.

@Test
public void testMergeAtScaleInAndMaxCounterAtRecovery() throws Exception {
    final File outDir = TEMP_FOLDER.newFolder();
    final Path path = new Path(outDir.toURI());
    final RollingPolicy<String, String> onCheckpointRP = DefaultRollingPolicy.builder().withMaxPartSize(// roll with 2 elements
    new MemorySize(7L)).build();
    final MockListState<byte[]> bucketStateContainerOne = new MockListState<>();
    final MockListState<byte[]> bucketStateContainerTwo = new MockListState<>();
    final MockListState<Long> partCounterContainerOne = new MockListState<>();
    final MockListState<Long> partCounterContainerTwo = new MockListState<>();
    final Buckets<String, String> bucketsOne = createBuckets(path, onCheckpointRP, 0);
    final Buckets<String, String> bucketsTwo = createBuckets(path, onCheckpointRP, 1);
    bucketsOne.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    bucketsOne.snapshotState(0L, bucketStateContainerOne, partCounterContainerOne);
    Assert.assertEquals(1L, bucketsOne.getMaxPartCounter());
    // make sure we have one in-progress file here
    Assert.assertNotNull(bucketsOne.getActiveBuckets().get("test1").getInProgressPart());
    // add a couple of in-progress files so that the part counter increases.
    bucketsTwo.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    bucketsTwo.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    bucketsTwo.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    bucketsTwo.snapshotState(0L, bucketStateContainerTwo, partCounterContainerTwo);
    Assert.assertEquals(2L, bucketsTwo.getMaxPartCounter());
    // make sure we have one in-progress file here and a pending
    Assert.assertEquals(1L, bucketsTwo.getActiveBuckets().get("test1").getPendingFileRecoverablesPerCheckpoint().size());
    Assert.assertNotNull(bucketsTwo.getActiveBuckets().get("test1").getInProgressPart());
    final ListState<byte[]> mergedBucketStateContainer = new MockListState<>();
    final ListState<Long> mergedPartCounterContainer = new MockListState<>();
    mergedBucketStateContainer.addAll(bucketStateContainerOne.getBackingList());
    mergedBucketStateContainer.addAll(bucketStateContainerTwo.getBackingList());
    mergedPartCounterContainer.addAll(partCounterContainerOne.getBackingList());
    mergedPartCounterContainer.addAll(partCounterContainerTwo.getBackingList());
    final Buckets<String, String> restoredBuckets = restoreBuckets(path, onCheckpointRP, 0, mergedBucketStateContainer, mergedPartCounterContainer);
    // we get the maximum of the previous tasks
    Assert.assertEquals(2L, restoredBuckets.getMaxPartCounter());
    final Map<String, Bucket<String, String>> activeBuckets = restoredBuckets.getActiveBuckets();
    Assert.assertEquals(1L, activeBuckets.size());
    Assert.assertTrue(activeBuckets.keySet().contains("test1"));
    final Bucket<String, String> bucket = activeBuckets.get("test1");
    Assert.assertEquals("test1", bucket.getBucketId());
    Assert.assertEquals(new Path(path, "test1"), bucket.getBucketPath());
    // the restored part file
    Assert.assertNotNull(bucket.getInProgressPart());
    // this is due to the Bucket#merge(). The in progress file of one
    // of the previous tasks is put in the list of pending files.
    Assert.assertEquals(1L, bucket.getPendingFileRecoverablesForCurrentCheckpoint().size());
    // we commit the pending for previous checkpoints
    Assert.assertTrue(bucket.getPendingFileRecoverablesPerCheckpoint().isEmpty());
}

Example 3 with MockListState

use of org.apache.flink.streaming.api.functions.sink.filesystem.TestUtils.MockListState in project flink by apache.

the class BucketsTest method testSnapshotAndRestore.

@Test
public void testSnapshotAndRestore() throws Exception {
    final File outDir = TEMP_FOLDER.newFolder();
    final Path path = new Path(outDir.toURI());
    final RollingPolicy<String, String> onCheckpointRollingPolicy = OnCheckpointRollingPolicy.build();
    final Buckets<String, String> buckets = createBuckets(path, onCheckpointRollingPolicy, 0);
    final ListState<byte[]> bucketStateContainer = new MockListState<>();
    final ListState<Long> partCounterContainer = new MockListState<>();
    buckets.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    buckets.snapshotState(0L, bucketStateContainer, partCounterContainer);
    assertThat(buckets.getActiveBuckets().get("test1"), hasSinglePartFileToBeCommittedOnCheckpointAck(path, "test1"));
    buckets.onElement("test2", new TestUtils.MockSinkContext(null, 1L, 2L));
    buckets.snapshotState(1L, bucketStateContainer, partCounterContainer);
    assertThat(buckets.getActiveBuckets().get("test1"), hasSinglePartFileToBeCommittedOnCheckpointAck(path, "test1"));
    assertThat(buckets.getActiveBuckets().get("test2"), hasSinglePartFileToBeCommittedOnCheckpointAck(path, "test2"));
    Buckets<String, String> restoredBuckets = restoreBuckets(path, onCheckpointRollingPolicy, 0, bucketStateContainer, partCounterContainer);
    final Map<String, Bucket<String, String>> activeBuckets = restoredBuckets.getActiveBuckets();
    // because we commit pending files for previous checkpoints upon recovery
    Assert.assertTrue(activeBuckets.isEmpty());
}

Example 4 with MockListState

use of org.apache.flink.streaming.api.functions.sink.filesystem.TestUtils.MockListState in project flink by apache.

the class BucketsTest method testBucketLifeCycleListenerOnCreatingAndInactive.

@Test
public void testBucketLifeCycleListenerOnCreatingAndInactive() throws Exception {
    File outDir = TEMP_FOLDER.newFolder();
    Path path = new Path(outDir.toURI());
    OnProcessingTimePolicy<String, String> rollOnProcessingTimeCountingPolicy = new OnProcessingTimePolicy<>(2L);
    RecordBucketLifeCycleListener bucketLifeCycleListener = new RecordBucketLifeCycleListener();
    Buckets<String, String> buckets = createBuckets(path, rollOnProcessingTimeCountingPolicy, bucketLifeCycleListener, null, 0, OutputFileConfig.builder().build());
    ListState<byte[]> bucketStateContainer = new MockListState<>();
    ListState<Long> partCounterContainer = new MockListState<>();
    buckets.onElement("test1", new TestUtils.MockSinkContext(null, 1L, 2L));
    buckets.onElement("test2", new TestUtils.MockSinkContext(null, 1L, 3L));
    // Will close the part file writer of the bucket "test1".
    buckets.onProcessingTime(4);
    buckets.snapshotState(0, bucketStateContainer, partCounterContainer);
    buckets.commitUpToCheckpoint(0);
    // Will close the part file writer of the bucket "test2".
    buckets.onProcessingTime(6);
    buckets.snapshotState(1, bucketStateContainer, partCounterContainer);
    buckets.commitUpToCheckpoint(1);
    List<Tuple2<RecordBucketLifeCycleListener.EventType, String>> expectedEvents = Arrays.asList(new Tuple2<>(RecordBucketLifeCycleListener.EventType.CREATED, "test1"), new Tuple2<>(RecordBucketLifeCycleListener.EventType.CREATED, "test2"), new Tuple2<>(RecordBucketLifeCycleListener.EventType.INACTIVE, "test1"), new Tuple2<>(RecordBucketLifeCycleListener.EventType.INACTIVE, "test2"));
    Assert.assertEquals(expectedEvents, bucketLifeCycleListener.getEvents());
}