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

Example 6 with List

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

the class FnHarnessTest method testLaunchFnHarnessAndTeardownCleanly.

@Test
// failure will cause test to timeout.
@SuppressWarnings("FutureReturnValueIgnored")
public void testLaunchFnHarnessAndTeardownCleanly() throws Exception {
    Function<String, String> environmentVariableMock = mock(Function.class);
    PipelineOptions options = PipelineOptionsFactory.create();
    when(environmentVariableMock.apply("HARNESS_ID")).thenReturn("id");
    when(environmentVariableMock.apply("PIPELINE_OPTIONS")).thenReturn(PipelineOptionsTranslation.toJson(options));
    List<BeamFnApi.LogEntry> logEntries = new ArrayList<>();
    List<BeamFnApi.InstructionResponse> instructionResponses = mock(List.class);
    BeamFnLoggingGrpc.BeamFnLoggingImplBase loggingService = new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<LogControl> responseObserver) {
            return TestStreams.withOnNext((BeamFnApi.LogEntry.List entries) -> logEntries.addAll(entries.getLogEntriesList())).withOnCompleted(responseObserver::onCompleted).build();
        }
    };
    BeamFnControlGrpc.BeamFnControlImplBase controlService = new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<InstructionResponse> control(StreamObserver<InstructionRequest> responseObserver) {
            CountDownLatch waitForResponses = new CountDownLatch(1);
            options.as(GcsOptions.class).getExecutorService().submit(() -> {
                responseObserver.onNext(INSTRUCTION_REQUEST);
                Uninterruptibles.awaitUninterruptibly(waitForResponses);
                responseObserver.onCompleted();
            });
            return TestStreams.withOnNext((InstructionResponse t) -> {
                instructionResponses.add(t);
                waitForResponses.countDown();
            }).withOnCompleted(waitForResponses::countDown).build();
        }
    };
    Server loggingServer = ServerBuilder.forPort(0).addService(loggingService).build();
    loggingServer.start();
    try {
        Server controlServer = ServerBuilder.forPort(0).addService(controlService).build();
        controlServer.start();
        try {
            Endpoints.ApiServiceDescriptor loggingDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + loggingServer.getPort()).build();
            Endpoints.ApiServiceDescriptor controlDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + controlServer.getPort()).build();
            when(environmentVariableMock.apply("LOGGING_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(loggingDescriptor));
            when(environmentVariableMock.apply("CONTROL_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(controlDescriptor));
            FnHarness.main(environmentVariableMock);
        } finally {
            controlServer.shutdownNow();
        }
    } finally {
        loggingServer.shutdownNow();
    }
    // Verify that we first run onStartup functions before even reading the environment, and that
    // we then call beforeProcessing functions before executing instructions.
    InOrder inOrder = inOrder(onStartupMock, beforeProcessingMock, environmentVariableMock, instructionResponses);
    inOrder.verify(onStartupMock).run();
    inOrder.verify(environmentVariableMock, atLeastOnce()).apply(any());
    inOrder.verify(beforeProcessingMock).accept(any());
    inOrder.verify(instructionResponses).add(INSTRUCTION_RESPONSE);
}

Example 7 with List

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

the class FakeBeamFnStateClient method handle.

@Override
public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
    // The id should never be filled out
    assertEquals("", requestBuilder.getId());
    requestBuilder.setId(generateId());
    StateRequest request = requestBuilder.build();
    StateKey key = request.getStateKey();
    StateResponse.Builder response;
    assertNotEquals(RequestCase.REQUEST_NOT_SET, request.getRequestCase());
    assertNotEquals(TypeCase.TYPE_NOT_SET, key.getTypeCase());
    // multimap side input and runner based state keys only support get requests
    if (key.getTypeCase() == TypeCase.MULTIMAP_SIDE_INPUT || key.getTypeCase() == TypeCase.RUNNER) {
        assertEquals(RequestCase.GET, request.getRequestCase());
    }
    switch(request.getRequestCase()) {
        case GET:
            List<ByteString> byteStrings = data.getOrDefault(request.getStateKey(), Collections.singletonList(ByteString.EMPTY));
            int block = 0;
            if (request.getGet().getContinuationToken().size() > 0) {
                block = Integer.parseInt(request.getGet().getContinuationToken().toStringUtf8());
            }
            ByteString returnBlock = byteStrings.get(block);
            ByteString continuationToken = ByteString.EMPTY;
            if (byteStrings.size() > block + 1) {
                continuationToken = ByteString.copyFromUtf8(Integer.toString(block + 1));
            }
            response = StateResponse.newBuilder().setGet(StateGetResponse.newBuilder().setData(returnBlock).setContinuationToken(continuationToken));
            break;
        case CLEAR:
            data.remove(request.getStateKey());
            response = StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance());
            break;
        case APPEND:
            List<ByteString> previousValue = data.computeIfAbsent(request.getStateKey(), (unused) -> new ArrayList<>());
            previousValue.add(request.getAppend().getData());
            response = StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance());
            break;
        default:
            throw new IllegalStateException(String.format("Unknown request type %s", request.getRequestCase()));
    }
    return CompletableFuture.completedFuture(response.setId(requestBuilder.getId()).build());
}

