Search in sources :

Example 1 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class DefaultJobBundleFactoryTest method loadBalancesBundles.

@Test
public void loadBalancesBundles() throws Exception {
    PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    portableOptions.setSdkWorkerParallelism(2);
    portableOptions.setLoadBalanceBundles(true);
    Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), envFactoryProviderMap, stageIdGenerator, serverInfo)) {
        OutputReceiverFactory orf = mock(OutputReceiverFactory.class);
        StateRequestHandler srh = mock(StateRequestHandler.class);
        when(srh.getCacheTokens()).thenReturn(Collections.emptyList());
        StageBundleFactory sbf = bundleFactory.forStage(getExecutableStage(environment));
        RemoteBundle b1 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        verify(envFactory, Mockito.times(1)).createEnvironment(eq(environment), any());
        final RemoteBundle b2 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        verify(envFactory, Mockito.times(2)).createEnvironment(eq(environment), any());
        AtomicBoolean b2Closing = new AtomicBoolean(false);
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        ScheduledFuture<Optional<Exception>> closingFuture = executor.schedule(() -> {
            try {
                b2Closing.compareAndSet(false, true);
                b2.close();
                return Optional.empty();
            } catch (Exception e) {
                return Optional.of(e);
            }
        }, 100, TimeUnit.MILLISECONDS);
        assertThat(b2Closing.get(), equalTo(false));
        // This call should block until closingFuture has finished closing b2 (100ms)
        RemoteBundle b3 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        // ensure the previous call waited for close
        assertThat(b2Closing.get(), equalTo(true));
        // Join closingFuture and check if an exception occurred
        Optional<Exception> closingException = closingFuture.get();
        if (closingException.isPresent()) {
            throw new AssertionError("Exception occurred while closing b2", closingException.get());
        }
        verify(envFactory, Mockito.times(2)).createEnvironment(eq(environment), any());
        b3.close();
        b1.close();
        executor.shutdown();
    }
}
Also used : StateRequestHandler(org.apache.beam.runners.fnexecution.state.StateRequestHandler) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) Optional(java.util.Optional) ExpectedException(org.junit.rules.ExpectedException) Struct(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) PortablePipelineOptions(org.apache.beam.sdk.options.PortablePipelineOptions) Test(org.junit.Test)

Example 2 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class PipelineOptionsTranslation method toProto.

/**
 * Converts the provided {@link PipelineOptions} to a {@link Struct}.
 */
public static Struct toProto(PipelineOptions options) {
    Struct.Builder builder = Struct.newBuilder();
    try {
        // TODO: Officially define URNs for options and their scheme.
        JsonNode treeNode = MAPPER.valueToTree(options);
        JsonNode rootOptions = treeNode.get("options");
        Iterator<Map.Entry<String, JsonNode>> optionsEntries = rootOptions.fields();
        if (!optionsEntries.hasNext()) {
            // (observed with version 2.2.3)
            throw new RuntimeException("Unable to convert pipeline options, please check for outdated jackson-core version in the classpath.");
        }
        Map<String, TreeNode> optionsUsingUrns = new HashMap<>();
        while (optionsEntries.hasNext()) {
            Map.Entry<String, JsonNode> entry = optionsEntries.next();
            optionsUsingUrns.put("beam:option:" + CaseFormat.LOWER_CAMEL.to(CaseFormat.LOWER_UNDERSCORE, entry.getKey()) + ":v1", entry.getValue());
        }
        // The JSON format of a Protobuf Struct is the JSON object that is equivalent to that struct
        // (with values encoded in a standard json-codeable manner). See Beam PR 3719 for more.
        JsonFormat.parser().merge(MAPPER.writeValueAsString(optionsUsingUrns), builder);
        return builder.build();
    } catch (IOException e) {
        throw new RuntimeException("Failed to convert PipelineOptions to Protocol", e);
    }
}
Also used : HashMap(java.util.HashMap) JsonNode(com.fasterxml.jackson.databind.JsonNode) IOException(java.io.IOException) Struct(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct) TreeNode(com.fasterxml.jackson.core.TreeNode) HashMap(java.util.HashMap) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map)

Example 3 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct 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));
}
Also used : ProvisionInfo(org.apache.beam.model.fnexecution.v1.ProvisionApi.ProvisionInfo) GetProvisionInfoResponse(org.apache.beam.model.fnexecution.v1.ProvisionApi.GetProvisionInfoResponse) ProvisionServiceBlockingStub(org.apache.beam.model.fnexecution.v1.ProvisionServiceGrpc.ProvisionServiceBlockingStub) Struct(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct) Test(org.junit.Test)

Example 4 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project beam by apache.

the class InMemoryJobService method prepare.

