Examples with SplitEnumerator org.apache.flink.api.connector.source.SplitEnumerator used on opensource projects

Example 1 with SplitEnumerator

use of org.apache.flink.api.connector.source.SplitEnumerator in project flink by apache.

the class SourceCoordinatorTest method testBlockOnClose.

@Test
public void testBlockOnClose() throws Exception {
    // It is possible that the split enumerator submits some heavy-duty work to the
    // coordinator executor which blocks the coordinator closure.
    final CountDownLatch latch = new CountDownLatch(1);
    try (final MockSplitEnumeratorContext<MockSourceSplit> enumeratorContext = new MockSplitEnumeratorContext<>(1);
        final MockSplitEnumerator splitEnumerator = new MockSplitEnumerator(1, enumeratorContext) {

            @Override
            public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
                context.callAsync(() -> 1L, (ignored, t) -> {
                    latch.countDown();
                    // Submit a callable that will never return.
                    try {
                        Thread.sleep(Long.MAX_VALUE);
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                });
            }
        };
        final SourceCoordinator<?, ?> coordinator = new SourceCoordinator<>(OPERATOR_NAME, new EnumeratorCreatingSource<>(() -> splitEnumerator), context, new CoordinatorStoreImpl())) {
        coordinator.start();
        coordinator.handleEventFromOperator(1, new SourceEventWrapper(new SourceEvent() {
        }));
        // Wait until the coordinator executor blocks.
        latch.await();
        CompletableFuture<?> future = ComponentClosingUtils.closeAsyncWithTimeout("testBlockOnClose", (ThrowingRunnable<Exception>) coordinator::close, Duration.ofMillis(1));
        future.exceptionally(e -> {
            assertTrue(e instanceof TimeoutException);
            return null;
        }).get();
        waitUtil(splitEnumerator::closed, Duration.ofSeconds(5), "Split enumerator was not closed in 5 seconds.");
    }
}