Examples with InstructionRequestHandler org.apache.beam.runners.fnexecution.control.InstructionRequestHandler used on opensource projects

Example 11 with InstructionRequestHandler

use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.

the class EmbeddedEnvironmentFactory method createEnvironment.

@Override
// no need to monitor shutdown thread
@SuppressWarnings("FutureReturnValueIgnored")
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
    ExecutorService executor = Executors.newSingleThreadExecutor();
    Future<?> fnHarness = executor.submit(() -> {
        try {
            FnHarness.main(workerId, options, // Runner capabilities.
            Collections.emptySet(), loggingServer.getApiServiceDescriptor(), controlServer.getApiServiceDescriptor(), null, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.clientDirect(), Caches.fromOptions(options));
        } catch (NoClassDefFoundError e) {
            // TODO: https://issues.apache.org/jira/browse/BEAM-4384 load the FnHarness in a
            // Restricted classpath that we control for any user.
            LOG.error("{} while executing an in-process FnHarness. " + "To use the {}, " + "the 'org.apache.beam:beam-sdks-java-harness' artifact " + "and its dependencies must be on the classpath", NoClassDefFoundError.class.getSimpleName(), EmbeddedEnvironmentFactory.class.getSimpleName(), e);
            throw e;
        }
        return null;
    });
    executor.submit(() -> {
        try {
            fnHarness.get();
        } catch (Throwable t) {
            executor.shutdownNow();
        }
    });
    InstructionRequestHandler handler = null;
    // Wait on a client from the gRPC server.
    while (handler == null) {
        try {
            // If the thread is not alive anymore, we abort.
            if (executor.isShutdown()) {
                throw new IllegalStateException("FnHarness startup failed");
            }
            handler = clientSource.take(workerId, Duration.ofSeconds(5L));
        } catch (TimeoutException timeoutEx) {
            LOG.info("Still waiting for startup of FnHarness");
        } catch (InterruptedException interruptEx) {
            Thread.currentThread().interrupt();
            throw new RuntimeException(interruptEx);
        }
    }
    return RemoteEnvironment.forHandler(environment, handler);
}

Example 12 with InstructionRequestHandler

use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.

the class ExternalEnvironmentFactory method createEnvironment.

/**
 * Creates a new, active {@link RemoteEnvironment} backed by an unmanaged worker.
 */
@Override
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
    Preconditions.checkState(environment.getUrn().equals(BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.EXTERNAL)), "The passed environment does not contain an ExternalPayload.");
    final RunnerApi.ExternalPayload externalPayload = RunnerApi.ExternalPayload.parseFrom(environment.getPayload());
    BeamFnApi.StartWorkerRequest startWorkerRequest = BeamFnApi.StartWorkerRequest.newBuilder().setWorkerId(workerId).setControlEndpoint(controlServiceServer.getApiServiceDescriptor()).setLoggingEndpoint(loggingServiceServer.getApiServiceDescriptor()).setArtifactEndpoint(retrievalServiceServer.getApiServiceDescriptor()).setProvisionEndpoint(provisioningServiceServer.getApiServiceDescriptor()).putAllParams(externalPayload.getParamsMap()).build();
    LOG.debug("Requesting worker ID {}", workerId);
    final ManagedChannel managedChannel = ManagedChannelFactory.createDefault().forDescriptor(externalPayload.getEndpoint());
    BeamFnApi.StartWorkerResponse startWorkerResponse = BeamFnExternalWorkerPoolGrpc.newBlockingStub(managedChannel).startWorker(startWorkerRequest);
    if (!startWorkerResponse.getError().isEmpty()) {
        throw new RuntimeException(startWorkerResponse.getError());
    }
    // Wait on a client from the gRPC server.
    InstructionRequestHandler instructionHandler = null;
    while (instructionHandler == null) {
        try {
            instructionHandler = clientSource.take(workerId, Duration.ofMinutes(2));
        } catch (TimeoutException timeoutEx) {
            LOG.info("Still waiting for startup of environment from {} for worker id {}", externalPayload.getEndpoint().getUrl(), workerId);
        } catch (InterruptedException interruptEx) {
            Thread.currentThread().interrupt();
            managedChannel.shutdownNow();
            throw new RuntimeException(interruptEx);
        }
    }
    final InstructionRequestHandler finalInstructionHandler = instructionHandler;
    return new RemoteEnvironment() {

        @Override
        public Environment getEnvironment() {
            return environment;
        }

        @Override
        public InstructionRequestHandler getInstructionRequestHandler() {
            return finalInstructionHandler;
        }

        @Override
        public void close() throws Exception {
            try {
                finalInstructionHandler.close();
                BeamFnApi.StopWorkerRequest stopWorkerRequest = BeamFnApi.StopWorkerRequest.newBuilder().setWorkerId(workerId).build();
                LOG.debug("Closing worker ID {}", workerId);
                BeamFnApi.StopWorkerResponse stopWorkerResponse = BeamFnExternalWorkerPoolGrpc.newBlockingStub(managedChannel).stopWorker(stopWorkerRequest);
                if (!stopWorkerResponse.getError().isEmpty()) {
                    throw new RuntimeException(stopWorkerResponse.getError());
                }
            } finally {
                managedChannel.shutdown();
                managedChannel.awaitTermination(10, TimeUnit.SECONDS);
                if (!managedChannel.isTerminated()) {
                    managedChannel.shutdownNow();
                }
            }
        }
    };
}

Example 13 with InstructionRequestHandler

use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.

the class ProcessEnvironmentTest method closeClosesInstructionRequestHandler.

@Test
public void closeClosesInstructionRequestHandler() throws Exception {
    InstructionRequestHandler handler = mock(InstructionRequestHandler.class);
    RemoteEnvironment env = ProcessEnvironment.create(mock(ProcessManager.class), Environment.getDefaultInstance(), "1", handler);
    env.close();
    verify(handler).close();
}

Example 14 with InstructionRequestHandler

use of org.apache.beam.runners.fnexecution.control.InstructionRequestHandler in project beam by apache.

the class RemoteEnvironmentTest method closeClosesInstructionRequestHandler.

@Test
public void closeClosesInstructionRequestHandler() throws Exception {
    InstructionRequestHandler handler = mock(InstructionRequestHandler.class);
    RemoteEnvironment env = new RemoteEnvironment() {

        @Override
        public Environment getEnvironment() {
            throw new UnsupportedOperationException();
        }

        @Override
        public InstructionRequestHandler getInstructionRequestHandler() {
            return handler;
        }
    };
    env.close();
    verify(handler).close();
}