Examples with ThrowingFunction org.apache.beam.sdk.function.ThrowingFunction used on opensource projects

Example 1 with ThrowingFunction

use of org.apache.beam.sdk.function.ThrowingFunction in project beam by apache.

the class AssignWindowsRunnerTest method factoryCreatesFromJavaWindowFn.

@Test
public void factoryCreatesFromJavaWindowFn() throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    PTransform windowPTransform = PTransform.newBuilder().putInputs("in", "input").putOutputs("out", "output").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN).setPayload(WindowIntoPayload.newBuilder().setWindowFn(WindowingStrategyTranslation.toProto(new TestWindowFn(), components)).build().toByteString()).build()).build();
    ThrowingFunction<WindowedValue<?>, WindowedValue<?>> fn = (ThrowingFunction) factory.forPTransform("transform", windowPTransform);
    assertThat(fn.apply(WindowedValue.of(22L, new Instant(5), new IntervalWindow(new Instant(0L), new Instant(20027L)), PaneInfo.ON_TIME_AND_ONLY_FIRING)), equalTo(WindowedValue.of(22L, new Instant(5), new TestWindowFn().assignWindow(new Instant(5)), PaneInfo.ON_TIME_AND_ONLY_FIRING)));
}

Example 2 with ThrowingFunction

use of org.apache.beam.sdk.function.ThrowingFunction in project beam by apache.

the class AssignWindowsRunnerTest method factoryCreatesFromKnownWindowFn.

@Test
public void factoryCreatesFromKnownWindowFn() throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    PTransform windowPTransform = PTransform.newBuilder().putInputs("in", "input").putOutputs("out", "output").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN).setPayload(WindowIntoPayload.newBuilder().setWindowFn(WindowingStrategyTranslation.toProto(Sessions.withGapDuration(Duration.standardMinutes(12L)), components)).build().toByteString()).build()).build();
    ThrowingFunction<WindowedValue<?>, WindowedValue<?>> fn = (ThrowingFunction) factory.forPTransform("transform", windowPTransform);
    WindowedValue<?> output = fn.apply(WindowedValue.of(22L, new Instant(5), new IntervalWindow(new Instant(0L), new Instant(20027L)), PaneInfo.ON_TIME_AND_ONLY_FIRING));
    assertThat(output, equalTo(WindowedValue.of(22L, new Instant(5), new IntervalWindow(new Instant(5L), Duration.standardMinutes(12L)), PaneInfo.ON_TIME_AND_ONLY_FIRING)));
}

Example 3 with ThrowingFunction

use of org.apache.beam.sdk.function.ThrowingFunction 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();
    }
}

Example 4 with ThrowingFunction

use of org.apache.beam.sdk.function.ThrowingFunction in project beam by apache.

the class BeamFnControlClientTest method testJavaErrorResponse.

@Test
public void testJavaErrorResponse() throws Exception {
    BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>();
    BlockingQueue<Throwable> error = new LinkedBlockingQueue<>();
    CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.<BeamFnApi.InstructionResponse>withOnNext(response -> fail(String.format("Unexpected Response %s", response))).withOnError(error::add).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.REGISTER, value -> {
            assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
            throw new Error("Test Error");
        });
        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();
        // Ensure that all exceptions are caught and translated to failures
        outboundServerObserver.onNext(InstructionRequest.newBuilder().setInstructionId("0").setRegister(RegisterRequest.getDefaultInstance()).build());
        // There should be an error reported to the StreamObserver.
        assertThat(error.take(), not(nullValue()));
        // Ensure that the client shuts down when an Error is thrown from the harness
        try {
            client.waitForTermination();
            throw new IllegalStateException("The future should have terminated with an error");
        } catch (ExecutionException errorWrapper) {
            assertThat(errorWrapper.getCause().getMessage(), containsString("Test Error"));
        }
    } finally {
        server.shutdownNow();
    }
}

Example 5 with ThrowingFunction

use of org.apache.beam.sdk.function.ThrowingFunction in project beam by apache.

the class FnHarness method main.

/**
 * Run a FnHarness with the given id and options that attaches to the specified logging and
 * control API service descriptors using the given channel factory and outbound observer factory.
 *
 * @param id Harness ID
 * @param options The options for this pipeline
 * @param runnerCapabilites
 * @param loggingApiServiceDescriptor
 * @param controlApiServiceDescriptor
 * @param statusApiServiceDescriptor
 * @param channelFactory
 * @param outboundObserverFactory
 * @param processWideCache
 * @throws Exception
 */
