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

Example 6 with InstructionRequest

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

the class RegisterAndProcessBundleOperation method getProcessBundleProgress.

/**
 * Returns the compound metrics recorded, by issuing a request to the SDK harness.
 *
 * <p>This includes key progress indicators as well as user-defined metrics.
 *
 * <p>Use {@link #getInputElementsConsumed} on the future value to extract the elements consumed
 * from the upstream read operation.
 *
 * <p>May be called at any time, including before start() and after finish().
 */
public CompletionStage<BeamFnApi.ProcessBundleProgressResponse> getProcessBundleProgress() {
    // processBundleId may be reset if this bundle finishes asynchronously.
    String processBundleId = this.processBundleId;
    if (processBundleId == null) {
        return CompletableFuture.completedFuture(BeamFnApi.ProcessBundleProgressResponse.getDefaultInstance());
    }
    InstructionRequest processBundleRequest = InstructionRequest.newBuilder().setInstructionId(idGenerator.getId()).setProcessBundleProgress(ProcessBundleProgressRequest.newBuilder().setInstructionId(processBundleId)).build();
    return instructionRequestHandler.handle(processBundleRequest).thenApply(InstructionResponse::getProcessBundleProgress);
}

Example 7 with InstructionRequest

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

the class RegisterAndProcessBundleOperation method start.

@Override
public void start() throws Exception {
    try (Closeable scope = context.enterStart()) {
        super.start();
        // Only register once by using the presence of the future as a signal.
        if (registerFuture == null) {
            InstructionRequest request = InstructionRequest.newBuilder().setInstructionId(idGenerator.getId()).setRegister(registerRequest).build();
            registerFuture = instructionRequestHandler.handle(request);
        }
        checkState(registerRequest.getProcessBundleDescriptorCount() == 1, "Only one bundle registration at a time currently supported.");
        InstructionRequest processBundleRequest = InstructionRequest.newBuilder().setInstructionId(getProcessBundleInstructionId()).setProcessBundle(ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId(registerRequest.getProcessBundleDescriptor(0).getId())).build();
        deregisterStateHandler = beamFnStateDelegator.registerForProcessBundleInstructionId(getProcessBundleInstructionId(), this::delegateByStateKeyType);
        processBundleResponse = getRegisterResponse(registerFuture).thenCompose(registerResponse -> instructionRequestHandler.handle(processBundleRequest));
    }
}

Example 8 with InstructionRequest

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

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

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