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

Example 1 with StateRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest 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 StateRequest

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest 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 StateRequest

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

the class MultimapUserState method startStateApiWrites.

@SuppressWarnings("FutureReturnValueIgnored")
private void startStateApiWrites() {
    // Clear currently persisted key-values
    if (isCleared) {
        beamFnStateClient.handle(keysStateRequest.toBuilder().setClear(StateClearRequest.getDefaultInstance()));
    } else if (!pendingRemoves.isEmpty()) {
        for (K key : pendingRemoves.values()) {
            StateRequest request = createUserStateRequest(key);
            beamFnStateClient.handle(request.toBuilder().setClear(StateClearRequest.getDefaultInstance()));
        }
    }
    // Persist pending key-values
    if (!pendingAdds.isEmpty()) {
        for (KV<K, List<V>> entry : pendingAdds.values()) {
            StateRequest request = createUserStateRequest(entry.getKey());
            beamFnStateClient.handle(request.toBuilder().setAppend(StateAppendRequest.newBuilder().setData(encodeValues(entry.getValue()))));
        }
    }
}

Example 4 with StateRequest

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

the class RegisterAndProcessBundleOperationTest method testProcessingBundleHandlesUserStateRequests.

@Test
public void testProcessingBundleHandlesUserStateRequests() throws Exception {
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    InMemoryStateInternals<ByteString> stateInternals = InMemoryStateInternals.forKey(ByteString.EMPTY);
    DataflowStepContext mockStepContext = mock(DataflowStepContext.class);
    DataflowStepContext mockUserStepContext = mock(DataflowStepContext.class);
    when(mockStepContext.namespacedToUser()).thenReturn(mockUserStepContext);
    when(mockUserStepContext.stateInternals()).thenReturn(stateInternals);
    InstructionRequestHandler instructionRequestHandler = new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            switch(request.getRequestCase()) {
                case REGISTER:
                    return CompletableFuture.completedFuture(responseFor(request).build());
                case PROCESS_BUNDLE:
                    return MoreFutures.supplyAsync(() -> {
                        StateRequest partialRequest = StateRequest.newBuilder().setStateKey(StateKey.newBuilder().setBagUserState(StateKey.BagUserState.newBuilder().setTransformId("testPTransformId").setWindow(ByteString.EMPTY).setUserStateId("testUserStateId"))).buildPartial();
                        StateRequest get = partialRequest.toBuilder().setGet(StateGetRequest.getDefaultInstance()).build();
                        StateRequest clear = partialRequest.toBuilder().setClear(StateClearRequest.getDefaultInstance()).build();
                        StateRequest append = partialRequest.toBuilder().setAppend(StateAppendRequest.newBuilder().setData(ByteString.copyFromUtf8("ABC"))).build();
                        StateRequestHandler stateHandler = stateHandlerCaptor.getValue();
                        StateResponse.Builder getWhenEmptyResponse = MoreFutures.get(stateHandler.handle(get));
                        assertEquals(ByteString.EMPTY, getWhenEmptyResponse.getGet().getData());
                        StateResponse.Builder appendWhenEmptyResponse = MoreFutures.get(stateHandler.handle(append));
                        assertNotNull(appendWhenEmptyResponse);
                        StateResponse.Builder appendWhenEmptyResponse2 = MoreFutures.get(stateHandler.handle(append));
                        assertNotNull(appendWhenEmptyResponse2);
                        StateResponse.Builder getWhenHasValueResponse = MoreFutures.get(stateHandler.handle(get));
                        assertEquals(ByteString.copyFromUtf8("ABC").concat(ByteString.copyFromUtf8("ABC")), getWhenHasValueResponse.getGet().getData());
                        StateResponse.Builder clearResponse = MoreFutures.get(stateHandler.handle(clear));
                        assertNotNull(clearResponse);
                        return responseFor(request).build();
                    });
                default:
                    // block forever
                    return new CompletableFuture<>();
            }
        }
    };
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, instructionRequestHandler, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of("testPTransformId", mockStepContext), ImmutableMap.of(), ImmutableTable.of(), ImmutableMap.of(), mockContext);
    operation.start();
    verify(mockBeamFnStateDelegator).registerForProcessBundleInstructionId(eq("778"), stateHandlerCaptor.capture());
    // This method blocks till the requests are completed
    operation.finish();
    // Ensure that the number of reigstrations matches the number of deregistrations
    assertEquals(stateServiceRegisterCounter.get(), stateServiceDeregisterCounter.get());
    assertEquals(0, stateServiceAbortCounter.get());
}

