Examples with NameContext org.apache.beam.runners.dataflow.worker.counters.NameContext used on opensource projects

Example 6 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class MapTaskExecutorTest method testPerElementProcessingTimeCounters.

/**
 * Verify counts for the per-element-output-time counter are correct.
 */
@Test
public void testPerElementProcessingTimeCounters() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(DataflowPipelineDebugOptions.class).setExperiments(Lists.newArrayList(DataflowElementExecutionTracker.TIME_PER_ELEMENT_EXPERIMENT));
    ExecutionStateSampler stateSampler = ExecutionStateSampler.newForTest();
    DataflowExecutionStateTracker stateTracker = new DataflowExecutionStateTracker(stateSampler, new TestDataflowExecutionState(NameContext.forStage("test-stage"), "other", null, /* requestingStepName */
    null, /* sideInputIndex */
    null, /* metricsContainer */
    NoopProfileScope.NOOP), counterSet, options, "test-work-item-id");
    NameContext parDoName = nameForStep("s1");
    // Wire a read operation with 3 elements to a ParDoOperation and assert that we count
    // the correct number of elements.
    ReadOperation read = ReadOperation.forTest(new TestReader("a", "b", "c"), new OutputReceiver(), TestOperationContext.create(counterSet, nameForStep("s0"), null, stateTracker));
    ParDoOperation parDo = new ParDoOperation(new NoopParDoFn(), new OutputReceiver[0], TestOperationContext.create(counterSet, parDoName, null, stateTracker));
    parDo.attachInput(read, 0);
    List<Operation> operations = Lists.newArrayList(read, parDo);
    try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) {
        executor.execute();
    }
    stateSampler.doSampling(100L);
    CounterName counterName = CounterName.named("per-element-processing-time").withOriginalName(parDoName);
    Counter<Long, CounterDistribution> counter = (Counter<Long, CounterDistribution>) counterSet.getExistingCounter(counterName);
    assertThat(counter.getAggregate().getCount(), equalTo(3L));
}

Example 7 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class StreamingModeExecutionContextTest method testTimerInternalsProcessingTimeSkew.

@Test
public void testTimerInternalsProcessingTimeSkew() {
    Windmill.WorkItemCommitRequest.Builder outputBuilder = Windmill.WorkItemCommitRequest.newBuilder();
    NameContext nameContext = NameContextsForTests.nameContextForTest();
    DataflowOperationContext operationContext = executionContext.createOperationContext(nameContext);
    StreamingModeExecutionContext.StepContext stepContext = executionContext.getStepContext(operationContext);
    Windmill.WorkItem.Builder workItemBuilder = Windmill.WorkItem.newBuilder().setKey(ByteString.EMPTY).setWorkToken(17L);
    Windmill.Timer.Builder timerBuilder = workItemBuilder.getTimersBuilder().addTimersBuilder();
    // Trigger a realtime timer that with clock skew but ensure that it would
    // still fire.
    Instant now = Instant.now();
    long offsetMillis = 60 * 1000;
    Instant timerTimestamp = now.plus(Duration.millis(offsetMillis));
    timerBuilder.setTag(ByteString.copyFromUtf8("a")).setTimestamp(timerTimestamp.getMillis() * 1000).setType(Windmill.Timer.Type.REALTIME);
    executionContext.start("key", workItemBuilder.build(), // input watermark
    new Instant(1000), // output watermark
    null, // synchronized processing time
    null, stateReader, stateFetcher, outputBuilder);
    TimerInternals timerInternals = stepContext.timerInternals();
    assertTrue(timerTimestamp.isBefore(timerInternals.currentProcessingTime()));
}

Example 8 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class TestOperationContext method create.

public static TestOperationContext create() {
    CounterSet counterSet = new CounterSet();
    NameContext nameContext = NameContextsForTests.nameContextForTest();
    return create(counterSet, nameContext);
}

Example 9 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class DataflowSideInputReadCounterTest method testEnterEntersStateIfCalledFromTrackedThread.

