Examples with BoundedQueueExecutor org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor used on opensource projects

Example 1 with BoundedQueueExecutor

use of org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor in project beam by apache.

the class StreamingDataflowWorkerTest method testActiveWorkForShardedKeys.

@Test
public void testActiveWorkForShardedKeys() throws Exception {
    BoundedQueueExecutor mockExecutor = Mockito.mock(BoundedQueueExecutor.class);
    StreamingDataflowWorker.ComputationState computationState = new StreamingDataflowWorker.ComputationState("computation", defaultMapTask(Arrays.asList(makeSourceInstruction(StringUtf8Coder.of()))), mockExecutor, ImmutableMap.of(), null);
    ShardedKey key1Shard1 = ShardedKey.create(ByteString.copyFromUtf8("key1"), 1);
    ShardedKey key1Shard2 = ShardedKey.create(ByteString.copyFromUtf8("key1"), 2);
    MockWork m1 = new MockWork(1);
    assertTrue(computationState.activateWork(key1Shard1, m1));
    Mockito.verify(mockExecutor).execute(m1, m1.getWorkItem().getSerializedSize());
    computationState.completeWork(key1Shard1, 1);
    Mockito.verifyNoMoreInteractions(mockExecutor);
    // Verify work queues.
    MockWork m2 = new MockWork(2);
    assertTrue(computationState.activateWork(key1Shard1, m2));
    Mockito.verify(mockExecutor).execute(m2, m2.getWorkItem().getSerializedSize());
    MockWork m3 = new MockWork(3);
    assertTrue(computationState.activateWork(key1Shard1, m3));
    Mockito.verifyNoMoreInteractions(mockExecutor);
    // Verify a different shard of key is a separate queue.
    MockWork m4 = new MockWork(3);
    assertFalse(computationState.activateWork(key1Shard1, m4));
    Mockito.verifyNoMoreInteractions(mockExecutor);
    assertTrue(computationState.activateWork(key1Shard2, m4));
    Mockito.verify(mockExecutor).execute(m4, m4.getWorkItem().getSerializedSize());
    // Verify duplicate work dropped
    assertFalse(computationState.activateWork(key1Shard2, m4));
    computationState.completeWork(key1Shard2, 3);
    Mockito.verifyNoMoreInteractions(mockExecutor);
}

Example 2 with BoundedQueueExecutor

use of org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor in project beam by apache.

the class StreamingDataflowWorkerTest method testActiveWork.

@Test
public void testActiveWork() throws Exception {
    BoundedQueueExecutor mockExecutor = Mockito.mock(BoundedQueueExecutor.class);
    StreamingDataflowWorker.ComputationState computationState = new StreamingDataflowWorker.ComputationState("computation", defaultMapTask(Arrays.asList(makeSourceInstruction(StringUtf8Coder.of()))), mockExecutor, ImmutableMap.of(), null);
    ShardedKey key1 = ShardedKey.create(ByteString.copyFromUtf8("key1"), 1);
    ShardedKey key2 = ShardedKey.create(ByteString.copyFromUtf8("key2"), 2);
    MockWork m1 = new MockWork(1);
    assertTrue(computationState.activateWork(key1, m1));
    Mockito.verify(mockExecutor).execute(m1, m1.getWorkItem().getSerializedSize());
    computationState.completeWork(key1, 1);
    Mockito.verifyNoMoreInteractions(mockExecutor);
    // Verify work queues.
    MockWork m2 = new MockWork(2);
    assertTrue(computationState.activateWork(key1, m2));
    Mockito.verify(mockExecutor).execute(m2, m2.getWorkItem().getSerializedSize());
    MockWork m3 = new MockWork(3);
    assertTrue(computationState.activateWork(key1, m3));
    Mockito.verifyNoMoreInteractions(mockExecutor);
    // Verify another key is a separate queue.
    MockWork m4 = new MockWork(4);
    assertTrue(computationState.activateWork(key2, m4));
    Mockito.verify(mockExecutor).execute(m4, m4.getWorkItem().getSerializedSize());
    computationState.completeWork(key2, 4);
    Mockito.verifyNoMoreInteractions(mockExecutor);
    computationState.completeWork(key1, 2);
    Mockito.verify(mockExecutor).forceExecute(m3, m3.getWorkItem().getSerializedSize());
    computationState.completeWork(key1, 3);
    Mockito.verifyNoMoreInteractions(mockExecutor);
    // Verify duplicate work dropped.
    MockWork m5 = new MockWork(5);
    computationState.activateWork(key1, m5);
    Mockito.verify(mockExecutor).execute(m5, m5.getWorkItem().getSerializedSize());
    assertFalse(computationState.activateWork(key1, m5));
    Mockito.verifyNoMoreInteractions(mockExecutor);
    computationState.completeWork(key1, 5);
    Mockito.verifyNoMoreInteractions(mockExecutor);
}