Examples with ApiServiceDescriptor org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor used on opensource projects

Example 6 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class ProcessBundleDescriptors method addStageOutput.

private static OutputEncoding addStageOutput(ApiServiceDescriptor dataEndpoint, Components.Builder components, PCollectionNode outputPCollection, WireCoderSetting wireCoderSetting) throws IOException {
    String outputWireCoderId = WireCoders.addSdkWireCoder(outputPCollection, components, wireCoderSetting);
    @SuppressWarnings("unchecked") Coder<WindowedValue<?>> wireCoder = (Coder) WireCoders.instantiateRunnerWireCoder(outputPCollection, components.build(), wireCoderSetting);
    RemoteGrpcPort outputPort = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(dataEndpoint).setCoderId(outputWireCoderId).build();
    RemoteGrpcPortWrite outputWrite = RemoteGrpcPortWrite.writeToPort(outputPCollection.getId(), outputPort);
    String outputId = uniqueId(String.format("fn/write/%s", outputPCollection.getId()), components::containsTransforms);
    PTransform outputTransform = outputWrite.toPTransform();
    components.putTransforms(outputId, outputTransform);
    return new AutoValue_ProcessBundleDescriptors_OutputEncoding(outputId, wireCoder);
}

Example 7 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class BeamFnLoggingClientTest method testLogging.

@Test
public void testLogging() throws Exception {
    BeamFnLoggingMDC.setInstructionId("instruction-1");
    AtomicBoolean clientClosedStream = new AtomicBoolean();
    Collection<BeamFnApi.LogEntry> values = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.LogControl>> outboundServerObserver = new AtomicReference<>();
    CallStreamObserver<BeamFnApi.LogEntry.List> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.LogEntry.List logEntries) -> values.addAll(logEntries.getLogEntriesList())).withOnCompleted(() -> {
        // Remember that the client told us that this stream completed
        clientClosedStream.set(true);
        outboundServerObserver.get().onCompleted();
    }).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<BeamFnApi.LogControl> outboundObserver) {
            outboundServerObserver.set(outboundObserver);
            return inboundServerObserver;
        }
    }).build();
    server.start();
    ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
    try {
        BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
        // Keep a strong reference to the loggers in this block. Otherwise the call to client.close()
        // removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
        Logger rootLogger = LogManager.getLogManager().getLogger("");
        Logger configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
        // Ensure that log levels were correctly set.
        assertEquals(Level.OFF, rootLogger.getLevel());
        assertEquals(Level.FINE, configuredLogger.getLevel());
        // Should be filtered because the default log level override is OFF
        rootLogger.log(FILTERED_RECORD);
        // Should not be filtered because the default log level override for ConfiguredLogger is DEBUG
        configuredLogger.log(TEST_RECORD);
        configuredLogger.log(TEST_RECORD_WITH_EXCEPTION);
        client.close();
        // Verify that after close, log levels are reset.
        assertEquals(Level.INFO, rootLogger.getLevel());
        assertNull(configuredLogger.getLevel());
        assertTrue(clientClosedStream.get());
        assertTrue(channel.isShutdown());
        assertThat(values, contains(TEST_ENTRY, TEST_ENTRY_WITH_EXCEPTION));
    } finally {
        server.shutdownNow();
    }
}

Example 8 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class BeamFnLoggingClientTest method testWhenServerHangsUpEarlyThatClientIsAbleCleanup.

@Test
public void testWhenServerHangsUpEarlyThatClientIsAbleCleanup() throws Exception {
    BeamFnLoggingMDC.setInstructionId("instruction-1");
    Collection<BeamFnApi.LogEntry> values = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.LogControl>> outboundServerObserver = new AtomicReference<>();
    CallStreamObserver<BeamFnApi.LogEntry.List> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.LogEntry.List logEntries) -> values.addAll(logEntries.getLogEntriesList())).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<BeamFnApi.LogControl> outboundObserver) {
            outboundServerObserver.set(outboundObserver);
            outboundObserver.onCompleted();
            return inboundServerObserver;
        }
    }).build();
    server.start();
    ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
    try {
        BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
        // Keep a strong reference to the loggers in this block. Otherwise the call to client.close()
        // removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
        Logger rootLogger = LogManager.getLogManager().getLogger("");
        Logger configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
        client.close();
        // Verify that after close, log levels are reset.
        assertEquals(Level.INFO, rootLogger.getLevel());
        assertNull(configuredLogger.getLevel());
    } finally {
        assertTrue(channel.isShutdown());
        server.shutdownNow();
    }
}

