Search in sources :

Example 21 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserDistributionMonitoringInfoToCounterUpdateTransformerTest method testTransformThrowsIfMonitoringInfoWithWrongUrnPrefixReceived.

@Test
public void testTransformThrowsIfMonitoringInfoWithWrongUrnPrefixReceived() {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.ELEMENT_COUNT).setType(TypeUrns.SUM_INT64_TYPE).build();
    UserDistributionMonitoringInfoToCounterUpdateTransformer testObject = new UserDistributionMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    exception.expect(RuntimeException.class);
    testObject.transform(monitoringInfo);
}
Also used : MonitoringInfo(org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo) HashMap(java.util.HashMap) DataflowStepContext(org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext) Test(org.junit.Test)

Example 22 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserDistributionMonitoringInfoToCounterUpdateTransformerTest method testTransformReturnsNullIfMonitoringInfoWithUnknownPTransformLabelPresent.

@Test
public void testTransformReturnsNullIfMonitoringInfoWithUnknownPTransformLabelPresent() {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.USER_DISTRIBUTION_INT64).setType(TypeUrns.DISTRIBUTION_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.NAME, "anyName").putLabels(MonitoringInfoConstants.Labels.NAMESPACE, "anyNamespace").putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "anyValue").build();
    UserDistributionMonitoringInfoToCounterUpdateTransformer testObject = new UserDistributionMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    assertNull(testObject.transform(monitoringInfo));
}
Also used : MonitoringInfo(org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo) HashMap(java.util.HashMap) DataflowStepContext(org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext) Test(org.junit.Test)

Example 23 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class ExecutionTimeMonitoringInfoToCounterUpdateTransformerTest method testTransformReturnsValidCounterUpdateWhenValidMSecMonitoringInfoReceived.

@Test
public void testTransformReturnsValidCounterUpdateWhenValidMSecMonitoringInfoReceived() throws Exception {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    NameContext nc = NameContext.create("anyStageName", "anyOriginalName", "anySystemName", "anyUserName");
    DataflowStepContext dsc = mock(DataflowStepContext.class);
    when(dsc.getNameContext()).thenReturn(nc);
    stepContextMapping.put("anyValue", dsc);
    ExecutionTimeMonitoringInfoToCounterUpdateTransformer testObject = new ExecutionTimeMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    // Execute
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.PROCESS_BUNDLE_MSECS).setType(TypeUrns.SUM_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "anyValue").setPayload(encodeInt64Counter(1L)).build();
    CounterUpdate result = testObject.transform(monitoringInfo);
    // Validate
    assertNotNull(result);
    assertEqualsAsJson("{cumulative:true, integer:{highBits:0, lowBits:1}, " + "structuredNameAndMetadata:{metadata:{kind:'SUM'}, " + "name:{executionStepName:'anyStageName', name:'process-msecs', origin:'SYSTEM', " + "originalStepName:'anyOriginalName'}}}", result);
}
Also used : MonitoringInfo(org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo) HashMap(java.util.HashMap) NameContext(org.apache.beam.runners.dataflow.worker.counters.NameContext) DataflowStepContext(org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext) CounterUpdate(com.google.api.services.dataflow.model.CounterUpdate) Test(org.junit.Test)

Example 24 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class ExecutionTimeMonitoringInfoToCounterUpdateTransformerTest method testTransformThrowsIfMonitoringInfoWithUnknownUrnReceived.

@Test
public void testTransformThrowsIfMonitoringInfoWithUnknownUrnReceived() {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.USER_SUM_INT64).setType(TypeUrns.SUM_INT64_TYPE).build();
    ExecutionTimeMonitoringInfoToCounterUpdateTransformer testObject = new ExecutionTimeMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    exception.expect(RuntimeException.class);
    testObject.transform(monitoringInfo);
}
Also used : MonitoringInfo(org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo) HashMap(java.util.HashMap) DataflowStepContext(org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext) Test(org.junit.Test)

Example 25 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserParDoFnFactoryTest method testCleanupWorks.

