Examples with ServerFactory org.apache.beam.sdk.fn.server.ServerFactory used on opensource projects

Example 1 with ServerFactory

use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.

the class ServerFactoryTest method usesUrlFactory.

@Test
public void usesUrlFactory() throws Exception {
    ServerFactory serverFactory = ServerFactory.createWithUrlFactory((host, port) -> "foo");
    CallStreamObserver<Elements> observer = TestStreams.withOnNext((Elements unused) -> {
    }).withOnCompleted(() -> {
    }).build();
    TestDataService service = new TestDataService(observer);
    ApiServiceDescriptor.Builder descriptorBuilder = ApiServiceDescriptor.newBuilder();
    grpcCleanupRule.register(serverFactory.allocateAddressAndCreate(ImmutableList.of(service), descriptorBuilder));
    assertThat(descriptorBuilder.getUrl(), is("foo"));
}

Example 2 with ServerFactory

use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.

the class DefaultJobBundleFactoryTest method createsMultipleEnvironmentOfSingleType.

@Test
public void createsMultipleEnvironmentOfSingleType() throws Exception {
    ServerFactory serverFactory = ServerFactory.createDefault();
    Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[1])).build();
    Environment environmentAA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[2])).build();
    EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
    when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
    when(envFactoryA.createEnvironment(eq(environmentAA), any())).thenReturn(remoteEnvironment);
    EnvironmentFactory.Provider environmentProviderFactoryA = mock(EnvironmentFactory.Provider.class);
    when(environmentProviderFactoryA.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryA);
    when(environmentProviderFactoryA.getServerFactory()).thenReturn(serverFactory);
    Environment environmentB = Environment.newBuilder().setUrn("env:urn:b").build();
    EnvironmentFactory envFactoryB = mock(EnvironmentFactory.class);
    when(envFactoryB.createEnvironment(eq(environmentB), any())).thenReturn(remoteEnvironment);
    EnvironmentFactory.Provider environmentProviderFactoryB = mock(EnvironmentFactory.Provider.class);
    when(environmentProviderFactoryB.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryB);
    when(environmentProviderFactoryB.getServerFactory()).thenReturn(serverFactory);
    Map<String, Provider> environmentFactoryProviderMap = ImmutableMap.of(environmentA.getUrn(), environmentProviderFactoryA, environmentB.getUrn(), environmentProviderFactoryB);
    try (DefaultJobBundleFactory bundleFactory = createDefaultJobBundleFactory(environmentFactoryProviderMap)) {
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(1)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(environmentProviderFactoryB, Mockito.times(0)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
        verify(envFactoryA, Mockito.times(0)).createEnvironment(eq(environmentAA), any());
        bundleFactory.forStage(getExecutableStage(environmentAA));
        verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(environmentProviderFactoryB, Mockito.times(0)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentAA), any());
    }
}

Example 3 with ServerFactory

use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.

the class DataflowRunnerHarness method main.

/**
 * Fetches and processes work units from the Dataflow service.
 */