@Override
public void prepare(PrepareJobRequest request, StreamObserver<PrepareJobResponse> responseObserver) {
    try {
        LOG.trace("{} {}", PrepareJobRequest.class.getSimpleName(), request);
        // insert preparation
        String preparationId = String.format("%s_%s", request.getJobName(), UUID.randomUUID().toString());
        Struct pipelineOptions = request.getPipelineOptions();
        if (pipelineOptions == null) {
            throw new NullPointerException("Encountered null pipeline options.");
        }
        LOG.trace("PIPELINE OPTIONS {} {}", pipelineOptions.getClass(), pipelineOptions);
        JobPreparation preparation = JobPreparation.builder().setId(preparationId).setPipeline(request.getPipeline()).setOptions(pipelineOptions).build();
        JobPreparation previous = preparations.putIfAbsent(preparationId, preparation);
        if (previous != null) {
            // this should never happen with a UUID
            String errMessage = String.format("A job with the preparation ID \"%s\" already exists.", preparationId);
            StatusException exception = Status.NOT_FOUND.withDescription(errMessage).asException();
            responseObserver.onError(exception);
            return;
        }
        String stagingSessionToken = stagingServiceTokenProvider.apply(preparationId);
        stagingSessionTokens.putIfAbsent(preparationId, stagingSessionToken);
        stagingService.getService().registerJob(stagingSessionToken, Maps.transformValues(request.getPipeline().getComponents().getEnvironmentsMap(), RunnerApi.Environment::getDependenciesList));
        // send response
        PrepareJobResponse response = PrepareJobResponse.newBuilder().setPreparationId(preparationId).setArtifactStagingEndpoint(stagingServiceDescriptor).setStagingSessionToken(stagingSessionToken).build();
        responseObserver.onNext(response);
        responseObserver.onCompleted();
    } catch (Exception e) {
        LOG.error("Could not prepare job with name {}", request.getJobName(), e);
        responseObserver.onError(Status.INTERNAL.withCause(e).asException());
    }
}
Also used : StatusException(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.StatusException) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) PrepareJobResponse(org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobResponse) PrepareJobRequest(org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobRequest) StatusRuntimeException(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.StatusRuntimeException) StatusException(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.StatusException) Struct(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct)

Example 5 with Struct

use of org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct in project grpc-java by grpc.

the class ClientXdsClientDataTest method parseHttpFilter_typedStructMigration.

@Test
public void parseHttpFilter_typedStructMigration() {
    filterRegistry.register(new TestFilter());
    Struct rawStruct = Struct.newBuilder().putFields("name", Value.newBuilder().setStringValue("default").build()).build();
    HttpFilter httpFilter = HttpFilter.newBuilder().setIsOptional(true).setTypedConfig(Any.pack(com.github.udpa.udpa.type.v1.TypedStruct.newBuilder().setTypeUrl("test-url").setValue(rawStruct).build())).build();
    FilterConfig config = ClientXdsClient.parseHttpFilter(httpFilter, filterRegistry, true).getStruct();
    assertThat(((SimpleFilterConfig) config).getConfig()).isEqualTo(rawStruct);
    HttpFilter httpFilterNewTypeStruct = HttpFilter.newBuilder().setIsOptional(true).setTypedConfig(Any.pack(TypedStruct.newBuilder().setTypeUrl("test-url").setValue(rawStruct).build())).build();
    config = ClientXdsClient.parseHttpFilter(httpFilterNewTypeStruct, filterRegistry, true).getStruct();
    assertThat(((SimpleFilterConfig) config).getConfig()).isEqualTo(rawStruct);
}
Also used : HttpFilter(io.envoyproxy.envoy.extensions.filters.network.http_connection_manager.v3.HttpFilter) FilterConfig(io.grpc.xds.Filter.FilterConfig) Struct(com.google.protobuf.Struct) TypedStruct(com.github.udpa.udpa.type.v1.TypedStruct) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)13 Struct (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.Struct)9 Struct (com.google.protobuf.Struct)8 PortablePipelineOptions (org.apache.beam.sdk.options.PortablePipelineOptions)5 ListValue (com.google.protobuf.ListValue)3 TestStruct (com.google.protobuf.util.JsonTestProto.TestStruct)3 JobInfo (org.apache.beam.runners.fnexecution.provisioning.JobInfo)3 StateRequestHandler (org.apache.beam.runners.fnexecution.state.StateRequestHandler)3 TypedStruct (com.github.udpa.udpa.type.v1.TypedStruct)2 Any (com.google.protobuf.Any)2 BoolValue (com.google.protobuf.BoolValue)2 BytesValue (com.google.protobuf.BytesValue)2 DoubleValue (com.google.protobuf.DoubleValue)2 FloatValue (com.google.protobuf.FloatValue)2 Int32Value (com.google.protobuf.Int32Value)2 Int64Value (com.google.protobuf.Int64Value)2 NullValue (com.google.protobuf.NullValue)2 StringValue (com.google.protobuf.StringValue)2 UInt32Value (com.google.protobuf.UInt32Value)2 UInt64Value (com.google.protobuf.UInt64Value)2