@Test
public void testCleanupWorks() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    CounterSet counters = new CounterSet();
    DoFn<?, ?> initialFn = new TestStatefulDoFn();
    CloudObject cloudObject = getCloudObject(initialFn, WindowingStrategy.of(FixedWindows.of(Duration.millis(10))));
    StateInternals stateInternals = InMemoryStateInternals.forKey("dummy");
    // The overarching step context that only ParDoFn gets
    DataflowStepContext stepContext = mock(DataflowStepContext.class);
    // The user step context that the DoFnRunner gets a handle on
    DataflowStepContext userStepContext = mock(DataflowStepContext.class);
    when(stepContext.namespacedToUser()).thenReturn(userStepContext);
    when(stepContext.stateInternals()).thenReturn(stateInternals);
    when(userStepContext.stateInternals()).thenReturn((StateInternals) stateInternals);
    DataflowExecutionContext<DataflowStepContext> executionContext = mock(DataflowExecutionContext.class);
    TestOperationContext operationContext = TestOperationContext.create(counters);
    when(executionContext.getStepContext(operationContext)).thenReturn(stepContext);
    when(executionContext.getSideInputReader(any(), any(), any())).thenReturn(NullSideInputReader.empty());
    ParDoFn parDoFn = factory.create(options, cloudObject, Collections.emptyList(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), executionContext, operationContext);
    Receiver rcvr = new OutputReceiver();
    parDoFn.startBundle(rcvr);
    IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(9));
    IntervalWindow secondWindow = new IntervalWindow(new Instant(10), new Instant(19));
    Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();
    StateNamespace firstWindowNamespace = StateNamespaces.window(windowCoder, firstWindow);
    StateNamespace secondWindowNamespace = StateNamespaces.window(windowCoder, secondWindow);
    StateTag<ValueState<String>> tag = StateTags.tagForSpec(TestStatefulDoFn.STATE_ID, StateSpecs.value(StringUtf8Coder.of()));
    // Set up non-empty state. We don't mock + verify calls to clear() but instead
    // check that state is actually empty. We musn't care how it is accomplished.
    stateInternals.state(firstWindowNamespace, tag).write("first");
    stateInternals.state(secondWindowNamespace, tag).write("second");
    when(userStepContext.getNextFiredTimer(windowCoder)).thenReturn(null);
    when(stepContext.getNextFiredTimer(windowCoder)).thenReturn(TimerData.of(SimpleParDoFn.CLEANUP_TIMER_ID, firstWindowNamespace, firstWindow.maxTimestamp().plus(Duration.millis(1L)), firstWindow.maxTimestamp().plus(Duration.millis(1L)), TimeDomain.EVENT_TIME)).thenReturn(null);
    // This should fire the timer to clean up the first window
    parDoFn.processTimers();
    assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue());
    assertThat(stateInternals.state(secondWindowNamespace, tag).read(), equalTo("second"));
    when(stepContext.getNextFiredTimer((Coder) windowCoder)).thenReturn(TimerData.of(SimpleParDoFn.CLEANUP_TIMER_ID, secondWindowNamespace, secondWindow.maxTimestamp().plus(Duration.millis(1L)), secondWindow.maxTimestamp().plus(Duration.millis(1L)), TimeDomain.EVENT_TIME)).thenReturn(null);
    // And this should clean up the second window
    parDoFn.processTimers();
    assertThat(stateInternals.state(firstWindowNamespace, tag).read(), nullValue());
    assertThat(stateInternals.state(secondWindowNamespace, tag).read(), nullValue());
}
Also used : Coder(org.apache.beam.sdk.coders.Coder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) Instant(org.joda.time.Instant) Receiver(org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver) OutputReceiver(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver) OutputReceiver(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver) ParDoFn(org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn) DataflowStepContext(org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext) StateNamespace(org.apache.beam.runners.core.StateNamespace) ValueState(org.apache.beam.sdk.state.ValueState) CounterSet(org.apache.beam.runners.dataflow.worker.counters.CounterSet) CloudObject(org.apache.beam.runners.dataflow.util.CloudObject) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) InMemoryStateInternals(org.apache.beam.runners.core.InMemoryStateInternals) StateInternals(org.apache.beam.runners.core.StateInternals) IntervalWindow(org.apache.beam.sdk.transforms.windowing.IntervalWindow) Test(org.junit.Test)

Aggregations

DataflowStepContext (org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext)26 Test (org.junit.Test)16 HashMap (java.util.HashMap)12 MonitoringInfo (org.apache.beam.model.pipeline.v1.MetricsApi.MonitoringInfo)11 CounterUpdate (com.google.api.services.dataflow.model.CounterUpdate)7 NameContext (org.apache.beam.runners.dataflow.worker.counters.NameContext)5 CounterMetadata (com.google.api.services.dataflow.model.CounterMetadata)3 CounterStructuredName (com.google.api.services.dataflow.model.CounterStructuredName)3 CounterStructuredNameAndMetadata (com.google.api.services.dataflow.model.CounterStructuredNameAndMetadata)3 InstructionRequest (org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest)3 InstructionResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse)3 CloudObject (org.apache.beam.runners.dataflow.util.CloudObject)3 ParDoFn (org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn)3 InstructionRequestHandler (org.apache.beam.runners.fnexecution.control.InstructionRequestHandler)3 Nullable (org.checkerframework.checker.nullness.qual.Nullable)3 Instant (org.joda.time.Instant)3 CompletableFuture (java.util.concurrent.CompletableFuture)2 CountDownLatch (java.util.concurrent.CountDownLatch)2 StateKey (org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey)2 StateRequest (org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest)2