public static void main(String[] unusedArgs) throws Exception {
    RunnerApi.@Nullable Pipeline pipeline = DataflowWorkerHarnessHelper.getPipelineFromEnv();
    // This descriptor is used for all services except logging. They are isolated to keep
    // critical traffic protected from best effort traffic.
    ApiServiceDescriptor controlApiService = DataflowWorkerHarnessHelper.getControlDescriptor();
    ApiServiceDescriptor loggingApiService = DataflowWorkerHarnessHelper.getLoggingDescriptor();
    ApiServiceDescriptor statusApiService = DataflowWorkerHarnessHelper.getStatusDescriptor();
    LOG.info("{} started, using port {} for control, {} for logging.", DataflowRunnerHarness.class, controlApiService, loggingApiService);
    DataflowWorkerHarnessHelper.initializeLogging(DataflowRunnerHarness.class);
    DataflowWorkerHarnessOptions pipelineOptions = DataflowWorkerHarnessHelper.initializeGlobalStateAndPipelineOptions(DataflowRunnerHarness.class);
    DataflowWorkerHarnessHelper.configureLogging(pipelineOptions);
    // Initialized registered file systems.˜
    FileSystems.setDefaultPipelineOptions(pipelineOptions);
    DataflowPipelineDebugOptions dataflowOptions = pipelineOptions.as(DataflowPipelineDebugOptions.class);
    ServerFactory serverFactory;
    if (DataflowRunner.hasExperiment(dataflowOptions, "beam_fn_api_epoll_domain_socket")) {
        serverFactory = ServerFactory.createEpollDomainSocket();
    } else if (DataflowRunner.hasExperiment(dataflowOptions, "beam_fn_api_epoll")) {
        serverFactory = ServerFactory.createEpollSocket();
    } else {
        serverFactory = ServerFactory.createDefault();
    }
    ServerStreamObserverFactory streamObserverFactory = ServerStreamObserverFactory.fromOptions(pipelineOptions);
    Server servicesServer = null;
    Server loggingServer = null;
    Server statusServer = null;
    try (BeamFnLoggingService beamFnLoggingService = new BeamFnLoggingService(loggingApiService, DataflowWorkerLoggingInitializer.getSdkLoggingHandler()::publish, streamObserverFactory::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        BeamFnControlService beamFnControlService = new BeamFnControlService(controlApiService, streamObserverFactory::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        BeamFnDataGrpcService beamFnDataService = new BeamFnDataGrpcService(pipelineOptions, controlApiService, streamObserverFactory::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        BeamWorkerStatusGrpcService beamWorkerStatusGrpcService = statusApiService == null ? null : BeamWorkerStatusGrpcService.create(statusApiService, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        GrpcStateService beamFnStateService = GrpcStateService.create()) {
        servicesServer = serverFactory.create(ImmutableList.of(beamFnControlService, beamFnDataService, beamFnStateService), controlApiService);
        loggingServer = serverFactory.create(ImmutableList.of(beamFnLoggingService), loggingApiService);
        // gRPC server for obtaining SDK harness runtime status information.
        if (beamWorkerStatusGrpcService != null) {
            statusServer = serverFactory.create(ImmutableList.of(beamWorkerStatusGrpcService), statusApiService);
        }
        start(pipeline, pipelineOptions, beamFnControlService, beamFnDataService, controlApiService, beamFnStateService, beamWorkerStatusGrpcService);
        if (statusServer != null) {
            statusServer.shutdown();
        }
        servicesServer.shutdown();
        loggingServer.shutdown();
        // wait 30 secs for outstanding requests to finish.
        if (statusServer != null) {
            statusServer.awaitTermination(30, TimeUnit.SECONDS);
        }
        servicesServer.awaitTermination(30, TimeUnit.SECONDS);
        loggingServer.awaitTermination(30, TimeUnit.SECONDS);
    } finally {
        if (statusServer != null && !statusServer.isTerminated()) {
            statusServer.shutdownNow();
        }
        if (servicesServer != null && !servicesServer.isTerminated()) {
            servicesServer.shutdownNow();
        }
        if (loggingServer != null && !loggingServer.isTerminated()) {
            loggingServer.shutdownNow();
        }
    }
}

Example 4 with ServerFactory

use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.

the class DefaultJobBundleFactory method createEnvironmentCaches.

private ImmutableList<EnvironmentCacheAndLock> createEnvironmentCaches(ThrowingFunction<ServerFactory, ServerInfo> serverInfoCreator, int count) {
    ImmutableList.Builder<EnvironmentCacheAndLock> caches = ImmutableList.builder();
    for (int i = 0; i < count; i++) {
        final Lock refLock;
        if (environmentExpirationMillis > 0) {
            // The lock ensures there is no race condition between expiring an environment and a client
            // still attempting to use it, hence referencing it.
            refLock = new ReentrantLock(true);
        } else {
            refLock = NoopLock.get();
        }
        CacheBuilder<Environment, WrappedSdkHarnessClient> cacheBuilder = CacheBuilder.newBuilder().removalListener(notification -> {
            WrappedSdkHarnessClient client = notification.getValue();
            final int refCount;
            // We need to use a lock here to ensure we are not causing the environment to
            // be removed if beforehand a StageBundleFactory has retrieved it but not yet
            // issued ref() on it.
            refLock.lock();
            try {
                refCount = client.unref();
            } finally {
                refLock.unlock();
            }
            if (refCount > 0) {
                LOG.warn("Expiring environment {} with {} remaining bundle references. Taking note to clean it up during shutdown if the references are not removed by then.", notification.getKey(), refCount);
                evictedActiveClients.add(client);
            }
        });
        if (environmentExpirationMillis > 0) {
            cacheBuilder.expireAfterWrite(environmentExpirationMillis, TimeUnit.MILLISECONDS);
        }
        LoadingCache<Environment, WrappedSdkHarnessClient> cache = cacheBuilder.build(new CacheLoader<Environment, WrappedSdkHarnessClient>() {

            @Override
            public WrappedSdkHarnessClient load(Environment environment) throws Exception {
                EnvironmentFactory.Provider environmentFactoryProvider = environmentFactoryProviderMap.get(environment.getUrn());
                ServerFactory serverFactory = environmentFactoryProvider.getServerFactory();
                ServerInfo serverInfo = serverInfoCreator.apply(serverFactory);
                String workerId = stageIdGenerator.getId();
                serverInfo.getProvisioningServer().getService().registerEnvironment(workerId, environment);
                EnvironmentFactory environmentFactory = environmentFactoryProvider.createEnvironmentFactory(serverInfo.getControlServer(), serverInfo.getLoggingServer(), serverInfo.getRetrievalServer(), serverInfo.getProvisioningServer(), clientPool, stageIdGenerator);
                return WrappedSdkHarnessClient.wrapping(environmentFactory.createEnvironment(environment, workerId), serverInfo);
            }
        });
        caches.add(new EnvironmentCacheAndLock(cache, refLock));
    }
    return caches.build();
}

Example 5 with ServerFactory

use of org.apache.beam.sdk.fn.server.ServerFactory in project beam by apache.

the class ServerFactoryTest method urlFactoryWithPortSupplier.

@Test
public void urlFactoryWithPortSupplier() throws Exception {
    ServerFactory serverFactory = ServerFactory.createWithUrlFactoryAndPortSupplier((host, port) -> "foo" + ":" + port, () -> 65535);
    CallStreamObserver<Elements> observer = TestStreams.withOnNext((Elements unused) -> {
    }).withOnCompleted(() -> {
    }).build();
    TestDataService service = new TestDataService(observer);
    ApiServiceDescriptor.Builder descriptorBuilder = ApiServiceDescriptor.newBuilder();
    grpcCleanupRule.register(serverFactory.allocateAddressAndCreate(ImmutableList.of(service), descriptorBuilder));
    assertThat(descriptorBuilder.getUrl(), is("foo:65535"));
}