use of org.apache.flink.runtime.util.BlockingFSDataInputStream 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);
}
Aggregations