use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server 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();
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class PortableRunner method run.
@Override
public PipelineResult run(Pipeline pipeline) {
Runnable cleanup;
if (Environments.ENVIRONMENT_LOOPBACK.equals(options.as(PortablePipelineOptions.class).getDefaultEnvironmentType())) {
GrpcFnServer<ExternalWorkerService> workerService;
try {
workerService = new ExternalWorkerService(options).start();
} catch (Exception exn) {
throw new RuntimeException("Failed to start GrpcFnServer for ExternalWorkerService", exn);
}
LOG.info("Starting worker service at {}", workerService.getApiServiceDescriptor().getUrl());
options.as(PortablePipelineOptions.class).setDefaultEnvironmentConfig(workerService.getApiServiceDescriptor().getUrl());
cleanup = () -> {
try {
LOG.warn("closing worker service {}", workerService);
workerService.close();
} catch (Exception exn) {
throw new RuntimeException(exn);
}
};
} else {
cleanup = null;
}
ImmutableList.Builder<String> filesToStageBuilder = ImmutableList.builder();
List<String> stagingFiles = options.as(PortablePipelineOptions.class).getFilesToStage();
if (stagingFiles == null) {
List<String> classpathResources = detectClassPathResourcesToStage(Environments.class.getClassLoader(), options);
if (classpathResources.isEmpty()) {
throw new IllegalArgumentException("No classpath elements found.");
}
LOG.debug("PortablePipelineOptions.filesToStage was not specified. " + "Defaulting to files from the classpath: {}", classpathResources.size());
filesToStageBuilder.addAll(classpathResources);
} else {
filesToStageBuilder.addAll(stagingFiles);
}
// TODO(heejong): remove jar_packages experimental flag when cross-language dependency
// management is implemented for all runners.
List<String> experiments = options.as(ExperimentalOptions.class).getExperiments();
if (experiments != null) {
Optional<String> jarPackages = experiments.stream().filter((String flag) -> flag.startsWith("jar_packages=")).findFirst();
jarPackages.ifPresent(s -> filesToStageBuilder.addAll(Arrays.asList(s.replaceFirst("jar_packages=", "").split(","))));
}
options.as(PortablePipelineOptions.class).setFilesToStage(filesToStageBuilder.build());
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, SdkComponents.create(options));
pipelineProto = DefaultArtifactResolver.INSTANCE.resolveArtifacts(pipelineProto);
PrepareJobRequest prepareJobRequest = PrepareJobRequest.newBuilder().setJobName(options.getJobName()).setPipeline(pipelineProto).setPipelineOptions(PipelineOptionsTranslation.toProto(options)).build();
LOG.info("Using job server endpoint: {}", endpoint);
ManagedChannel jobServiceChannel = channelFactory.forDescriptor(ApiServiceDescriptor.newBuilder().setUrl(endpoint).build());
JobServiceBlockingStub jobService = JobServiceGrpc.newBlockingStub(jobServiceChannel);
try (CloseableResource<JobServiceBlockingStub> wrappedJobService = CloseableResource.of(jobService, unused -> jobServiceChannel.shutdown())) {
final int jobServerTimeout = options.as(PortablePipelineOptions.class).getJobServerTimeout();
PrepareJobResponse prepareJobResponse = jobService.withDeadlineAfter(jobServerTimeout, TimeUnit.SECONDS).withWaitForReady().prepare(prepareJobRequest);
LOG.info("PrepareJobResponse: {}", prepareJobResponse);
ApiServiceDescriptor artifactStagingEndpoint = prepareJobResponse.getArtifactStagingEndpoint();
String stagingSessionToken = prepareJobResponse.getStagingSessionToken();
try (CloseableResource<ManagedChannel> artifactChannel = CloseableResource.of(channelFactory.forDescriptor(artifactStagingEndpoint), ManagedChannel::shutdown)) {
ArtifactStagingService.offer(new ArtifactRetrievalService(), ArtifactStagingServiceGrpc.newStub(artifactChannel.get()), stagingSessionToken);
} catch (CloseableResource.CloseException e) {
LOG.warn("Error closing artifact staging channel", e);
// CloseExceptions should only be thrown while closing the channel.
} catch (Exception e) {
throw new RuntimeException("Error staging files.", e);
}
RunJobRequest runJobRequest = RunJobRequest.newBuilder().setPreparationId(prepareJobResponse.getPreparationId()).build();
// Run the job and wait for a result, we don't set a timeout here because
// it may take a long time for a job to complete and streaming
// jobs never return a response.
RunJobResponse runJobResponse = jobService.run(runJobRequest);
LOG.info("RunJobResponse: {}", runJobResponse);
ByteString jobId = runJobResponse.getJobIdBytes();
return new JobServicePipelineResult(jobId, jobServerTimeout, wrappedJobService.transfer(), cleanup);
} catch (CloseException e) {
throw new RuntimeException(e);
}
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class PortableRunnerTest method createJobServer.
private void createJobServer(JobState.Enum jobState, JobApi.MetricResults metricResults) throws IOException {
ArtifactStagingService stagingService = new ArtifactStagingService(new ArtifactStagingService.ArtifactDestinationProvider() {
@Override
public ArtifactStagingService.ArtifactDestination getDestination(String stagingToken, String name) throws IOException {
return ArtifactStagingService.ArtifactDestination.create("/dev/null", ByteString.EMPTY, ByteStreams.nullOutputStream());
}
@Override
public void removeStagedArtifacts(String stagingToken) {
}
});
stagingService.registerJob("TestStagingToken", ImmutableMap.of());
Server server = grpcCleanupRule.register(InProcessServerBuilder.forName(ENDPOINT_URL).addService(new TestJobService(ENDPOINT_DESCRIPTOR, "prepId", "jobId", jobState, metricResults)).addService(stagingService).build());
server.start();
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class StaticGrpcProvisionServiceTest method returnsProvisionInfo.
@Test
public void returnsProvisionInfo() throws Exception {
Struct options = Struct.newBuilder().putFields("foo", Value.newBuilder().setBoolValue(true).build()).putFields("bar", Value.newBuilder().setNumberValue(2.5).build()).putFields("baz", Value.newBuilder().setListValue(ListValue.newBuilder().addValues(Value.newBuilder().setStructValue(Struct.newBuilder().putFields("spam", Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build()))).build()).build()).build();
ProvisionInfo info = ProvisionInfo.newBuilder().setPipelineOptions(options).build();
GrpcFnServer<StaticGrpcProvisionService> server = GrpcFnServer.allocatePortAndCreateFor(StaticGrpcProvisionService.create(info, GrpcContextHeaderAccessorProvider.getHeaderAccessor()), InProcessServerFactory.create());
ProvisionServiceBlockingStub stub = ProvisionServiceGrpc.newBlockingStub(InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).intercept(AddHarnessIdInterceptor.create("test_worker")).build());
GetProvisionInfoResponse provisionResponse = stub.getProvisionInfo(GetProvisionInfoRequest.getDefaultInstance());
assertThat(provisionResponse.getInfo(), equalTo(info));
}
use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server in project beam by apache.
the class GrpcLoggingServiceTest method testMultipleClientsSuccessfullyProcessed.
@Test
public void testMultipleClientsSuccessfullyProcessed() throws Exception {
ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
GrpcLoggingService service = GrpcLoggingService.forWriter(new CollectionAppendingLogWriter(logs));
try (GrpcFnServer<GrpcLoggingService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
Collection<Callable<Void>> tasks = new ArrayList<>();
for (int i = 1; i <= 3; ++i) {
final int instructionId = i;
tasks.add(() -> {
CountDownLatch waitForServerHangup = new CountDownLatch(1);
String url = server.getApiServiceDescriptor().getUrl();
ManagedChannel channel = InProcessChannelBuilder.forName(url).build();
StreamObserver<LogEntry.List> outboundObserver = BeamFnLoggingGrpc.newStub(channel).logging(TestStreams.withOnNext(messageDiscarder).withOnCompleted(new CountDown(waitForServerHangup)).build());
outboundObserver.onNext(createLogsWithIds(instructionId, -instructionId));
outboundObserver.onCompleted();
waitForServerHangup.await();
return null;
});
}
ExecutorService executorService = Executors.newCachedThreadPool();
executorService.invokeAll(tasks);
assertThat(logs, containsInAnyOrder(createLogWithId(1L), createLogWithId(2L), createLogWithId(3L), createLogWithId(-1L), createLogWithId(-2L), createLogWithId(-3L)));
}
}
Aggregations