Examples with FlinkMetricContainer org.apache.beam.runners.flink.metrics.FlinkMetricContainer used on opensource projects

Example 1 with FlinkMetricContainer

use of org.apache.beam.runners.flink.metrics.FlinkMetricContainer in project beam by apache.

the class FlinkStatefulDoFnFunction method open.

@Override
public void open(Configuration parameters) {
    // Note that the SerializablePipelineOptions already initialize FileSystems in the readObject()
    // deserialization method. However, this is a hack, and we want to properly initialize the
    // options where they are needed.
    PipelineOptions options = serializedOptions.get();
    FileSystems.setDefaultPipelineOptions(options);
    metricContainer = new FlinkMetricContainer(getRuntimeContext());
    doFnInvoker = DoFnInvokers.tryInvokeSetupFor(dofn, options);
}

Example 2 with FlinkMetricContainer

use of org.apache.beam.runners.flink.metrics.FlinkMetricContainer in project beam by apache.

the class FlinkExecutableStageFunctionTest method testAccumulatorRegistrationOnOperatorClose.

@Test
public void testAccumulatorRegistrationOnOperatorClose() throws Exception {
    FlinkExecutableStageFunction<Integer> function = getFunction(Collections.emptyMap());
    function.open(new Configuration());
    String metricContainerFieldName = "metricContainer";
    FlinkMetricContainer monitoredContainer = Mockito.spy((FlinkMetricContainer) Whitebox.getInternalState(function, metricContainerFieldName));
    Whitebox.setInternalState(function, metricContainerFieldName, monitoredContainer);
    function.close();
    Mockito.verify(monitoredContainer).registerMetricsForPipelineResult();
}

Example 3 with FlinkMetricContainer

use of org.apache.beam.runners.flink.metrics.FlinkMetricContainer in project beam by apache.

the class BoundedSourceWrapper method run.

@Override
public void run(SourceContext<WindowedValue<OutputT>> ctx) throws Exception {
    // figure out which split sources we're responsible for
    int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
    int numSubtasks = getRuntimeContext().getNumberOfParallelSubtasks();
    List<BoundedSource<OutputT>> localSources = new ArrayList<>();
    for (int i = 0; i < splitSources.size(); i++) {
        if (i % numSubtasks == subtaskIndex) {
            localSources.add(splitSources.get(i));
        }
    }
    LOG.info("Bounded Flink Source {}/{} is reading from sources: {}", subtaskIndex, numSubtasks, localSources);
    FlinkMetricContainer metricContainer = new FlinkMetricContainer(getRuntimeContext());
    ReaderInvocationUtil<OutputT, BoundedSource.BoundedReader<OutputT>> readerInvoker = new ReaderInvocationUtil<>(stepName, serializedOptions.getPipelineOptions(), metricContainer);
    readers = new ArrayList<>();
    // initialize readers from scratch
    for (BoundedSource<OutputT> source : localSources) {
        readers.add(source.createReader(serializedOptions.getPipelineOptions()));
    }
    if (readers.size() == 1) {
        // the easy case, we just read from one reader
        BoundedSource.BoundedReader<OutputT> reader = readers.get(0);
        boolean dataAvailable = readerInvoker.invokeStart(reader);
        if (dataAvailable) {
            emitElement(ctx, reader);
        }
        while (isRunning) {
            dataAvailable = readerInvoker.invokeAdvance(reader);
            if (dataAvailable) {
                emitElement(ctx, reader);
            } else {
                break;
            }
        }
    } else {
        // a bit more complicated, we are responsible for several readers
        // loop through them and sleep if none of them had any data
        int currentReader = 0;
        // start each reader and emit data if immediately available
        for (BoundedSource.BoundedReader<OutputT> reader : readers) {
            boolean dataAvailable = readerInvoker.invokeStart(reader);
            if (dataAvailable) {
                emitElement(ctx, reader);
            }
        }
        // a flag telling us whether any of the readers had data
        // if no reader had data, sleep for bit
        boolean hadData = false;
        while (isRunning && !readers.isEmpty()) {
            BoundedSource.BoundedReader<OutputT> reader = readers.get(currentReader);
            boolean dataAvailable = readerInvoker.invokeAdvance(reader);
            if (dataAvailable) {
                emitElement(ctx, reader);
                hadData = true;
            } else {
                readers.remove(currentReader);
                currentReader--;
                if (readers.isEmpty()) {
                    break;
                }
            }
            currentReader = (currentReader + 1) % readers.size();
            if (currentReader == 0 && !hadData) {
                Thread.sleep(50);
            } else if (currentReader == 0) {
                hadData = false;
            }
        }
    }
    // emit final Long.MAX_VALUE watermark, just to be sure
    ctx.emitWatermark(new Watermark(Long.MAX_VALUE));
}

Example 4 with FlinkMetricContainer

use of org.apache.beam.runners.flink.metrics.FlinkMetricContainer in project beam by apache.

the class DoFnOperatorTest method testAccumulatorRegistrationOnOperatorClose.

@Test
public void testAccumulatorRegistrationOnOperatorClose() throws Exception {
    DoFnOperator<String, String> doFnOperator = getOperatorForCleanupInspection();
    OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness = new OneInputStreamOperatorTestHarness<>(doFnOperator);
    testHarness.open();
    String metricContainerFieldName = "flinkMetricContainer";
    FlinkMetricContainer monitoredContainer = Mockito.spy((FlinkMetricContainer) Whitebox.getInternalState(doFnOperator, metricContainerFieldName));
    Whitebox.setInternalState(doFnOperator, metricContainerFieldName, monitoredContainer);
    // Closes and disposes the operator
    testHarness.close();
    // Ensure that dispose has the metrics code
    doFnOperator.cleanUp();
    Mockito.verify(monitoredContainer, Mockito.times(2)).registerMetricsForPipelineResult();
}

Example 5 with FlinkMetricContainer

use of org.apache.beam.runners.flink.metrics.FlinkMetricContainer in project beam by apache.

the class SourceInputFormatTest method testAccumulatorRegistrationOnOperatorClose.

@Test
public void testAccumulatorRegistrationOnOperatorClose() throws Exception {
    SourceInputFormat<Long> sourceInputFormat = new TestSourceInputFormat<>("step", CountingSource.upTo(10), PipelineOptionsFactory.create());
    sourceInputFormat.open(sourceInputFormat.createInputSplits(1)[0]);
    String metricContainerFieldName = "metricContainer";
    FlinkMetricContainer monitoredContainer = Mockito.spy((FlinkMetricContainer) Whitebox.getInternalState(sourceInputFormat, metricContainerFieldName));
    Whitebox.setInternalState(sourceInputFormat, metricContainerFieldName, monitoredContainer);
    sourceInputFormat.close();
    Mockito.verify(monitoredContainer).registerMetricsForPipelineResult();
}