Example 5 with StateRequest

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

the class RegisterAndProcessBundleOperationTest method testProcessingBundleHandlesMultimapSideInputRequests.

@Test
public void testProcessingBundleHandlesMultimapSideInputRequests() throws Exception {
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    DataflowStepContext mockStepContext = mock(DataflowStepContext.class);
    DataflowStepContext mockUserStepContext = mock(DataflowStepContext.class);
    when(mockStepContext.namespacedToUser()).thenReturn(mockUserStepContext);
    CountDownLatch waitForStateHandler = new CountDownLatch(1);
    // Issues state calls to the Runner after a process bundle request is sent.
    InstructionRequestHandler fakeClient = new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            switch(request.getRequestCase()) {
                case REGISTER:
                    return CompletableFuture.completedFuture(responseFor(request).build());
                case PROCESS_BUNDLE:
                    return MoreFutures.supplyAsync(() -> {
                        StateKey getKey = StateKey.newBuilder().setMultimapSideInput(StateKey.MultimapSideInput.newBuilder().setTransformId("testPTransformId").setSideInputId("testSideInputId").setWindow(ByteString.copyFrom(CoderUtils.encodeToByteArray(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE))).setKey(ByteString.copyFrom(CoderUtils.encodeToByteArray(ByteArrayCoder.of(), "ABC".getBytes(StandardCharsets.UTF_8), Coder.Context.NESTED)))).build();
                        StateRequest getRequest = StateRequest.newBuilder().setStateKey(getKey).setGet(StateGetRequest.getDefaultInstance()).build();
                        waitForStateHandler.await();
                        StateRequestHandler stateHandler = stateHandlerCaptor.getValue();
                        StateResponse.Builder getResponse = MoreFutures.get(stateHandler.handle(getRequest));
                        assertEquals(encodeAndConcat(Arrays.asList("X", "Y", "Z"), StringUtf8Coder.of()), getResponse.getGet().getData());
                        return responseFor(request).build();
                    });
                default:
                    // block forever on other request types
                    return new CompletableFuture<>();
            }
        }
    };
    SideInputReader fakeSideInputReader = new SideInputReader() {

        @Override
        @Nullable
        public <T> T get(PCollectionView<T> view, BoundedWindow window) {
            assertEquals(GlobalWindow.INSTANCE, window);
            assertEquals("testSideInputId", view.getTagInternal().getId());
            return (T) InMemoryMultimapSideInputView.fromIterable(ByteArrayCoder.of(), ImmutableList.of(KV.of("ABC".getBytes(StandardCharsets.UTF_8), "X"), KV.of("ABC".getBytes(StandardCharsets.UTF_8), "Y"), KV.of("ABC".getBytes(StandardCharsets.UTF_8), "Z")));
        }

        @Override
        public <T> boolean contains(PCollectionView<T> view) {
            return "testSideInputId".equals(view.getTagInternal().getId());
        }

        @Override
        public boolean isEmpty() {
            return false;
        }
    };
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, fakeClient, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of("testPTransformId", mockStepContext), ImmutableMap.of("testPTransformId", fakeSideInputReader), ImmutableTable.of("testPTransformId", "testSideInputId", DataflowPortabilityPCollectionView.with(new TupleTag<>("testSideInputId"), FullWindowedValueCoder.of(KvCoder.of(ByteArrayCoder.of(), StringUtf8Coder.of()), GlobalWindow.Coder.INSTANCE))), ImmutableMap.of(), mockContext);
    operation.start();
    verify(mockBeamFnStateDelegator).registerForProcessBundleInstructionId(eq("778"), stateHandlerCaptor.capture());
    waitForStateHandler.countDown();
    // This method blocks till the requests are completed
    operation.finish();
    // Ensure that the number of reigstrations matches the number of deregistrations
    assertEquals(stateServiceRegisterCounter.get(), stateServiceDeregisterCounter.get());
    assertEquals(0, stateServiceAbortCounter.get());
}