@Test
public void testEnterEntersStateIfCalledFromTrackedThread() {
    DataflowExecutionContext mockedExecutionContext = mock(DataflowExecutionContext.class);
    DataflowOperationContext mockedOperationContext = mock(DataflowOperationContext.class);
    final int siIndexId = 3;
    ExecutionStateTracker mockedExecutionStateTracker = mock(ExecutionStateTracker.class);
    when(mockedExecutionContext.getExecutionStateTracker()).thenReturn(mockedExecutionStateTracker);
    Thread mockedThreadObject = Thread.currentThread();
    when(mockedExecutionStateTracker.getTrackedThread()).thenReturn(mockedThreadObject);
    DataflowExecutionState mockedExecutionState = mock(DataflowExecutionState.class);
    when(mockedExecutionStateTracker.getCurrentState()).thenReturn(mockedExecutionState);
    NameContext mockedNameContext = mock(NameContext.class);
    when(mockedExecutionState.getStepName()).thenReturn(mockedNameContext);
    when(mockedNameContext.originalName()).thenReturn("DummyName");
    NameContext mockedDeclaringNameContext = mock(NameContext.class);
    when(mockedOperationContext.nameContext()).thenReturn(mockedDeclaringNameContext);
    when(mockedDeclaringNameContext.originalName()).thenReturn("DummyDeclaringName");
    CounterFactory mockedCounterFactory = mock(CounterFactory.class);
    when(mockedExecutionContext.getCounterFactory()).thenReturn(mockedCounterFactory);
    Counter<Long, Long> mockedCounter = mock(Counter.class);
    when(mockedCounterFactory.longSum(any())).thenReturn(mockedCounter);
    DataflowExecutionStateRegistry mockedExecutionStateRegistry = mock(DataflowExecutionStateRegistry.class);
    when(mockedExecutionContext.getExecutionStateRegistry()).thenReturn(mockedExecutionStateRegistry);
    DataflowExecutionState mockedCounterExecutionState = mock(DataflowExecutionState.class);
    when(mockedExecutionStateRegistry.getIOState(any(), any(), any(), any(), any(), any())).thenReturn(mockedCounterExecutionState);
    DataflowSideInputReadCounter testObject = new DataflowSideInputReadCounter(mockedExecutionContext, mockedOperationContext, siIndexId);
    testObject.enter();
    verify(mockedExecutionStateTracker).enterState(mockedCounterExecutionState);
}

Example 10 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class DataflowSideInputReadCounterTest method testAddBytesReadUpdatesCounter.

@Test
public void testAddBytesReadUpdatesCounter() {
    DataflowExecutionContext mockedExecutionContext = mock(DataflowExecutionContext.class);
    DataflowOperationContext mockedOperationContext = mock(DataflowOperationContext.class);
    final int siIndexId = 3;
    ExecutionStateTracker mockedExecutionStateTracker = mock(ExecutionStateTracker.class);
    when(mockedExecutionContext.getExecutionStateTracker()).thenReturn(mockedExecutionStateTracker);
    Thread mockedThreadObject = mock(Thread.class);
    when(mockedExecutionStateTracker.getTrackedThread()).thenReturn(mockedThreadObject);
    DataflowExecutionState mockedExecutionState = mock(DataflowExecutionState.class);
    when(mockedExecutionStateTracker.getCurrentState()).thenReturn(mockedExecutionState);
    NameContext mockedNameContext = mock(NameContext.class);
    when(mockedExecutionState.getStepName()).thenReturn(mockedNameContext);
    when(mockedNameContext.originalName()).thenReturn("DummyName");
    NameContext mockedDeclaringNameContext = mock(NameContext.class);
    when(mockedOperationContext.nameContext()).thenReturn(mockedDeclaringNameContext);
    when(mockedDeclaringNameContext.originalName()).thenReturn("DummyDeclaringName");
    CounterFactory mockedCounterFactory = mock(CounterFactory.class);
    when(mockedExecutionContext.getCounterFactory()).thenReturn(mockedCounterFactory);
    Counter<Long, Long> mockedCounter = mock(Counter.class);
    when(mockedCounterFactory.longSum(any())).thenReturn(mockedCounter);
    when(mockedExecutionContext.getExecutionStateRegistry()).thenReturn(mock(DataflowExecutionStateRegistry.class));
    DataflowSideInputReadCounter testObject = new DataflowSideInputReadCounter(mockedExecutionContext, mockedOperationContext, siIndexId);
    testObject.addBytesRead(10l);
    verify(mockedCounter).addValue(10l);
}