| /* |
| * 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.vendor.guava.v20_0.com.google.common.base.Preconditions.checkState; |
| |
| import java.io.Serializable; |
| import java.util.HashMap; |
| import java.util.Map; |
| import java.util.concurrent.ExecutionException; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.Future; |
| import java.util.concurrent.ThreadFactory; |
| import javax.annotation.Nullable; |
| import org.apache.beam.fn.harness.FnHarness; |
| import org.apache.beam.model.fnexecution.v1.BeamFnApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.runners.core.InMemoryTimerInternals; |
| import org.apache.beam.runners.core.StateInternals; |
| import org.apache.beam.runners.core.TimerInternals; |
| import org.apache.beam.runners.core.construction.PipelineTranslation; |
| import org.apache.beam.runners.core.construction.graph.ExecutableStage; |
| import org.apache.beam.runners.core.construction.graph.FusedPipeline; |
| import org.apache.beam.runners.core.construction.graph.GreedyPipelineFuser; |
| import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext; |
| import org.apache.beam.runners.dataflow.worker.counters.NameContext; |
| import org.apache.beam.runners.fnexecution.GrpcContextHeaderAccessorProvider; |
| import org.apache.beam.runners.fnexecution.GrpcFnServer; |
| import org.apache.beam.runners.fnexecution.InProcessServerFactory; |
| import org.apache.beam.runners.fnexecution.control.*; |
| import org.apache.beam.runners.fnexecution.data.GrpcDataService; |
| import org.apache.beam.runners.fnexecution.logging.GrpcLoggingService; |
| import org.apache.beam.runners.fnexecution.logging.Slf4jLogWriter; |
| import org.apache.beam.runners.fnexecution.provisioning.JobInfo; |
| import org.apache.beam.runners.fnexecution.state.GrpcStateService; |
| import org.apache.beam.runners.fnexecution.state.StateRequestHandler; |
| import org.apache.beam.sdk.Pipeline; |
| import org.apache.beam.sdk.coders.*; |
| import org.apache.beam.sdk.fn.data.FnDataReceiver; |
| import org.apache.beam.sdk.fn.stream.OutboundObserverFactory; |
| import org.apache.beam.sdk.fn.test.InProcessManagedChannelFactory; |
| import org.apache.beam.sdk.options.PipelineOptionsFactory; |
| import org.apache.beam.sdk.state.TimeDomain; |
| import org.apache.beam.sdk.state.Timer; |
| import org.apache.beam.sdk.state.TimerSpec; |
| import org.apache.beam.sdk.state.TimerSpecs; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.Impulse; |
| import org.apache.beam.sdk.transforms.ParDo; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.transforms.windowing.GlobalWindow; |
| import org.apache.beam.sdk.transforms.windowing.PaneInfo; |
| import org.apache.beam.sdk.util.CoderUtils; |
| import org.apache.beam.sdk.util.WindowedValue; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.vendor.grpc.v1p13p1.com.google.protobuf.Struct; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.base.Optional; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.util.concurrent.ThreadFactoryBuilder; |
| import org.joda.time.Duration; |
| import org.joda.time.Instant; |
| import org.junit.After; |
| import org.junit.Before; |
| import org.junit.Test; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| import org.mockito.Mockito; |
| |
| @RunWith(JUnit4.class) |
| public class TimerReceiverTest implements Serializable { |
| private transient GrpcFnServer<FnApiControlClientPoolService> controlServer; |
| private transient GrpcFnServer<GrpcDataService> dataServer; |
| private transient GrpcFnServer<GrpcLoggingService> loggingServer; |
| private transient GrpcStateService stateDelegator; |
| private transient SdkHarnessClient client; |
| private transient ExecutorService sdkHarnessExecutor; |
| private transient Future<?> sdkHarnessExecutorFuture; |
| |
| @Before |
| public void setUp() throws Exception { |
| ThreadFactory threadFactory = new ThreadFactoryBuilder().setDaemon(true).build(); |
| ExecutorService serverExecutor = Executors.newCachedThreadPool(threadFactory); |
| InProcessServerFactory serverFactory = InProcessServerFactory.create(); |
| dataServer = |
| GrpcFnServer.allocatePortAndCreateFor( |
| GrpcDataService.create(serverExecutor, OutboundObserverFactory.serverDirect()), |
| serverFactory); |
| loggingServer = |
| GrpcFnServer.allocatePortAndCreateFor( |
| GrpcLoggingService.forWriter(Slf4jLogWriter.getDefault()), serverFactory); |
| stateDelegator = GrpcStateService.create(); |
| ControlClientPool clientPool = MapControlClientPool.create(); |
| controlServer = |
| GrpcFnServer.allocatePortAndCreateFor( |
| FnApiControlClientPoolService.offeringClientsToPool( |
| clientPool.getSink(), GrpcContextHeaderAccessorProvider.getHeaderAccessor()), |
| serverFactory); |
| |
| // Create the SDK harness, and wait until it connects |
| sdkHarnessExecutor = Executors.newSingleThreadExecutor(threadFactory); |
| sdkHarnessExecutorFuture = |
| sdkHarnessExecutor.submit( |
| () -> { |
| try { |
| FnHarness.main( |
| "id", |
| PipelineOptionsFactory.create(), |
| loggingServer.getApiServiceDescriptor(), |
| controlServer.getApiServiceDescriptor(), |
| InProcessManagedChannelFactory.create(), |
| OutboundObserverFactory.clientDirect()); |
| } catch (Exception e) { |
| throw new RuntimeException(e); |
| } |
| }); |
| |
| InstructionRequestHandler controlClient = |
| clientPool.getSource().take("", java.time.Duration.ofSeconds(2)); |
| client = SdkHarnessClient.usingFnApiClient(controlClient, dataServer.getService()); |
| } |
| |
| @After |
| public void tearDown() throws Exception { |
| controlServer.close(); |
| dataServer.close(); |
| loggingServer.close(); |
| sdkHarnessExecutor.shutdownNow(); |
| try { |
| sdkHarnessExecutorFuture.get(); |
| } catch (ExecutionException e) { |
| if (e.getCause() instanceof RuntimeException |
| && e.getCause().getCause() instanceof InterruptedException) { |
| // expected |
| } else { |
| throw e; |
| } |
| } |
| } |
| |
| /* |
| Tests that we can schedule a single timer to fire, and that it fires. |
| */ |
| @Test |
| public void testSingleTimerScheduling() throws Exception { |
| final String timerId = "timerId"; |
| |
| Pipeline p = Pipeline.create(); |
| PCollection<Integer> output = |
| p.apply("impulse", Impulse.create()) |
| .apply( |
| "create", |
| ParDo.of( |
| new DoFn<byte[], KV<String, Integer>>() { |
| @ProcessElement |
| public void process(ProcessContext ctxt) {} |
| })) |
| .apply( |
| "timer", |
| ParDo.of( |
| new DoFn<KV<String, Integer>, Integer>() { |
| @TimerId(timerId) |
| private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME); |
| |
| @ProcessElement |
| public void processElement( |
| @TimerId(timerId) Timer timer, OutputReceiver<Integer> r) { |
| timer.offset(Duration.standardSeconds(1)).setRelative(); |
| } |
| |
| @OnTimer(timerId) |
| public void onTimer( |
| @TimerId(timerId) Timer timer, |
| TimeDomain timeDomain, |
| OutputReceiver<Integer> r) { |
| r.output(0); |
| } |
| })); |
| |
| RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p); |
| FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto); |
| Optional<ExecutableStage> optionalStage = |
| Iterables.tryFind( |
| fused.getFusedStages(), (ExecutableStage stage) -> !stage.getTimers().isEmpty()); |
| checkState(optionalStage.isPresent(), "Expected a stage with timers."); |
| ExecutableStage stage = optionalStage.get(); |
| |
| ProcessBundleDescriptors.ExecutableProcessBundleDescriptor descriptor = |
| ProcessBundleDescriptors.fromExecutableStage( |
| "test_stage", stage, dataServer.getApiServiceDescriptor()); |
| |
| TimerReceiver timerReceiver = |
| Mockito.spy( |
| new TimerReceiver( |
| stage.getComponents(), |
| buildDataflowStepContext(), |
| buildStageBundleFactory(client, descriptor, stateDelegator))); |
| |
| Map<String, ProcessBundleDescriptors.TimerSpec> timerSpecMap = new HashMap<>(); |
| descriptor |
| .getTimerSpecs() |
| .values() |
| .forEach( |
| transformTimerMap -> { |
| for (ProcessBundleDescriptors.TimerSpec timerSpec : transformTimerMap.values()) { |
| timerSpecMap.put(timerSpec.timerId(), timerSpec); |
| } |
| }); |
| |
| String timerOutputPCollection = timerSpecMap.get(timerId).outputCollectionId(); |
| String timerInputPCollection = timerSpecMap.get(timerId).inputCollectionId(); |
| |
| // Arbitrary offset. |
| long testTimerOffset = 123456; |
| // Arbitrary key. |
| Object timer = timerBytes("X", testTimerOffset); |
| Object windowedTimer = WindowedValue.valueInGlobalWindow(timer); |
| |
| // Simulate the SDK Harness sending a timer element to the Runner Harness. |
| org.junit.Assert.assertTrue(timerReceiver.receive(timerOutputPCollection, windowedTimer)); |
| |
| // Expect that we get a timer element when we finish. |
| Object expected = |
| WindowedValue.of( |
| timer, |
| BoundedWindow.TIMESTAMP_MIN_VALUE.plus(testTimerOffset), |
| GlobalWindow.INSTANCE, |
| PaneInfo.NO_FIRING); |
| |
| Mockito.verify(timerReceiver, Mockito.never()) |
| .fireTimer( |
| timerInputPCollection, |
| (WindowedValue<KV<Object, org.apache.beam.runners.core.construction.Timer>>) expected); |
| |
| // Simulate firing the timer. Expect that the fired timer is exactly the one we received |
| // originally (with additional details). |
| timerReceiver.finish(); |
| Mockito.verify(timerReceiver) |
| .fireTimer( |
| timerInputPCollection, |
| (WindowedValue<KV<Object, org.apache.beam.runners.core.construction.Timer>>) expected); |
| } |
| |
| /* |
| Tests that we can schedule multiple timers to fire, and that they all fire. |
| */ |
| @Test |
| public void testMultiTimerScheduling() throws Exception { |
| final String timerId1 = "timerId1"; |
| final String timerId2 = "timerId2"; |
| |
| Pipeline p = Pipeline.create(); |
| PCollection<Integer> output = |
| p.apply("impulse", Impulse.create()) |
| .apply( |
| "create", |
| ParDo.of( |
| new DoFn<byte[], KV<String, Integer>>() { |
| @ProcessElement |
| public void process(ProcessContext ctxt) {} |
| })) |
| .apply( |
| "timer", |
| ParDo.of( |
| new DoFn<KV<String, Integer>, Integer>() { |
| @TimerId(timerId1) |
| private final TimerSpec timer1 = TimerSpecs.timer(TimeDomain.EVENT_TIME); |
| |
| @TimerId(timerId2) |
| private final TimerSpec timer2 = TimerSpecs.timer(TimeDomain.EVENT_TIME); |
| |
| @ProcessElement |
| public void processElement( |
| @TimerId(timerId1) Timer timer1, |
| @TimerId(timerId2) Timer timer2, |
| OutputReceiver<Integer> r) { |
| timer1.offset(Duration.standardSeconds(1)).setRelative(); |
| timer2.offset(Duration.standardSeconds(2)).setRelative(); |
| } |
| |
| @OnTimer(timerId1) |
| public void onTimer1( |
| @TimerId(timerId1) Timer timer, |
| TimeDomain timeDomain, |
| OutputReceiver<Integer> r) { |
| r.output(1); |
| } |
| |
| @OnTimer(timerId2) |
| public void onTimer2( |
| @TimerId(timerId2) Timer timer, |
| TimeDomain timeDomain, |
| OutputReceiver<Integer> r) { |
| r.output(2); |
| } |
| })); |
| |
| RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p); |
| FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto); |
| Optional<ExecutableStage> optionalStage = |
| Iterables.tryFind( |
| fused.getFusedStages(), (ExecutableStage stage) -> !stage.getTimers().isEmpty()); |
| checkState(optionalStage.isPresent(), "Expected a stage with timers."); |
| ExecutableStage stage = optionalStage.get(); |
| |
| ProcessBundleDescriptors.ExecutableProcessBundleDescriptor descriptor = |
| ProcessBundleDescriptors.fromExecutableStage( |
| "test_stage", stage, dataServer.getApiServiceDescriptor()); |
| |
| TimerReceiver timerReceiver = |
| Mockito.spy( |
| new TimerReceiver( |
| stage.getComponents(), |
| buildDataflowStepContext(), |
| buildStageBundleFactory(client, descriptor, stateDelegator))); |
| |
| Map<String, ProcessBundleDescriptors.TimerSpec> timerSpecMap = new HashMap<>(); |
| descriptor |
| .getTimerSpecs() |
| .values() |
| .forEach( |
| transformTimerMap -> { |
| for (ProcessBundleDescriptors.TimerSpec timerSpec : transformTimerMap.values()) { |
| timerSpecMap.put(timerSpec.timerId(), timerSpec); |
| } |
| }); |
| |
| // Arbitrary offset. |
| long testTimerOffset = 123456; |
| // Arbitrary key. |
| Object timer1 = timerBytes("X", testTimerOffset); |
| Object windowedTimer1 = WindowedValue.valueInGlobalWindow(timer1); |
| |
| Object timer2 = timerBytes("Y", testTimerOffset); |
| Object windowedTimer2 = WindowedValue.valueInGlobalWindow(timer2); |
| |
| // Simulate the SDK Harness sending a timer element to the Runner Harness. |
| org.junit.Assert.assertTrue( |
| timerReceiver.receive(timerSpecMap.get(timerId1).outputCollectionId(), windowedTimer1)); |
| org.junit.Assert.assertTrue( |
| timerReceiver.receive(timerSpecMap.get(timerId2).outputCollectionId(), windowedTimer2)); |
| |
| // Expect that we get a timer element when we finish. |
| Object expectedTimer1 = |
| WindowedValue.of( |
| timer1, |
| BoundedWindow.TIMESTAMP_MIN_VALUE.plus(testTimerOffset), |
| GlobalWindow.INSTANCE, |
| PaneInfo.NO_FIRING); |
| |
| Object expectedTimer2 = |
| WindowedValue.of( |
| timer2, |
| BoundedWindow.TIMESTAMP_MIN_VALUE.plus(testTimerOffset), |
| GlobalWindow.INSTANCE, |
| PaneInfo.NO_FIRING); |
| |
| Mockito.verify(timerReceiver, Mockito.never()) |
| .fireTimer( |
| timerSpecMap.get(timerId1).inputCollectionId(), |
| (WindowedValue<KV<Object, org.apache.beam.runners.core.construction.Timer>>) |
| expectedTimer1); |
| Mockito.verify(timerReceiver, Mockito.never()) |
| .fireTimer( |
| timerSpecMap.get(timerId2).inputCollectionId(), |
| (WindowedValue<KV<Object, org.apache.beam.runners.core.construction.Timer>>) |
| expectedTimer2); |
| |
| // Simulate firing the timer. Expect that the fired timer is exactly the one we received |
| // originally (with additional details). |
| timerReceiver.finish(); |
| Mockito.verify(timerReceiver) |
| .fireTimer( |
| timerSpecMap.get(timerId1).inputCollectionId(), |
| (WindowedValue<KV<Object, org.apache.beam.runners.core.construction.Timer>>) |
| expectedTimer1); |
| Mockito.verify(timerReceiver) |
| .fireTimer( |
| timerSpecMap.get(timerId2).inputCollectionId(), |
| (WindowedValue<KV<Object, org.apache.beam.runners.core.construction.Timer>>) |
| expectedTimer2); |
| } |
| |
| private static class SimpleStageBundleFactory implements StageBundleFactory { |
| private final SdkHarnessClient client; |
| private final SdkHarnessClient.BundleProcessor processor; |
| private final ProcessBundleDescriptors.ExecutableProcessBundleDescriptor |
| processBundleDescriptor; |
| |
| SimpleStageBundleFactory( |
| SdkHarnessClient client, |
| SdkHarnessClient.BundleProcessor processor, |
| ProcessBundleDescriptors.ExecutableProcessBundleDescriptor processBundleDescriptor) { |
| this.client = client; |
| this.processor = processor; |
| this.processBundleDescriptor = processBundleDescriptor; |
| } |
| |
| @Override |
| public RemoteBundle getBundle( |
| OutputReceiverFactory outputReceiverFactory, |
| StateRequestHandler stateRequestHandler, |
| BundleProgressHandler progressHandler) |
| throws Exception { |
| ImmutableMap.Builder<BeamFnApi.Target, RemoteOutputReceiver<?>> outputReceivers = |
| ImmutableMap.builder(); |
| for (Map.Entry<BeamFnApi.Target, Coder<WindowedValue<?>>> targetCoder : |
| processBundleDescriptor.getOutputTargetCoders().entrySet()) { |
| BeamFnApi.Target target = targetCoder.getKey(); |
| Coder<WindowedValue<?>> coder = targetCoder.getValue(); |
| String bundleOutputPCollection = |
| Iterables.getOnlyElement( |
| processBundleDescriptor |
| .getProcessBundleDescriptor() |
| .getTransformsOrThrow(target.getPrimitiveTransformReference()) |
| .getInputsMap() |
| .values()); |
| FnDataReceiver<WindowedValue<?>> outputReceiver = |
| outputReceiverFactory.create(bundleOutputPCollection); |
| outputReceivers.put(target, RemoteOutputReceiver.of(coder, outputReceiver)); |
| } |
| return processor.newBundle(outputReceivers.build(), stateRequestHandler, progressHandler); |
| } |
| |
| @Override |
| public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor getProcessBundleDescriptor() { |
| return processBundleDescriptor; |
| } |
| |
| @Override |
| public void close() throws Exception {} |
| } |
| |
| private static class TestStepContext extends DataflowExecutionContext.DataflowStepContext { |
| private InMemoryTimerInternals timerInternals; |
| |
| public TestStepContext(NameContext nameContext) { |
| super(nameContext); |
| |
| timerInternals = new InMemoryTimerInternals(); |
| } |
| |
| @Nullable |
| @Override |
| public <W extends BoundedWindow> TimerInternals.TimerData getNextFiredTimer( |
| Coder<W> windowCoder) { |
| try { |
| timerInternals.advanceInputWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE); |
| } catch (Exception e) { |
| throw new IllegalStateException("Exception thrown advancing watermark", e); |
| } |
| |
| return timerInternals.removeNextEventTimer(); |
| } |
| |
| @Override |
| public DataflowExecutionContext.DataflowStepContext namespacedToUser() { |
| return this; |
| } |
| |
| @Override |
| public TimerInternals timerInternals() { |
| return timerInternals; |
| } |
| |
| @Override |
| public <W extends BoundedWindow> void setStateCleanupTimer( |
| String timerId, W window, Coder<W> windowCoder, Instant cleanupTime) { |
| throw new UnsupportedOperationException(); |
| } |
| |
| @Override |
| public StateInternals stateInternals() { |
| throw new UnsupportedOperationException(); |
| } |
| } |
| |
| private KV<String, org.apache.beam.runners.core.construction.Timer<byte[]>> timerBytes( |
| String key, long timestampOffset) throws CoderException { |
| return KV.of( |
| key, |
| org.apache.beam.runners.core.construction.Timer.of( |
| BoundedWindow.TIMESTAMP_MIN_VALUE.plus(timestampOffset), |
| CoderUtils.encodeToByteArray(VoidCoder.of(), null, Coder.Context.NESTED))); |
| } |
| |
| private static DataflowExecutionContext.DataflowStepContext buildDataflowStepContext() { |
| return new TestStepContext(NameContext.create("", "", "", "")); |
| } |
| |
| private static StageBundleFactory buildStageBundleFactory( |
| SdkHarnessClient client, |
| ProcessBundleDescriptors.ExecutableProcessBundleDescriptor processBundleDescriptor, |
| GrpcStateService stateDelegator) { |
| return new SimpleStageBundleFactory( |
| client, |
| client.getProcessor( |
| processBundleDescriptor.getProcessBundleDescriptor(), |
| processBundleDescriptor.getRemoteInputDestinations(), |
| stateDelegator), |
| processBundleDescriptor); |
| } |
| |
| private static JobInfo jobInfo() { |
| return JobInfo.create("job_id", "job_name", "", Struct.newBuilder().build()); |
| } |
| } |
