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

Example 6 with Counter

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

the class DataflowExecutionStateTrackerTest method assertElementProcessingTimeCounter.

private void assertElementProcessingTimeCounter(NameContext step, int millis, int bucketOffset) {
    CounterName counterName = ElementExecutionTracker.COUNTER_NAME.withOriginalName(step);
    Counter<?, CounterDistribution> counter = (Counter<?, CounterFactory.CounterDistribution>) counterSet.getExistingCounter(counterName);
    assertNotNull(counter);
    CounterFactory.CounterDistribution distribution = counter.getAggregate();
    assertThat(distribution, equalTo(CounterFactory.CounterDistribution.builder().minMax(millis, millis).count(1).sum(millis).sumOfSquares(millis * millis).buckets(bucketOffset, Lists.newArrayList(1L)).build()));
}

Example 7 with Counter

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

the class GroupingShuffleEntryIteratorTest method testCopyValuesIterator.

@Test
public void testCopyValuesIterator() throws Exception {
    setCurrentExecutionState(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1);
    MockitoAnnotations.initMocks(this);
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(DataflowPipelineDebugOptions.class).setExperiments(Lists.newArrayList(Experiment.IntertransformIO.getName()));
    ArrayList<ShuffleEntry> entries = new ArrayList<>();
    entries.add(shuffleEntry("k1", "v11"));
    entries.add(shuffleEntry("k1", "v12"));
    entries.add(shuffleEntry("k1", "v13"));
    when(reader.read(START_POSITION, END_POSITION)).thenReturn(new ListReiterator<>(entries, 0));
    final ShuffleReadCounter shuffleReadCounter = new ShuffleReadCounter(ORIGINAL_SHUFFLE_STEP_NAME, true, null);
    iterator = new GroupingShuffleEntryIterator(reader, START_POSITION, END_POSITION) {

        @Override
        protected void notifyElementRead(long byteSize) {
        // nothing
        }

        @Override
        protected void commitBytesRead(long bytes) {
            shuffleReadCounter.addBytesRead(bytes);
        }
    };
    assertTrue(iterator.advance());
    KeyGroupedShuffleEntries k1Entries = iterator.getCurrent();
    Reiterator<ShuffleEntry> values1 = k1Entries.values.iterator();
    assertTrue(values1.hasNext());
    Reiterator<ShuffleEntry> values1Copy1 = values1.copy();
    Reiterator<ShuffleEntry> values1Copy2 = values1.copy();
    ShuffleEntry expectedEntry = values1.next();
    // Advance the iterator again to record bytes read.
    assertFalse(iterator.advance());
    // Test that the copy works two ways: 1) if we call hasNext, and 2) if we don't.
    assertTrue(values1Copy1.hasNext());
    assertEquals(expectedEntry, values1Copy1.next());
    assertEquals(expectedEntry, values1Copy2.next());
    Map<String, Long> expectedReadBytesMap = new HashMap<>();
    expectedReadBytesMap.put(MOCK_ORIGINAL_NAME_FOR_EXECUTING_STEP1, 15L);
    // Verify that each executing step used when reading from the GroupingShuffleReader
    // has a counter with a bytes read value.
    assertEquals(expectedReadBytesMap.size(), (long) shuffleReadCounter.counterSet.size());
    Iterator it = expectedReadBytesMap.entrySet().iterator();
    while (it.hasNext()) {
        Map.Entry<String, Long> pair = (Map.Entry) it.next();
        Counter counter = shuffleReadCounter.counterSet.getExistingCounter(ShuffleReadCounter.generateCounterName(ORIGINAL_SHUFFLE_STEP_NAME, pair.getKey()));
        assertEquals(pair.getValue(), counter.getAggregate());
    }
}