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

Example 16 with List

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

the class BeamFnDataGrpcServiceTest method testMultipleClientsSendMessagesAreDirectedToProperConsumers.

@Test
public void testMultipleClientsSendMessagesAreDirectedToProperConsumers() throws Exception {
    LinkedBlockingQueue<BeamFnApi.Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    CountDownLatch waitForInboundElements = new CountDownLatch(1);
    for (int i = 0; i < 3; ++i) {
        String instructionId = Integer.toString(i);
        executorService.submit(() -> {
            ManagedChannel channel = ManagedChannelFactory.createDefault().withInterceptors(Arrays.asList(AddHarnessIdInterceptor.create(WORKER_ID))).forDescriptor(service.getApiServiceDescriptor());
            StreamObserver<BeamFnApi.Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
            outboundObserver.onNext(elementsWithData(instructionId));
            waitForInboundElements.await();
            outboundObserver.onCompleted();
            return null;
        });
    }
    List<Collection<WindowedValue<String>>> serverInboundValues = new ArrayList<>();
    Collection<InboundDataClient> inboundDataClients = new ArrayList<>();
    for (int i = 0; i < 3; ++i) {
        BlockingQueue<WindowedValue<String>> serverInboundValue = new LinkedBlockingQueue<>();
        serverInboundValues.add(serverInboundValue);
        inboundDataClients.add(service.getDataService(WORKER_ID).receive(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER, serverInboundValue::add));
    }
    // Waiting for the client provides the necessary synchronization for the elements to arrive.
    for (InboundDataClient inboundDataClient : inboundDataClients) {
        inboundDataClient.awaitCompletion();
    }
    waitForInboundElements.countDown();
    for (int i = 0; i < 3; ++i) {
        assertThat(serverInboundValues.get(i), contains(valueInGlobalWindow("A" + i), valueInGlobalWindow("B" + i), valueInGlobalWindow("C" + i)));
    }
    assertThat(clientInboundElements, empty());
}

Example 17 with List

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

the class RegisterAndProcessBundleOperationTest method testProcessingBundleBlocksOnFinishWithError.

@Test
public void testProcessingBundleBlocksOnFinishWithError() 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).setError("error").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(true, 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());
}

Example 18 with List

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

the class BeamFnLoggingServiceTest method testMultipleClientsFailingIsHandledGracefullyByServer.

@Test(timeout = 5000)
public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Exception {
    Collection<Callable<Void>> tasks = new ArrayList<>();
    ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
    try (BeamFnLoggingService service = new BeamFnLoggingService(findOpenPort(), logs::add, ServerStreamObserverFactory.fromOptions(PipelineOptionsFactory.create())::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor())) {
        server = ServerFactory.createDefault().create(Arrays.asList(service), service.getApiServiceDescriptor());
        CountDownLatch waitForTermination = new CountDownLatch(3);
        final BlockingQueue<StreamObserver<List>> outboundObservers = new LinkedBlockingQueue<>();
        for (int i = 1; i <= 3; ++i) {
            int instructionId = i;
            tasks.add(() -> {
                ManagedChannel channel = ManagedChannelFactory.createDefault().withInterceptors(Arrays.asList(AddHarnessIdInterceptor.create(WORKER_ID + instructionId))).forDescriptor(service.getApiServiceDescriptor());
                StreamObserver<BeamFnApi.LogEntry.List> outboundObserver = BeamFnLoggingGrpc.newStub(channel).logging(TestStreams.withOnNext(BeamFnLoggingServiceTest::discardMessage).withOnError(waitForTermination::countDown).build());
                outboundObserver.onNext(createLogsWithIds(instructionId, -instructionId));
                outboundObservers.add(outboundObserver);
                return null;
            });
        }
        ExecutorService executorService = Executors.newCachedThreadPool();
        executorService.invokeAll(tasks);
        for (int i = 1; i <= 3; ++i) {
            outboundObservers.take().onError(new RuntimeException("Client " + i));
        }
        waitForTermination.await();
    }
}

Example 19 with List

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

the class CommonCoderTest method parseField.

private static Object parseField(Object value, Schema.FieldType fieldType) {
    if (value == null) {
        return null;
    }
    switch(fieldType.getTypeName()) {
        case BYTE:
            return ((Number) value).byteValue();
        case INT16:
            return ((Number) value).shortValue();
        case INT32:
            return ((Number) value).intValue();
        case INT64:
            return ((Number) value).longValue();
        case FLOAT:
            return Float.parseFloat((String) value);
        case DOUBLE:
            return Double.parseDouble((String) value);
        case STRING:
            return (String) value;
        case BOOLEAN:
            return (Boolean) value;
        case BYTES:
            // extract String as byte[]
            return ((String) value).getBytes(StandardCharsets.ISO_8859_1);
        case ARRAY:
            return ((List<Object>) value).stream().map((element) -> parseField(element, fieldType.getCollectionElementType())).collect(toImmutableList());
        case MAP:
            Map<Object, Object> kvMap = new HashMap<>();
            ((Map<Object, Object>) value).entrySet().stream().forEach((entry) -> kvMap.put(parseField(entry.getKey(), fieldType.getMapKeyType()), parseField(entry.getValue(), fieldType.getMapValueType())));
            return kvMap;
        case ROW:
            // Clone map so we don't mutate the underlying value
            Map<String, Object> rowMap = new HashMap<>((Map<String, Object>) value);
            Schema schema = fieldType.getRowSchema();
            Row.Builder row = Row.withSchema(schema);
            for (Schema.Field field : schema.getFields()) {
                Object element = rowMap.remove(field.getName());
                if (element != null) {
                    element = parseField(element, field.getType());
                }
                row.addValue(element);
            }
            if (!rowMap.isEmpty()) {
                throw new IllegalArgumentException("Value contains keys that are not in the schema: " + rowMap.keySet());
            }
            return row.build();
        case LOGICAL_TYPE:
            // representation type, then convert to the base type.
            return fieldType.getLogicalType().toInputType(parseField(value, fieldType.getLogicalType().getBaseType()));
        default:
            // DECIMAL, DATETIME
            throw new IllegalArgumentException("Unsupported type name: " + fieldType.getTypeName());
    }
}

