Examples with ThrowingRunnable org.apache.beam.sdk.util.ThrowingRunnable used on opensource projects

Example 1 with ThrowingRunnable

use of org.apache.beam.sdk.util.ThrowingRunnable in project beam by apache.

the class RegisterAndProcessBundleOperationTest method testAbortCancelsAndCleansUpDuringProcessBundle.

@Test
public void testAbortCancelsAndCleansUpDuringProcessBundle() throws Exception {
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    ExecutorService executorService = Executors.newCachedThreadPool();
    CountDownLatch waitForAbortToComplete = new CountDownLatch(1);
    AtomicReference<ThrowingRunnable> abortReference = new AtomicReference<>();
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>();
            if (request.getRequestCase() == RequestCase.PROCESS_BUNDLE) {
                executorService.submit((Callable<Void>) () -> {
                    abortReference.get().run();
                    waitForAbortToComplete.countDown();
                    return null;
                });
            } else {
                completeFuture(request, responseFuture);
            }
            return responseFuture;
        }
    }, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
    abortReference.set(operation::abort);
    operation.start();
    waitForAbortToComplete.await();
    // Ensure that the number of registrations matches the number of aborts
    assertEquals(stateServiceRegisterCounter.get(), stateServiceAbortCounter.get());
    assertEquals(0, stateServiceDeregisterCounter.get());
}

Example 2 with ThrowingRunnable

use of org.apache.beam.sdk.util.ThrowingRunnable in project beam by apache.

the class RegisterAndProcessBundleOperationTest method testAbortCancelsAndCleansUpDuringRegister.

@Test
public void testAbortCancelsAndCleansUpDuringRegister() throws Exception {
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    ExecutorService executorService = Executors.newCachedThreadPool();
    CountDownLatch waitForAbortToComplete = new CountDownLatch(1);
    AtomicReference<ThrowingRunnable> abortReference = new AtomicReference<>();
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>();
            if (request.getRequestCase() == RequestCase.PROCESS_BUNDLE) {
                executorService.submit((Callable<Void>) () -> {
                    abortReference.get().run();
                    waitForAbortToComplete.countDown();
                    return null;
                });
            } else {
                completeFuture(request, responseFuture);
            }
            return responseFuture;
        }
    }, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
    abortReference.set(operation::abort);
    operation.start();
    waitForAbortToComplete.await();
    // Ensure that the number of registrations matches the number of aborts
    assertEquals(stateServiceRegisterCounter.get(), stateServiceAbortCounter.get());
    assertEquals(0, stateServiceDeregisterCounter.get());
}