public static void main(String id, PipelineOptions options, Set<String> runnerCapabilites, Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor, Endpoints.ApiServiceDescriptor controlApiServiceDescriptor, Endpoints.ApiServiceDescriptor statusApiServiceDescriptor, ManagedChannelFactory channelFactory, OutboundObserverFactory outboundObserverFactory, Cache<Object, Object> processWideCache) throws Exception {
    channelFactory = channelFactory.withInterceptors(ImmutableList.of(AddHarnessIdInterceptor.create(id)));
    IdGenerator idGenerator = IdGenerators.decrementingLongs();
    ShortIdMap metricsShortIds = new ShortIdMap();
    ExecutorService executorService = options.as(GcsOptions.class).getExecutorService();
    // intercepts logging and sends it to the logging service.
    try (BeamFnLoggingClient logging = new BeamFnLoggingClient(options, loggingApiServiceDescriptor, channelFactory::forDescriptor)) {
        LOG.info("Fn Harness started");
        // Register standard file systems.
        FileSystems.setDefaultPipelineOptions(options);
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        ManagedChannel channel = channelFactory.forDescriptor(controlApiServiceDescriptor);
        BeamFnControlGrpc.BeamFnControlStub controlStub = BeamFnControlGrpc.newStub(channel);
        BeamFnControlGrpc.BeamFnControlBlockingStub blockingControlStub = BeamFnControlGrpc.newBlockingStub(channel);
        BeamFnDataGrpcClient beamFnDataMultiplexer = new BeamFnDataGrpcClient(options, channelFactory::forDescriptor, outboundObserverFactory);
        BeamFnStateGrpcClientCache beamFnStateGrpcClientCache = new BeamFnStateGrpcClientCache(idGenerator, channelFactory, outboundObserverFactory);
        FinalizeBundleHandler finalizeBundleHandler = new FinalizeBundleHandler(options.as(GcsOptions.class).getExecutorService());
        Function<String, BeamFnApi.ProcessBundleDescriptor> getProcessBundleDescriptor = new Function<String, ProcessBundleDescriptor>() {

            private static final String PROCESS_BUNDLE_DESCRIPTORS = "ProcessBundleDescriptors";

            private final Cache<String, BeamFnApi.ProcessBundleDescriptor> cache = Caches.subCache(processWideCache, PROCESS_BUNDLE_DESCRIPTORS);

            @Override
            public BeamFnApi.ProcessBundleDescriptor apply(String id) {
                return cache.computeIfAbsent(id, this::loadDescriptor);
            }

            private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
                return blockingControlStub.getProcessBundleDescriptor(BeamFnApi.GetProcessBundleDescriptorRequest.newBuilder().setProcessBundleDescriptorId(id).build());
            }
        };
        MetricsEnvironment.setProcessWideContainer(MetricsContainerImpl.createProcessWideContainer());
        ProcessBundleHandler processBundleHandler = new ProcessBundleHandler(options, runnerCapabilites, getProcessBundleDescriptor, beamFnDataMultiplexer, beamFnStateGrpcClientCache, finalizeBundleHandler, metricsShortIds, processWideCache);
        BeamFnStatusClient beamFnStatusClient = null;
        if (statusApiServiceDescriptor != null) {
            beamFnStatusClient = new BeamFnStatusClient(statusApiServiceDescriptor, channelFactory::forDescriptor, processBundleHandler.getBundleProcessorCache(), options, processWideCache);
        }
        // TODO(BEAM-9729): Remove once runners no longer send this instruction.
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, request -> BeamFnApi.InstructionResponse.newBuilder().setRegister(BeamFnApi.RegisterResponse.getDefaultInstance()));
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.FINALIZE_BUNDLE, finalizeBundleHandler::finalizeBundle);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, processBundleHandler::processBundle);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE_PROGRESS, processBundleHandler::progress);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE_SPLIT, processBundleHandler::trySplit);
        handlers.put(InstructionRequest.RequestCase.MONITORING_INFOS, request -> BeamFnApi.InstructionResponse.newBuilder().setMonitoringInfos(BeamFnApi.MonitoringInfosMetadataResponse.newBuilder().putAllMonitoringInfo(StreamSupport.stream(request.getMonitoringInfos().getMonitoringInfoIdList().spliterator(), false).collect(Collectors.toMap(Function.identity(), metricsShortIds::get)))));
        HarnessMonitoringInfosInstructionHandler processWideHandler = new HarnessMonitoringInfosInstructionHandler(metricsShortIds);
        handlers.put(InstructionRequest.RequestCase.HARNESS_MONITORING_INFOS, processWideHandler::harnessMonitoringInfos);
        JvmInitializers.runBeforeProcessing(options);
        String samplingPeriodMills = ExperimentalOptions.getExperimentValue(options, ExperimentalOptions.STATE_SAMPLING_PERIOD_MILLIS);
        if (samplingPeriodMills != null) {
            ExecutionStateSampler.setSamplingPeriod(Integer.parseInt(samplingPeriodMills));
        }
        ExecutionStateSampler.instance().start();
        LOG.info("Entering instruction processing loop");
        // The control client immediately dispatches requests to an executor so we execute on the
        // direct executor. If we created separate channels for different stubs we could use
        // directExecutor() when building the channel.
        BeamFnControlClient control = new BeamFnControlClient(controlStub.withExecutor(MoreExecutors.directExecutor()), outboundObserverFactory, executorService, handlers);
        control.waitForTermination();
        if (beamFnStatusClient != null) {
            beamFnStatusClient.close();
        }
        processBundleHandler.shutdown();
    } finally {
        System.out.println("Shutting SDK harness down.");
        ExecutionStateSampler.instance().stop();
        executorService.shutdown();
    }
}