Example 20 with List

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

the class ProcessBundleHandlerTest method setupProcessBundleHandlerForSimpleRecordingDoFn.

private ProcessBundleHandler setupProcessBundleHandlerForSimpleRecordingDoFn(List<String> dataOutput, List<Timers> timerOutput, boolean enableOutputEmbedding) throws Exception {
    DoFnWithExecutionInformation doFnWithExecutionInformation = DoFnWithExecutionInformation.of(new SimpleDoFn(), SimpleDoFn.MAIN_OUTPUT_TAG, Collections.emptyMap(), DoFnSchemaInformation.create());
    RunnerApi.FunctionSpec functionSpec = RunnerApi.FunctionSpec.newBuilder().setUrn(ParDoTranslation.CUSTOM_JAVA_DO_FN_URN).setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(doFnWithExecutionInformation))).build();
    RunnerApi.ParDoPayload parDoPayload = ParDoPayload.newBuilder().setDoFn(functionSpec).putTimerFamilySpecs("tfs-" + SimpleDoFn.TIMER_FAMILY_ID, TimerFamilySpec.newBuilder().setTimeDomain(RunnerApi.TimeDomain.Enum.EVENT_TIME).setTimerFamilyCoderId("timer-coder").build()).build();
    BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = ProcessBundleDescriptor.newBuilder().putTransforms("2L", PTransform.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).putOutputs("2L-output", "2L-output-pc").build()).putTransforms("3L", PTransform.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN).setPayload(parDoPayload.toByteString())).putInputs("3L-input", "2L-output-pc").build()).putPcollections("2L-output-pc", PCollection.newBuilder().setWindowingStrategyId("window-strategy").setCoderId("2L-output-coder").setIsBounded(IsBounded.Enum.BOUNDED).build()).putWindowingStrategies("window-strategy", WindowingStrategy.newBuilder().setWindowCoderId("window-strategy-coder").setWindowFn(FunctionSpec.newBuilder().setUrn("beam:window_fn:global_windows:v1")).setOutputTime(OutputTime.Enum.END_OF_WINDOW).setAccumulationMode(AccumulationMode.Enum.ACCUMULATING).setTrigger(Trigger.newBuilder().setAlways(Always.getDefaultInstance())).setClosingBehavior(ClosingBehavior.Enum.EMIT_ALWAYS).setOnTimeBehavior(OnTimeBehavior.Enum.FIRE_ALWAYS).build()).setTimerApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("url").build()).putCoders("string_coder", CoderTranslation.toProto(StringUtf8Coder.of()).getCoder()).putCoders("2L-output-coder", Coder.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(ModelCoders.KV_CODER_URN).build()).addComponentCoderIds("string_coder").addComponentCoderIds("string_coder").build()).putCoders("window-strategy-coder", Coder.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(ModelCoders.GLOBAL_WINDOW_CODER_URN).build()).build()).putCoders("timer-coder", Coder.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(ModelCoders.TIMER_CODER_URN)).addComponentCoderIds("string_coder").addComponentCoderIds("window-strategy-coder").build()).build();
    Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
    Map<String, PTransformRunnerFactory> urnToPTransformRunnerFactoryMap = Maps.newHashMap(REGISTERED_RUNNER_FACTORIES);
    urnToPTransformRunnerFactoryMap.put(DATA_INPUT_URN, (PTransformRunnerFactory<Object>) (context) -> {
        context.addIncomingDataEndpoint(ApiServiceDescriptor.getDefaultInstance(), KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()), (input) -> {
            dataOutput.add(input.getValue());
        });
        return null;
    });
    Mockito.doAnswer((invocation) -> new BeamFnDataOutboundAggregator(PipelineOptionsFactory.create(), invocation.getArgument(1), new StreamObserver<Elements>() {

        @Override
        public void onNext(Elements elements) {
            for (Timers timer : elements.getTimersList()) {
                timerOutput.addAll(elements.getTimersList());
            }
        }

        @Override
        public void onError(Throwable throwable) {
        }

        @Override
        public void onCompleted() {
        }
    }, invocation.getArgument(2))).when(beamFnDataClient).createOutboundAggregator(any(), any(), anyBoolean());
    return new ProcessBundleHandler(PipelineOptionsFactory.create(), enableOutputEmbedding ? Collections.singleton(BeamUrns.getUrn(StandardRunnerProtocols.Enum.CONTROL_RESPONSE_ELEMENTS_EMBEDDING)) : Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, null, /* beamFnStateClient */
    null, /* finalizeBundleHandler */
    new ShortIdMap(), urnToPTransformRunnerFactoryMap, Caches.noop(), new BundleProcessorCache());
}