Example 8 with List

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

the class ProcessBundleHandler method embedOutboundElementsIfApplicable.

private void embedOutboundElementsIfApplicable(ProcessBundleResponse.Builder response, BundleProcessor bundleProcessor) {
    if (bundleProcessor.getOutboundAggregators().isEmpty()) {
        return;
    }
    List<Elements> collectedElements = new ArrayList<>(bundleProcessor.getOutboundAggregators().size());
    boolean hasFlushedAggregator = false;
    for (BeamFnDataOutboundAggregator aggregator : bundleProcessor.getOutboundAggregators().values()) {
        Elements elements = aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
        if (elements == null) {
            hasFlushedAggregator = true;
        }
        collectedElements.add(elements);
    }
    if (!hasFlushedAggregator) {
        Elements.Builder elementsToEmbed = Elements.newBuilder();
        for (Elements collectedElement : collectedElements) {
            elementsToEmbed.mergeFrom(collectedElement);
        }
        response.setElements(elementsToEmbed.build());
    } else {
        // Since there was at least one flushed aggregator, we have to use the aggregators that were
        // able to successfully collect their elements to emit them and can not send them as part of
        // the ProcessBundleResponse.
        int i = 0;
        for (BeamFnDataOutboundAggregator aggregator : bundleProcessor.getOutboundAggregators().values()) {
            Elements elements = collectedElements.get(i++);
            if (elements != null) {
                aggregator.sendElements(elements);
            }
        }
    }
}

Example 9 with List

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

the class BeamFnDataReadRunnerTest method executeSplit.

private static ProcessBundleSplitResponse executeSplit(BeamFnDataReadRunner<String> readRunner, String pTransformId, String bundleId, long index, double fractionOfRemainder, long inputElements, List<Long> allowedSplitPoints) throws Exception {
    for (long i = -1; i < index; i++) {
        readRunner.forwardElementToConsumer(valueInGlobalWindow(Long.valueOf(i).toString()));
    }
    ProcessBundleSplitRequest request = ProcessBundleSplitRequest.newBuilder().setInstructionId(bundleId).putDesiredSplits(pTransformId, DesiredSplit.newBuilder().setEstimatedInputElements(inputElements).setFractionOfRemainder(fractionOfRemainder).addAllAllowedSplitPoints(allowedSplitPoints).build()).build();
    ProcessBundleSplitResponse.Builder responseBuilder = ProcessBundleSplitResponse.newBuilder();
    readRunner.trySplit(request, responseBuilder);
    return responseBuilder.build();
}

Example 10 with List

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

the class MultimapUserState method startStateApiWrites.

@SuppressWarnings("FutureReturnValueIgnored")
private void startStateApiWrites() {
    // Clear currently persisted key-values
    if (isCleared) {
        beamFnStateClient.handle(keysStateRequest.toBuilder().setClear(StateClearRequest.getDefaultInstance()));
    } else if (!pendingRemoves.isEmpty()) {
        for (K key : pendingRemoves.values()) {
            StateRequest request = createUserStateRequest(key);
            beamFnStateClient.handle(request.toBuilder().setClear(StateClearRequest.getDefaultInstance()));
        }
    }
    // Persist pending key-values
    if (!pendingAdds.isEmpty()) {
        for (KV<K, List<V>> entry : pendingAdds.values()) {
            StateRequest request = createUserStateRequest(entry.getKey());
            beamFnStateClient.handle(request.toBuilder().setAppend(StateAppendRequest.newBuilder().setData(encodeValues(entry.getValue()))));
        }
    }
}