Search in sources :

Example 1 with BatchModeExecutionState

use of org.apache.beam.runners.dataflow.worker.BatchModeExecutionContext.BatchModeExecutionState in project beam by apache.

the class BatchModeExecutionContextTest method stateSamplingInBatch.

@Test(timeout = 2000)
public void stateSamplingInBatch() {
    // Test that when writing on one thread and reading from another, updates always eventually
    // reach the reading thread.
    BatchModeExecutionState state = new BatchModeExecutionState(NameContextsForTests.nameContextForTest(), "testState", null, /* requestingStepName */
    null, /* inputIndex */
    null, /* metricsContainer */
    NoopProfileScope.NOOP);
    ExecutionStateSampler sampler = ExecutionStateSampler.newForTest();
    try {
        sampler.start();
        ExecutionStateTracker tracker = new ExecutionStateTracker(sampler);
        Thread executionThread = new Thread();
        executionThread.setName("looping-thread-for-test");
        tracker.activate(executionThread);
        tracker.enterState(state);
        // Wait for the state to be incremented 3 times
        long value = 0;
        for (int i = 0; i < 3; i++) {
            CounterUpdate update = null;
            while (update == null) {
                update = state.extractUpdate(false);
            }
            long newValue = splitIntToLong(update.getInteger());
            assertThat(newValue, Matchers.greaterThan(value));
            value = newValue;
        }
    } finally {
        sampler.stop();
    }
}
Also used : ExecutionStateTracker(org.apache.beam.runners.core.metrics.ExecutionStateTracker) BatchModeExecutionState(org.apache.beam.runners.dataflow.worker.BatchModeExecutionContext.BatchModeExecutionState) ExecutionStateSampler(org.apache.beam.runners.core.metrics.ExecutionStateSampler) CounterUpdate(com.google.api.services.dataflow.model.CounterUpdate) Test(org.junit.Test)

Aggregations

CounterUpdate (com.google.api.services.dataflow.model.CounterUpdate)1 ExecutionStateSampler (org.apache.beam.runners.core.metrics.ExecutionStateSampler)1 ExecutionStateTracker (org.apache.beam.runners.core.metrics.ExecutionStateTracker)1 BatchModeExecutionState (org.apache.beam.runners.dataflow.worker.BatchModeExecutionContext.BatchModeExecutionState)1 Test (org.junit.Test)1