Examples with GrpcDataService org.apache.beam.runners.fnexecution.data.GrpcDataService used on opensource projects

Example 1 with GrpcDataService

use of org.apache.beam.runners.fnexecution.data.GrpcDataService in project beam by apache.

the class DefaultJobBundleFactoryTest method setUpMocks.

@Before
public void setUpMocks() throws Exception {
    MockitoAnnotations.initMocks(this);
    when(envFactory.createEnvironment(eq(environment), any())).thenReturn(remoteEnvironment);
    when(remoteEnvironment.getInstructionRequestHandler()).thenReturn(instructionHandler);
    when(instructionHandler.handle(any())).thenReturn(CompletableFuture.completedFuture(instructionResponse));
    when(dataServer.getApiServiceDescriptor()).thenReturn(ApiServiceDescriptor.getDefaultInstance());
    GrpcDataService dataService = mock(GrpcDataService.class);
    when(dataService.send(any(), any())).thenReturn(mock(CloseableFnDataReceiver.class));
    when(dataServer.getService()).thenReturn(dataService);
    when(stateServer.getApiServiceDescriptor()).thenReturn(ApiServiceDescriptor.getDefaultInstance());
    GrpcStateService stateService = mock(GrpcStateService.class);
    when(stateService.registerForProcessBundleInstructionId(any(), any())).thenReturn(mock(StateDelegator.Registration.class));
    when(stateServer.getService()).thenReturn(stateService);
    when(provisioningServer.getService()).thenReturn(provisionService);
    serverInfo = new AutoValue_DefaultJobBundleFactory_ServerInfo.Builder().setControlServer(controlServer).setLoggingServer(loggingServer).setRetrievalServer(retrievalServer).setProvisioningServer(provisioningServer).setDataServer(dataServer).setStateServer(stateServer).build();
}

Example 2 with GrpcDataService

use of org.apache.beam.runners.fnexecution.data.GrpcDataService in project beam by apache.

the class DefaultJobBundleFactory method createServerInfo.

private ServerInfo createServerInfo(JobInfo jobInfo, ServerFactory serverFactory) throws IOException {
    Preconditions.checkNotNull(serverFactory, "serverFactory can not be null");
    PortablePipelineOptions portableOptions = PipelineOptionsTranslation.fromProto(jobInfo.pipelineOptions()).as(PortablePipelineOptions.class);
    GrpcFnServer<FnApiControlClientPoolService> controlServer = GrpcFnServer.allocatePortAndCreateFor(FnApiControlClientPoolService.offeringClientsToPool(clientPool.getSink(), GrpcContextHeaderAccessorProvider.getHeaderAccessor()), serverFactory);
    GrpcFnServer<GrpcLoggingService> loggingServer = GrpcFnServer.allocatePortAndCreateFor(GrpcLoggingService.forWriter(Slf4jLogWriter.getDefault()), serverFactory);
    GrpcFnServer<ArtifactRetrievalService> retrievalServer = GrpcFnServer.allocatePortAndCreateFor(new ArtifactRetrievalService(), serverFactory);
    ProvisionApi.ProvisionInfo.Builder provisionInfo = jobInfo.toProvisionInfo().toBuilder();
    provisionInfo.setLoggingEndpoint(loggingServer.getApiServiceDescriptor());
    provisionInfo.setArtifactEndpoint(retrievalServer.getApiServiceDescriptor());
    provisionInfo.setControlEndpoint(controlServer.getApiServiceDescriptor());
    GrpcFnServer<StaticGrpcProvisionService> provisioningServer = GrpcFnServer.allocatePortAndCreateFor(StaticGrpcProvisionService.create(provisionInfo.build(), GrpcContextHeaderAccessorProvider.getHeaderAccessor()), serverFactory);
    GrpcFnServer<GrpcDataService> dataServer = GrpcFnServer.allocatePortAndCreateFor(GrpcDataService.create(portableOptions, executor, OutboundObserverFactory.serverDirect()), serverFactory);
    GrpcFnServer<GrpcStateService> stateServer = GrpcFnServer.allocatePortAndCreateFor(GrpcStateService.create(), serverFactory);
    ServerInfo serverInfo = new AutoValue_DefaultJobBundleFactory_ServerInfo.Builder().setControlServer(controlServer).setLoggingServer(loggingServer).setRetrievalServer(retrievalServer).setProvisioningServer(provisioningServer).setDataServer(dataServer).setStateServer(stateServer).build();
    return serverInfo;
}