Examples with StateResponse org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse used on opensource projects

Example 1 with StateResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse in project beam by apache.

the class FakeBeamFnStateClient method handle.

@Override
public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
    // The id should never be filled out
    assertEquals("", requestBuilder.getId());
    requestBuilder.setId(generateId());
    StateRequest request = requestBuilder.build();
    StateKey key = request.getStateKey();
    StateResponse.Builder response;
    assertNotEquals(RequestCase.REQUEST_NOT_SET, request.getRequestCase());
    assertNotEquals(TypeCase.TYPE_NOT_SET, key.getTypeCase());
    // multimap side input and runner based state keys only support get requests
    if (key.getTypeCase() == TypeCase.MULTIMAP_SIDE_INPUT || key.getTypeCase() == TypeCase.RUNNER) {
        assertEquals(RequestCase.GET, request.getRequestCase());
    }
    switch(request.getRequestCase()) {
        case GET:
            List<ByteString> byteStrings = data.getOrDefault(request.getStateKey(), Collections.singletonList(ByteString.EMPTY));
            int block = 0;
            if (request.getGet().getContinuationToken().size() > 0) {
                block = Integer.parseInt(request.getGet().getContinuationToken().toStringUtf8());
            }
            ByteString returnBlock = byteStrings.get(block);
            ByteString continuationToken = ByteString.EMPTY;
            if (byteStrings.size() > block + 1) {
                continuationToken = ByteString.copyFromUtf8(Integer.toString(block + 1));
            }
            response = StateResponse.newBuilder().setGet(StateGetResponse.newBuilder().setData(returnBlock).setContinuationToken(continuationToken));
            break;
        case CLEAR:
            data.remove(request.getStateKey());
            response = StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance());
            break;
        case APPEND:
            List<ByteString> previousValue = data.computeIfAbsent(request.getStateKey(), (unused) -> new ArrayList<>());
            previousValue.add(request.getAppend().getData());
            response = StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance());
            break;
        default:
            throw new IllegalStateException(String.format("Unknown request type %s", request.getRequestCase()));
    }
    return CompletableFuture.completedFuture(response.setId(requestBuilder.getId()).build());
}

Example 2 with StateResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse in project beam by apache.

the class BeamFnStateGrpcClientCacheTest method testCachingOfClient.

@Test
public void testCachingOfClient() throws Exception {
    Endpoints.ApiServiceDescriptor otherApiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(apiServiceDescriptor.getUrl() + "-other").build();
    Server testServer2 = InProcessServerBuilder.forName(otherApiServiceDescriptor.getUrl()).addService(new BeamFnStateGrpc.BeamFnStateImplBase() {

        @Override
        public StreamObserver<StateRequest> state(StreamObserver<StateResponse> outboundObserver) {
            throw new IllegalStateException("Unexpected in test.");
        }
    }).build();
    testServer2.start();
    try {
        assertSame(clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(apiServiceDescriptor));
        assertNotSame(clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(otherApiServiceDescriptor));
    } finally {
        testServer2.shutdownNow();
    }
}

Example 3 with StateResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse in project beam by apache.

the class ProcessBundleHandlerTest method testPendingStateCallsBlockTillCompletion.

