| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.dataflow.worker.fn.control; |
| |
| import static org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation.encodeAndConcat; |
| import static org.hamcrest.Matchers.empty; |
| import static org.hamcrest.Matchers.equalTo; |
| import static org.hamcrest.Matchers.not; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertSame; |
| import static org.junit.Assert.assertThat; |
| import static org.mockito.Matchers.any; |
| import static org.mockito.Matchers.eq; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.verify; |
| import static org.mockito.Mockito.when; |
| |
| import java.nio.charset.StandardCharsets; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.List; |
| import java.util.concurrent.Callable; |
| import java.util.concurrent.CompletableFuture; |
| import java.util.concurrent.CompletionStage; |
| import java.util.concurrent.CountDownLatch; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.atomic.AtomicInteger; |
| import java.util.concurrent.atomic.AtomicLong; |
| import java.util.concurrent.atomic.AtomicReference; |
| import javax.annotation.Nullable; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest.RequestCase; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleProgressResponse; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateAppendRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateClearRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateGetRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse; |
| import org.apache.beam.runners.core.InMemoryMultimapSideInputView; |
| import org.apache.beam.runners.core.InMemoryStateInternals; |
| import org.apache.beam.runners.core.SideInputReader; |
| import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext; |
| import org.apache.beam.runners.dataflow.worker.DataflowPortabilityPCollectionView; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.OperationContext; |
| import org.apache.beam.runners.fnexecution.control.InstructionRequestHandler; |
| import org.apache.beam.runners.fnexecution.state.StateDelegator; |
| import org.apache.beam.runners.fnexecution.state.StateRequestHandler; |
| import org.apache.beam.sdk.coders.ByteArrayCoder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.coders.StringUtf8Coder; |
| import org.apache.beam.sdk.fn.IdGenerator; |
| import org.apache.beam.sdk.fn.IdGenerators; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.transforms.windowing.GlobalWindow; |
| import org.apache.beam.sdk.util.CoderUtils; |
| import org.apache.beam.sdk.util.MoreFutures; |
| import org.apache.beam.sdk.util.ThrowingRunnable; |
| import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollectionView; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.ValueInSingleWindow.Coder; |
| import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableTable; |
| import org.junit.Before; |
| import org.junit.Test; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| import org.mockito.ArgumentCaptor; |
| import org.mockito.Captor; |
| import org.mockito.Mock; |
| import org.mockito.MockitoAnnotations; |
| import org.mockito.invocation.InvocationOnMock; |
| import org.mockito.stubbing.Answer; |
| |
| /** Tests for {@link RegisterAndProcessBundleOperation}. */ |
| @RunWith(JUnit4.class) |
| @SuppressWarnings("FutureReturnValueIgnored") |
| public class RegisterAndProcessBundleOperationTest { |
| private static final BeamFnApi.RegisterRequest REGISTER_REQUEST = |
| BeamFnApi.RegisterRequest.newBuilder() |
| .addProcessBundleDescriptor(BeamFnApi.ProcessBundleDescriptor.newBuilder().setId("555")) |
| .build(); |
| @Mock private OperationContext mockContext; |
| @Mock private StateDelegator mockBeamFnStateDelegator; |
| |
| @Captor private ArgumentCaptor<StateRequestHandler> stateHandlerCaptor; |
| |
| private AtomicInteger stateServiceRegisterCounter; |
| private AtomicInteger stateServiceDeregisterCounter; |
| private AtomicInteger stateServiceAbortCounter; |
| |
| @Before |
| public void setUp() { |
| MockitoAnnotations.initMocks(this); |
| stateServiceRegisterCounter = new AtomicInteger(); |
| stateServiceDeregisterCounter = new AtomicInteger(); |
| stateServiceAbortCounter = new AtomicInteger(); |
| when(mockBeamFnStateDelegator.registerForProcessBundleInstructionId( |
| any(String.class), any(StateRequestHandler.class))) |
| .thenAnswer( |
| new Answer<StateDelegator.Registration>() { |
| @Override |
| public StateDelegator.Registration answer(InvocationOnMock invocationOnMock) |
| throws Throwable { |
| stateServiceRegisterCounter.incrementAndGet(); |
| return new StateDelegator.Registration() { |
| |
| @Override |
| public void deregister() { |
| stateServiceDeregisterCounter.incrementAndGet(); |
| } |
| |
| @Override |
| public void abort() { |
| stateServiceAbortCounter.incrementAndGet(); |
| } |
| }; |
| } |
| }); |
| } |
| |
| private IdGenerator makeIdGeneratorStartingFrom(long initialValue) { |
| return new IdGenerator() { |
| AtomicLong longs = new AtomicLong(initialValue); |
| |
| @Override |
| public String getId() { |
| return Long.toString(longs.getAndIncrement()); |
| } |
| }; |
| } |
| |
| @Test |
| public void testSupportsRestart() { |
| new RegisterAndProcessBundleOperation( |
| IdGenerators.decrementingLongs(), |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>(); |
| completeFuture(request, responseFuture); |
| return responseFuture; |
| } |
| |
| @Override |
| public void close() {} |
| }, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext) |
| .supportsRestart(); |
| } |
| |
| @Test |
| public void testRegisterOnlyOnFirstBundle() throws Exception { |
| List<BeamFnApi.InstructionRequest> requests = new ArrayList<>(); |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| idGenerator, |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| requests.add(request); |
| switch (request.getRequestCase()) { |
| case REGISTER: |
| case PROCESS_BUNDLE: |
| return CompletableFuture.completedFuture(responseFor(request).build()); |
| default: |
| // block forever on other requests |
| return new CompletableFuture<>(); |
| } |
| } |
| |
| @Override |
| public void close() {} |
| }, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| |
| // Ensure that the first time we start we send the register and process bundle requests |
| assertThat(requests, empty()); |
| operation.start(); |
| assertEquals( |
| requests.get(0), |
| BeamFnApi.InstructionRequest.newBuilder() |
| .setInstructionId("777") |
| .setRegister(REGISTER_REQUEST) |
| .build()); |
| assertEquals( |
| requests.get(1), |
| BeamFnApi.InstructionRequest.newBuilder() |
| .setInstructionId("778") |
| .setProcessBundle( |
| BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("555")) |
| .build()); |
| operation.finish(); |
| |
| // Ensure on restart that we only send the process bundle request |
| operation.start(); |
| assertEquals( |
| requests.get(2), |
| BeamFnApi.InstructionRequest.newBuilder() |
| .setInstructionId("779") |
| .setProcessBundle( |
| BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("555")) |
| .build()); |
| operation.finish(); |
| } |
| |
| @Test |
| public void testTentativeUserMetrics() throws Exception { |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| |
| CountDownLatch processBundleLatch = new CountDownLatch(1); |
| |
| final String stepName = "fakeStepNameWithUserMetrics"; |
| final String namespace = "sdk/whatever"; |
| final String name = "someCounter"; |
| final long counterValue = 42; |
| |
| final BeamFnApi.Metrics.User.MetricName metricName = |
| BeamFnApi.Metrics.User.MetricName.newBuilder() |
| .setNamespace(namespace) |
| .setName(name) |
| .build(); |
| |
| InstructionRequestHandler instructionRequestHandler = |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| switch (request.getRequestCase()) { |
| case REGISTER: |
| return CompletableFuture.completedFuture(responseFor(request).build()); |
| case PROCESS_BUNDLE: |
| return MoreFutures.supplyAsync( |
| () -> { |
| processBundleLatch.await(); |
| return responseFor(request).build(); |
| }); |
| case PROCESS_BUNDLE_PROGRESS: |
| return CompletableFuture.completedFuture( |
| responseFor(request) |
| .setProcessBundleProgress( |
| BeamFnApi.ProcessBundleProgressResponse.newBuilder() |
| .setMetrics( |
| BeamFnApi.Metrics.newBuilder() |
| .putPtransforms( |
| stepName, |
| BeamFnApi.Metrics.PTransform.newBuilder() |
| .addUser( |
| BeamFnApi.Metrics.User.newBuilder() |
| .setMetricName(metricName) |
| .setCounterData( |
| BeamFnApi.Metrics.User.CounterData |
| .newBuilder() |
| .setValue(counterValue))) |
| .build()))) |
| .build()); |
| default: |
| // block forever |
| return new CompletableFuture<>(); |
| } |
| } |
| |
| @Override |
| public void close() {} |
| }; |
| |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| idGenerator, |
| instructionRequestHandler, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| |
| operation.start(); |
| |
| BeamFnApi.Metrics metrics = MoreFutures.get(operation.getProcessBundleProgress()).getMetrics(); |
| assertThat(metrics.getPtransformsOrThrow(stepName).getUserCount(), equalTo(1)); |
| |
| BeamFnApi.Metrics.User userMetric = metrics.getPtransformsOrThrow(stepName).getUser(0); |
| assertThat(userMetric.getMetricName(), equalTo(metricName)); |
| assertThat(userMetric.getCounterData().getValue(), equalTo(counterValue)); |
| |
| processBundleLatch.countDown(); |
| operation.finish(); |
| } |
| |
| @Test |
| public void testFinalUserMetrics() throws Exception { |
| List<BeamFnApi.InstructionRequest> requests = new ArrayList<>(); |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| ExecutorService executorService = Executors.newCachedThreadPool(); |
| |
| CountDownLatch processBundleLatch = new CountDownLatch(1); |
| |
| final String stepName = "fakeStepNameWithUserMetrics"; |
| final String namespace = "sdk/whatever"; |
| final String name = "someCounter"; |
| final long counterValue = 42; |
| final long finalCounterValue = 77; |
| |
| final BeamFnApi.Metrics.User.MetricName metricName = |
| BeamFnApi.Metrics.User.MetricName.newBuilder() |
| .setNamespace(namespace) |
| .setName(name) |
| .build(); |
| |
| InstructionRequestHandler instructionRequestHandler = |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| switch (request.getRequestCase()) { |
| case REGISTER: |
| return CompletableFuture.completedFuture(responseFor(request).build()); |
| case PROCESS_BUNDLE: |
| return MoreFutures.supplyAsync( |
| () -> { |
| processBundleLatch.await(); |
| return responseFor(request) |
| .setProcessBundle( |
| BeamFnApi.ProcessBundleResponse.newBuilder() |
| .setMetrics( |
| BeamFnApi.Metrics.newBuilder() |
| .putPtransforms( |
| stepName, |
| BeamFnApi.Metrics.PTransform.newBuilder() |
| .addUser( |
| BeamFnApi.Metrics.User.newBuilder() |
| .setMetricName(metricName) |
| .setCounterData( |
| BeamFnApi.Metrics.User.CounterData |
| .newBuilder() |
| .setValue(finalCounterValue))) |
| .build()))) |
| .build(); |
| }); |
| case PROCESS_BUNDLE_PROGRESS: |
| return CompletableFuture.completedFuture( |
| responseFor(request) |
| .setProcessBundleProgress( |
| BeamFnApi.ProcessBundleProgressResponse.newBuilder() |
| .setMetrics( |
| BeamFnApi.Metrics.newBuilder() |
| .putPtransforms( |
| stepName, |
| BeamFnApi.Metrics.PTransform.newBuilder() |
| .addUser( |
| BeamFnApi.Metrics.User.newBuilder() |
| .setMetricName(metricName) |
| .setCounterData( |
| BeamFnApi.Metrics.User.CounterData |
| .newBuilder() |
| .setValue(counterValue))) |
| .build()))) |
| .build()); |
| default: |
| // block forever |
| return new CompletableFuture<>(); |
| } |
| } |
| |
| @Override |
| public void close() {} |
| }; |
| |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| idGenerator, |
| instructionRequestHandler, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| |
| operation.start(); |
| |
| // Force some intermediate metrics to test crosstalk is not introduced |
| BeamFnApi.Metrics metrics = MoreFutures.get(operation.getProcessBundleProgress()).getMetrics(); |
| BeamFnApi.Metrics.User userMetric = metrics.getPtransformsOrThrow(stepName).getUser(0); |
| assertThat(userMetric.getMetricName(), equalTo(metricName)); |
| assertThat(userMetric.getCounterData().getValue(), not(equalTo(finalCounterValue))); |
| |
| processBundleLatch.countDown(); |
| operation.finish(); |
| |
| metrics = MoreFutures.get(operation.getFinalMetrics()); |
| |
| userMetric = metrics.getPtransformsOrThrow(stepName).getUser(0); |
| assertThat(userMetric.getMetricName(), equalTo(metricName)); |
| assertThat(userMetric.getCounterData().getValue(), equalTo(finalCounterValue)); |
| } |
| |
| @Test |
| public void testProcessingBundleBlocksOnFinish() throws Exception { |
| List<BeamFnApi.InstructionRequest> requests = new ArrayList<>(); |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| ExecutorService executorService = Executors.newCachedThreadPool(); |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| idGenerator, |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| requests.add(request); |
| switch (request.getRequestCase()) { |
| case REGISTER: |
| return CompletableFuture.completedFuture(responseFor(request).build()); |
| case PROCESS_BUNDLE: |
| CompletableFuture<InstructionResponse> responseFuture = |
| new CompletableFuture<>(); |
| executorService.submit( |
| () -> { |
| // Purposefully sleep simulating SDK harness doing work |
| Thread.sleep(100); |
| responseFuture.complete(responseFor(request).build()); |
| completeFuture(request, responseFuture); |
| return null; |
| }); |
| return responseFuture; |
| default: |
| // Anything else hangs; nothing else should be blocking |
| return new CompletableFuture<>(); |
| } |
| } |
| |
| @Override |
| public void close() {} |
| }, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| |
| operation.start(); |
| // This method blocks till the requests are completed |
| operation.finish(); |
| |
| // Ensure that the messages were received |
| assertEquals( |
| requests.get(0), |
| BeamFnApi.InstructionRequest.newBuilder() |
| .setInstructionId("777") |
| .setRegister(REGISTER_REQUEST) |
| .build()); |
| assertEquals( |
| requests.get(1), |
| BeamFnApi.InstructionRequest.newBuilder() |
| .setInstructionId("778") |
| .setProcessBundle( |
| BeamFnApi.ProcessBundleRequest.newBuilder().setProcessBundleDescriptorId("555")) |
| .build()); |
| } |
| |
| @Test |
| public void testProcessingBundleHandlesUserStateRequests() throws Exception { |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| ExecutorService executorService = Executors.newCachedThreadPool(); |
| |
| 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 InstructionRequestHandler() { |
| @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<>(); |
| } |
| } |
| |
| @Override |
| public void close() {} |
| }; |
| |
| 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()); |
| } |
| |
| @Test |
| public void testProcessingBundleHandlesMultimapSideInputRequests() throws Exception { |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| ExecutorService executorService = Executors.newCachedThreadPool(); |
| |
| 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 InstructionRequestHandler() { |
| @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<>(); |
| } |
| } |
| |
| @Override |
| public void close() {} |
| }; |
| |
| SideInputReader fakeSideInputReader = |
| new SideInputReader() { |
| @Nullable |
| @Override |
| 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()); |
| } |
| |
| @Test |
| public void testAbortCancelsAndCleansUpDuringRegister() throws Exception { |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| ExecutorService executorService = Executors.newCachedThreadPool(); |
| |
| CountDownLatch waitForAbortToComplete = new CountDownLatch(1); |
| AtomicReference<ThrowingRunnable> abortReference = new AtomicReference<>(); |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| idGenerator, |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>(); |
| if (request.getRequestCase() == RequestCase.PROCESS_BUNDLE) { |
| executorService.submit( |
| (Callable<Void>) |
| () -> { |
| abortReference.get().run(); |
| waitForAbortToComplete.countDown(); |
| return null; |
| }); |
| } else { |
| completeFuture(request, responseFuture); |
| } |
| return responseFuture; |
| } |
| |
| @Override |
| public void close() {} |
| }, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| abortReference.set(operation::abort); |
| operation.start(); |
| |
| waitForAbortToComplete.await(); |
| // Ensure that the number of registrations matches the number of aborts |
| assertEquals(stateServiceRegisterCounter.get(), stateServiceAbortCounter.get()); |
| assertEquals(0, stateServiceDeregisterCounter.get()); |
| } |
| |
| @Test |
| public void testAbortCancelsAndCleansUpDuringProcessBundle() throws Exception { |
| IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L); |
| ExecutorService executorService = Executors.newCachedThreadPool(); |
| |
| CountDownLatch waitForAbortToComplete = new CountDownLatch(1); |
| AtomicReference<ThrowingRunnable> abortReference = new AtomicReference<>(); |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| idGenerator, |
| new InstructionRequestHandler() { |
| @Override |
| public CompletionStage<InstructionResponse> handle(InstructionRequest request) { |
| CompletableFuture<InstructionResponse> responseFuture = new CompletableFuture<>(); |
| if (request.getRequestCase() == RequestCase.PROCESS_BUNDLE) { |
| executorService.submit( |
| (Callable<Void>) |
| () -> { |
| abortReference.get().run(); |
| waitForAbortToComplete.countDown(); |
| return null; |
| }); |
| } else { |
| completeFuture(request, responseFuture); |
| } |
| return responseFuture; |
| } |
| |
| @Override |
| public void close() {} |
| }, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| abortReference.set(operation::abort); |
| operation.start(); |
| |
| waitForAbortToComplete.await(); |
| // Ensure that the number of registrations matches the number of aborts |
| assertEquals(stateServiceRegisterCounter.get(), stateServiceAbortCounter.get()); |
| assertEquals(0, stateServiceDeregisterCounter.get()); |
| } |
| |
| private InstructionResponse.Builder responseFor(BeamFnApi.InstructionRequest request) { |
| BeamFnApi.InstructionResponse.Builder response = |
| BeamFnApi.InstructionResponse.newBuilder().setInstructionId(request.getInstructionId()); |
| if (request.hasRegister()) { |
| response.setRegister(BeamFnApi.RegisterResponse.getDefaultInstance()); |
| } else if (request.hasProcessBundle()) { |
| response.setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance()); |
| } else if (request.hasFinalizeBundle()) { |
| response.setFinalizeBundle(BeamFnApi.FinalizeBundleResponse.getDefaultInstance()); |
| } else if (request.hasProcessBundleProgress()) { |
| response.setProcessBundleProgress( |
| BeamFnApi.ProcessBundleProgressResponse.getDefaultInstance()); |
| } else if (request.hasProcessBundleSplit()) { |
| response.setProcessBundleSplit(BeamFnApi.ProcessBundleSplitResponse.getDefaultInstance()); |
| } |
| return response; |
| } |
| |
| private void completeFuture( |
| BeamFnApi.InstructionRequest request, CompletableFuture<InstructionResponse> response) { |
| response.complete(responseFor(request).build()); |
| } |
| |
| @Test |
| public void testGetProcessBundleProgressReturnsDefaultInstanceIfNoBundleIdCached() |
| throws Exception { |
| InstructionRequestHandler mockInstructionRequestHandler = mock(InstructionRequestHandler.class); |
| |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| IdGenerators.decrementingLongs(), |
| mockInstructionRequestHandler, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| |
| assertEquals( |
| ProcessBundleProgressResponse.getDefaultInstance(), |
| MoreFutures.get(operation.getProcessBundleProgress())); |
| } |
| |
| @Test |
| public void testGetProcessBundleProgressFetchesProgressResponseWhenBundleIdCached() |
| throws Exception { |
| InstructionRequestHandler mockInstructionRequestHandler = mock(InstructionRequestHandler.class); |
| |
| RegisterAndProcessBundleOperation operation = |
| new RegisterAndProcessBundleOperation( |
| IdGenerators.decrementingLongs(), |
| mockInstructionRequestHandler, |
| mockBeamFnStateDelegator, |
| REGISTER_REQUEST, |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableMap.of(), |
| ImmutableTable.of(), |
| ImmutableMap.of(), |
| mockContext); |
| |
| operation.getProcessBundleInstructionId(); // this generates and caches bundleId |
| |
| ProcessBundleProgressResponse expectedResult = |
| ProcessBundleProgressResponse.newBuilder().build(); |
| InstructionResponse instructionResponse = |
| InstructionResponse.newBuilder().setProcessBundleProgress(expectedResult).build(); |
| CompletableFuture resultFuture = CompletableFuture.completedFuture(instructionResponse); |
| when(mockInstructionRequestHandler.handle(any())).thenReturn(resultFuture); |
| |
| final ProcessBundleProgressResponse result = |
| MoreFutures.get(operation.getProcessBundleProgress()); |
| |
| assertSame("Return value from mockInstructionRequestHandler", expectedResult, result); |
| } |
| } |
