Examples with StateObjectCollection org.apache.flink.runtime.checkpoint.StateObjectCollection used on opensource projects

Example 11 with StateObjectCollection

use of org.apache.flink.runtime.checkpoint.StateObjectCollection in project flink by apache.

the class BackendRestorerProcedureTest method testCanBeCanceledViaRegistry.

/**
 * Test that the restore can be stopped via the provided closeable registry.
 */
@Test
public void testCanBeCanceledViaRegistry() throws Exception {
    CloseableRegistry closeableRegistry = new CloseableRegistry();
    OneShotLatch waitForBlock = new OneShotLatch();
    OneShotLatch unblock = new OneShotLatch();
    OperatorStateHandle blockingRestoreHandle = mock(OperatorStateHandle.class);
    when(blockingRestoreHandle.openInputStream()).thenReturn(new BlockingFSDataInputStream(waitForBlock, unblock));
    List<StateObjectCollection<OperatorStateHandle>> sortedRestoreOptions = Collections.singletonList(new StateObjectCollection<>(Collections.singletonList(blockingRestoreHandle)));
    BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> restorerProcedure = new BackendRestorerProcedure<>(backendSupplier, closeableRegistry, "test op state backend");
    AtomicReference<Exception> exceptionReference = new AtomicReference<>(null);
    Thread restoreThread = new Thread(() -> {
        try {
            restorerProcedure.createAndRestore(sortedRestoreOptions);
        } catch (Exception e) {
            exceptionReference.set(e);
        }
    });
    restoreThread.start();
    waitForBlock.await();
    closeableRegistry.close();
    unblock.trigger();
    restoreThread.join();
    Exception exception = exceptionReference.get();
    Assert.assertTrue(exception instanceof FlinkException);
}