@Test
public void testPendingStateCallsBlockTillCompletion() throws Exception {
    BeamFnApi.ProcessBundleDescriptor processBundleDescriptor = BeamFnApi.ProcessBundleDescriptor.newBuilder().putTransforms("2L", RunnerApi.PTransform.newBuilder().setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(DATA_INPUT_URN).build()).build()).setStateApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance()).build();
    Map<String, BeamFnApi.ProcessBundleDescriptor> fnApiRegistry = ImmutableMap.of("1L", processBundleDescriptor);
    CompletableFuture<StateResponse>[] successfulResponse = new CompletableFuture[1];
    CompletableFuture<StateResponse>[] unsuccessfulResponse = new CompletableFuture[1];
    BeamFnStateGrpcClientCache mockBeamFnStateGrpcClient = Mockito.mock(BeamFnStateGrpcClientCache.class);
    BeamFnStateClient mockBeamFnStateClient = Mockito.mock(BeamFnStateClient.class);
    when(mockBeamFnStateGrpcClient.forApiServiceDescriptor(any())).thenReturn(mockBeamFnStateClient);
    doAnswer(invocation -> {
        StateRequest.Builder stateRequestBuilder = (StateRequest.Builder) invocation.getArguments()[0];
        CompletableFuture<StateResponse> completableFuture = new CompletableFuture<>();
        new Thread(() -> {
            // Simulate sleeping which introduces a race which most of the time requires
            // the ProcessBundleHandler to block.
            Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
            switch(stateRequestBuilder.getInstructionId()) {
                case "SUCCESS":
                    completableFuture.complete(StateResponse.getDefaultInstance());
                    break;
                case "FAIL":
                    completableFuture.completeExceptionally(new RuntimeException("TEST ERROR"));
            }
        }).start();
        return completableFuture;
    }).when(mockBeamFnStateClient).handle(any());
    ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), fnApiRegistry::get, beamFnDataClient, mockBeamFnStateGrpcClient, null, /* finalizeBundleHandler */
    new ShortIdMap(), ImmutableMap.of(DATA_INPUT_URN, new PTransformRunnerFactory<Object>() {

        @Override
        public Object createRunnerForPTransform(Context context) throws IOException {
            BeamFnStateClient beamFnStateClient = context.getBeamFnStateClient();
            context.addStartBundleFunction(() -> doStateCalls(beamFnStateClient));
            return null;
        }

        private void doStateCalls(BeamFnStateClient beamFnStateClient) {
            successfulResponse[0] = beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("SUCCESS"));
            unsuccessfulResponse[0] = beamFnStateClient.handle(StateRequest.newBuilder().setInstructionId("FAIL"));
        }
    }), Caches.noop(), new BundleProcessorCache());
    handler.processBundle(BeamFnApi.InstructionRequest.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("1L")).build());
    assertTrue(successfulResponse[0].isDone());
    assertTrue(unsuccessfulResponse[0].isDone());
}

Example 4 with StateResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse in project beam by apache.

the class BeamFnStateGrpcClientCacheTest method testServerErrorCausesPendingAndFutureCallsToFail.

@Test
public void testServerErrorCausesPendingAndFutureCallsToFail() throws Exception {
    BeamFnStateClient client = clientCache.forApiServiceDescriptor(apiServiceDescriptor);
    CompletableFuture<StateResponse> inflight = client.handle(StateRequest.newBuilder().setInstructionId(SUCCESS));
    // Wait for the client to connect.
    StreamObserver<StateResponse> outboundServerObserver = outboundServerObservers.take();
    // Send an error from the server.
    outboundServerObserver.onError(new StatusRuntimeException(Status.INTERNAL.withDescription(SERVER_ERROR)));
    try {
        inflight.get();
        fail("Expected unsuccessful response due to server error");
    } catch (ExecutionException e) {
        assertThat(e.toString(), containsString(SERVER_ERROR));
    }
    try {
        inflight.get();
        fail("Expected unsuccessful response due to server error");
    } catch (ExecutionException e) {
        assertThat(e.toString(), containsString(SERVER_ERROR));
    }
}

Example 5 with StateResponse

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse in project beam by apache.

the class CommonCoderTest method instantiateCoder.

private static Coder<?> instantiateCoder(CommonCoder coder) {
    List<Coder<?>> components = new ArrayList<>();
    for (CommonCoder innerCoder : coder.getComponents()) {
        components.add(instantiateCoder(innerCoder));
    }
    // translator since we need to interact with a fake state client.
    if (coder.getUrn().equals(getUrn(StandardCoders.Enum.STATE_BACKED_ITERABLE))) {
        BeamFnStateClient stateClient = new BeamFnStateClient() {

            @Override
            public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
                checkState(requestBuilder.hasGet());
                checkState(requestBuilder.hasStateKey());
                checkState(requestBuilder.getStateKey().hasRunner());
                StateResponse.Builder rval = StateResponse.newBuilder();
                rval.setId(requestBuilder.getId());
                rval.setGet(StateGetResponse.newBuilder().setData(coder.getState().getOrDefault(requestBuilder.getStateKey().getRunner().getKey(), ByteString.EMPTY)));
                return CompletableFuture.completedFuture(rval.build());
            }
        };
        return new StateBackedIterable.Coder<>(() -> Caches.noop(), stateClient, () -> "instructionId", (Coder) Iterables.getOnlyElement(components));
    }
    Class<? extends Coder> coderType = ImmutableBiMap.copyOf(new ModelCoderRegistrar().getCoderURNs()).inverse().get(coder.getUrn());
    checkNotNull(coderType, "Unknown coder URN: " + coder.getUrn());
    CoderTranslator<?> translator = new ModelCoderRegistrar().getCoderTranslators().get(coderType);
    checkNotNull(translator, "No translator found for common coder class: " + coderType.getSimpleName());
    return translator.fromComponents(components, coder.getPayload(), new TranslationContext() {
    });
}