Examples with InstructionResponse org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse used on opensource projects

Example 16 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse in project beam by apache.

the class FnApiWindowMappingFn method loadIfNeeded.

private TargetWindowT loadIfNeeded(BoundedWindow mainWindow) {
    try {
        String processRequestInstructionId = idGenerator.getId();
        InstructionRequest processRequest = InstructionRequest.newBuilder().setInstructionId(processRequestInstructionId).setProcessBundle(ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId(registerIfRequired())).build();
        ConcurrentLinkedQueue<WindowedValue<KV<byte[], TargetWindowT>>> outputValue = new ConcurrentLinkedQueue<>();
        // Open the inbound consumer
        InboundDataClient waitForInboundTermination = beamFnDataService.receive(LogicalEndpoint.data(processRequestInstructionId, "write"), inboundCoder, outputValue::add);
        CompletionStage<InstructionResponse> processResponse = instructionRequestHandler.handle(processRequest);
        // Open the outbound consumer
        try (CloseableFnDataReceiver<WindowedValue<KV<byte[], BoundedWindow>>> outboundConsumer = beamFnDataService.send(LogicalEndpoint.data(processRequestInstructionId, "read"), outboundCoder)) {
            outboundConsumer.accept(WindowedValue.valueInGlobalWindow(KV.of(EMPTY_ARRAY, mainWindow)));
        }
        // Check to see if processing the request failed.
        MoreFutures.get(processResponse);
        waitForInboundTermination.awaitCompletion();
        WindowedValue<KV<byte[], TargetWindowT>> sideInputWindow = outputValue.poll();
        checkState(sideInputWindow != null, "Expected side input window to have been emitted by SDK harness.");
        checkState(sideInputWindow.getValue() != null, "Side input window emitted by SDK harness was a WindowedValue with no value in it.");
        checkState(sideInputWindow.getValue().getValue() != null, "Side input window emitted by SDK harness was a WindowedValue<KV<...>> with a null V.");
        checkState(outputValue.isEmpty(), "Expected only a single side input window to have been emitted by " + "the SDK harness but also received %s", outputValue);
        return sideInputWindow.getValue().getValue();
    } catch (Throwable e) {
        LOG.error("Unable to map main input window {} to side input window.", mainWindow, e);
        throw new IllegalStateException(e);
    }
}

Example 17 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse 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 18 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse 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());
}

Example 19 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse in project beam by apache.

the class RegisterAndProcessBundleOperationTest method testRegisterOnlyOnFirstBundle.

@Test
public void testRegisterOnlyOnFirstBundle() throws Exception {
    List<BeamFnApi.InstructionRequest> requests = new ArrayList<>();
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            requests.add(request);
            switch(request.getRequestCase()) {
                case REGISTER:
                case PROCESS_BUNDLE:
                    return CompletableFuture.completedFuture(responseFor(request).build());
                default:
                    // block forever on other requests
                    return new CompletableFuture<>();
            }
        }
    }, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
    // Ensure that the first time we start we send the register and process bundle requests
    assertThat(requests, empty());
    operation.start();
    assertEquals(requests.get(0), BeamFnApi.InstructionRequest.newBuilder().setInstructionId("777").setRegister(REGISTER_REQUEST).build());
    assertEquals(requests.get(1), BeamFnApi.InstructionRequest.newBuilder().setInstructionId("778").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("555")).build());
    operation.finish();
    assertEquals(false, operation.hasFailed());
    // Ensure on restart that we only send the process bundle request
    operation.start();
    assertEquals(requests.get(2), BeamFnApi.InstructionRequest.newBuilder().setInstructionId("779").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("555")).build());
    operation.finish();
    assertEquals(false, operation.hasFailed());
}

Example 20 with InstructionResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse in project beam by apache.

the class RegisterAndProcessBundleOperationTest method testProcessingBundleBlocksOnFinish.

@Test
public void testProcessingBundleBlocksOnFinish() throws Exception {
    List<BeamFnApi.InstructionRequest> requests = new ArrayList<>();
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    ExecutorService executorService = Executors.newCachedThreadPool();
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            requests.add(request);
            switch(request.getRequestCase()) {
                case REGISTER:
                    return CompletableFuture.completedFuture(responseFor(request).build());
                case PROCESS_BUNDLE:
                    CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>();
                    executorService.submit(() -> {
                        // Purposefully sleep simulating SDK harness doing work
                        Thread.sleep(100);
                        responseFuture.complete(responseFor(request).build());
                        return null;
                    });
                    return responseFuture;
                default:
                    // Anything else hangs; nothing else should be blocking
                    return new CompletableFuture<>();
            }
        }
    }, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of(), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
    operation.start();
    // This method blocks till the requests are completed
    operation.finish();
    assertEquals(false, operation.hasFailed());
    // Ensure that the messages were received
    assertEquals(requests.get(0), BeamFnApi.InstructionRequest.newBuilder().setInstructionId("777").setRegister(REGISTER_REQUEST).build());
    assertEquals(requests.get(1), BeamFnApi.InstructionRequest.newBuilder().setInstructionId("778").setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("555")).build());
}