Example 9 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class BeamFnDataWriteRunnerTest method testReuseForMultipleBundles.

@Test
public void testReuseForMultipleBundles() throws Exception {
    AtomicReference<String> bundleId = new AtomicReference<>("0");
    String localInputId = "inputPC";
    RunnerApi.PTransform pTransform = RemoteGrpcPortWrite.writeToPort(localInputId, PORT_SPEC).toPTransform();
    List<WindowedValue<String>> output0 = new ArrayList<>();
    List<WindowedValue<String>> output1 = new ArrayList<>();
    Map<ApiServiceDescriptor, BeamFnDataOutboundAggregator> aggregators = new HashMap<>();
    BeamFnDataOutboundAggregator aggregator = createRecordingAggregator(ImmutableMap.of("0", output0, "1", output1), bundleId::get);
    aggregators.put(PORT_SPEC.getApiServiceDescriptor(), aggregator);
    PTransformRunnerFactoryTestContext context = PTransformRunnerFactoryTestContext.builder(TRANSFORM_ID, pTransform).beamFnDataClient(mockBeamFnDataClient).processBundleInstructionIdSupplier(bundleId::get).outboundAggregators(aggregators).pCollections(ImmutableMap.of(localInputId, RunnerApi.PCollection.newBuilder().setCoderId(ELEM_CODER_ID).build())).coders(COMPONENTS.getCodersMap()).windowingStrategies(COMPONENTS.getWindowingStrategiesMap()).build();
    new BeamFnDataWriteRunner.Factory<String>().createRunnerForPTransform(context);
    assertThat(context.getPCollectionConsumers().keySet(), containsInAnyOrder(localInputId));
    FnDataReceiver<Object> pCollectionConsumer = context.getPCollectionConsumer(localInputId);
    pCollectionConsumer.accept(valueInGlobalWindow("ABC"));
    pCollectionConsumer.accept(valueInGlobalWindow("DEF"));
    assertThat(output0, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF")));
    output0.clear();
    // Process for bundle id 1
    bundleId.set("1");
    pCollectionConsumer.accept(valueInGlobalWindow("GHI"));
    pCollectionConsumer.accept(valueInGlobalWindow("JKL"));
    assertThat(output1, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
    verifyNoMoreInteractions(mockBeamFnDataClient);
}

Example 10 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class BeamFnControlClientTest method testDelegation.

@Test
public void testDelegation() throws Exception {
    AtomicBoolean clientClosedStream = new AtomicBoolean();
    BlockingQueue<BeamFnApi.InstructionResponse> values = new LinkedBlockingQueue<>();
    BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>();
    CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.withOnNext(values::add).withOnCompleted(() -> clientClosedStream.set(true)).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<BeamFnApi.InstructionResponse> control(StreamObserver<BeamFnApi.InstructionRequest> outboundObserver) {
            Uninterruptibles.putUninterruptibly(outboundServerObservers, outboundObserver);
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, value -> {
            assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
            return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance());
        });
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> {
            assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
            throw FAILURE;
        });
        ExecutorService executor = Executors.newCachedThreadPool();
        BeamFnControlClient client = new BeamFnControlClient(apiServiceDescriptor, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.trivial(), executor, handlers);
        // Get the connected client and attempt to send and receive an instruction
        StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take();
        outboundServerObserver.onNext(SUCCESSFUL_REQUEST);
        assertEquals(SUCCESSFUL_RESPONSE, values.take());
        // Ensure that conversion of an unknown request type is properly converted to a
        // failure response.
        outboundServerObserver.onNext(UNKNOWN_HANDLER_REQUEST);
        assertEquals(UNKNOWN_HANDLER_RESPONSE, values.take());
        // Ensure that all exceptions are caught and translated to failures
        outboundServerObserver.onNext(FAILURE_REQUEST);
        assertEquals(FAILURE_RESPONSE, values.take());
        // Ensure that the server completing the stream translates to the completable future
        // being completed allowing for a successful shutdown of the client.
        outboundServerObserver.onCompleted();
        client.waitForTermination();
    } finally {
        server.shutdownNow();
    }
}