Examples with StreamingModeExecutionState org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionState used on opensource projects

Example 1 with StreamingModeExecutionState

use of org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionState in project beam by apache.

the class StreamingModeExecutionContextTest method stateSamplingInStreaming.

@Test(timeout = 2000)
public void stateSamplingInStreaming() {
    // Test that when writing on one thread and reading from another, updates always eventually
    // reach the reading thread.
    StreamingModeExecutionState state = new StreamingModeExecutionState(NameContextsForTests.nameContextForTest(), "testState", null, NoopProfileScope.NOOP, null);
    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
        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(0L));
        }
    } finally {
        sampler.stop();
    }
}

Example 2 with StreamingModeExecutionState

use of org.apache.beam.runners.dataflow.worker.StreamingModeExecutionContext.StreamingModeExecutionState in project beam by apache.

the class StreamingModeExecutionContextTest method testAtomicExtractUpdate.

/**
 * Ensure that incrementing and extracting counter updates are correct under concurrent reader and
 * writer threads.
 */
@Test
public void testAtomicExtractUpdate() throws InterruptedException, ExecutionException {
    long numUpdates = 1_000_000;
    StreamingModeExecutionState state = new StreamingModeExecutionState(NameContextsForTests.nameContextForTest(), "testState", null, NoopProfileScope.NOOP, null);
    ExecutorService executor = Executors.newFixedThreadPool(2);
    AtomicBoolean doneWriting = new AtomicBoolean(false);
    Callable<Long> reader = () -> {
        long count = 0;
        boolean isLastRead;
        do {
            isLastRead = doneWriting.get();
            CounterUpdate update = state.extractUpdate(false);
            if (update != null) {
                count += splitIntToLong(update.getInteger());
            }
        } while (!isLastRead);
        return count;
    };
    Runnable writer = () -> {
        for (int i = 0; i < numUpdates; i++) {
            state.takeSample(1L);
        }
        doneWriting.set(true);
    };
    // NB: Reader is invoked before writer to ensure they execute concurrently.
    List<Future<Long>> results = executor.invokeAll(Lists.newArrayList(reader, Executors.callable(writer, 0L)), 2, TimeUnit.SECONDS);
    long count = results.get(0).get();
    assertThat(count, equalTo(numUpdates));
}