Examples with WatermarkMessage org.apache.samza.system.WatermarkMessage used on opensource projects

Example 6 with WatermarkMessage

use of org.apache.samza.system.WatermarkMessage in project samza by apache.

the class TestControlMessageSender method testSend.

@Test
public void testSend() {
    SystemStreamMetadata metadata = mock(SystemStreamMetadata.class);
    Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
    partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata);
    StreamMetadataCache metadataCache = mock(StreamMetadataCache.class);
    when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata);
    SystemStream systemStream = new SystemStream("test-system", "test-stream");
    Set<Integer> partitions = new HashSet<>();
    MessageCollector collector = mock(MessageCollector.class);
    doAnswer(invocation -> {
        OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0];
        partitions.add((Integer) envelope.getPartitionKey());
        assertEquals(envelope.getSystemStream(), systemStream);
        return null;
    }).when(collector).send(any());
    ControlMessageSender sender = new ControlMessageSender(metadataCache);
    WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0");
    sender.send(watermark, systemStream, collector);
    assertEquals(partitions.size(), 1);
}

Example 7 with WatermarkMessage

use of org.apache.samza.system.WatermarkMessage in project samza by apache.

the class TestControlMessageSender method testBroadcast.

@Test
public void testBroadcast() {
    SystemStreamMetadata metadata = mock(SystemStreamMetadata.class);
    Map<Partition, SystemStreamMetadata.SystemStreamPartitionMetadata> partitionMetadata = new HashMap<>();
    partitionMetadata.put(new Partition(0), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    partitionMetadata.put(new Partition(1), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    partitionMetadata.put(new Partition(2), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    partitionMetadata.put(new Partition(3), mock(SystemStreamMetadata.SystemStreamPartitionMetadata.class));
    when(metadata.getSystemStreamPartitionMetadata()).thenReturn(partitionMetadata);
    StreamMetadataCache metadataCache = mock(StreamMetadataCache.class);
    when(metadataCache.getSystemStreamMetadata(anyObject(), anyBoolean())).thenReturn(metadata);
    SystemStream systemStream = new SystemStream("test-system", "test-stream");
    Set<Integer> partitions = new HashSet<>();
    MessageCollector collector = mock(MessageCollector.class);
    doAnswer(invocation -> {
        OutgoingMessageEnvelope envelope = (OutgoingMessageEnvelope) invocation.getArguments()[0];
        partitions.add((Integer) envelope.getPartitionKey());
        assertEquals(envelope.getSystemStream(), systemStream);
        return null;
    }).when(collector).send(any());
    ControlMessageSender sender = new ControlMessageSender(metadataCache);
    WatermarkMessage watermark = new WatermarkMessage(System.currentTimeMillis(), "task 0");
    SystemStreamPartition ssp = new SystemStreamPartition(systemStream, new Partition(0));
    sender.broadcastToOtherPartitions(watermark, ssp, collector);
    assertEquals(partitions.size(), 3);
}

Example 8 with WatermarkMessage

use of org.apache.samza.system.WatermarkMessage in project samza by apache.

the class OperatorImpl method aggregateWatermark.

/**
 * Aggregate the {@link WatermarkMessage} from each ssp into a watermark. Then call onWatermark() if
 * a new watermark exits.
 * @param watermarkMessage a {@link WatermarkMessage} object
 * @param ssp {@link SystemStreamPartition} that the message is coming from.
 * @param collector message collector
 * @param coordinator task coordinator
 */
public final CompletionStage<Void> aggregateWatermark(WatermarkMessage watermarkMessage, SystemStreamPartition ssp, MessageCollector collector, TaskCoordinator coordinator) {
    LOG.debug("Received watermark {} from {}", watermarkMessage.getTimestamp(), ssp);
    watermarkStates.update(watermarkMessage, ssp);
    long watermark = watermarkStates.getWatermark(ssp.getSystemStream());
    CompletionStage<Void> watermarkFuture = CompletableFuture.completedFuture(null);
    if (currentWatermark < watermark) {
        LOG.debug("Got watermark {} from stream {}", watermark, ssp.getSystemStream());
        if (watermarkMessage.getTaskName() != null && shouldTaskBroadcastToOtherPartitions(ssp)) {
            // This is the aggregation task, which already received all the watermark messages from upstream
            // broadcast the watermark to all the peer partitions
            // additionally if elasiticty is enabled
            // then only one of the elastic tasks of the ssp will broadcast
            controlMessageSender.broadcastToOtherPartitions(new WatermarkMessage(watermark), ssp, collector);
        }
        // populate the watermark through the dag
        watermarkFuture = onWatermark(watermark, collector, coordinator).thenAccept(ignored -> watermarkStates.updateAggregateMetric(ssp, watermark));
    }
    return watermarkFuture;
}

Example 9 with WatermarkMessage

use of org.apache.samza.system.WatermarkMessage in project beam by apache.

the class SamzaImpulseSystemTest method testSamzaImpulseSystemConsumer.

@Test
public void testSamzaImpulseSystemConsumer() throws Exception {
    SystemConsumer consumer = new SamzaImpulseSystemFactory().getConsumer("default-system", new MapConfig(), null);
    Map<SystemStreamPartition, List<IncomingMessageEnvelope>> result = consumer.poll(Collections.singleton(sspForPartition(0)), 100);
    Assert.assertEquals(1, result.size());
    Assert.assertTrue(result.containsKey(sspForPartition(0)));
    List<IncomingMessageEnvelope> messageEnvelopes = result.get(sspForPartition(0));
    Assert.assertEquals(3, messageEnvelopes.size());
    Assert.assertTrue(messageEnvelopes.get(0).getMessage() instanceof OpMessage);
    OpMessage impulseEvent = (OpMessage) messageEnvelopes.get(0).getMessage();
    Assert.assertEquals(OpMessage.Type.ELEMENT, impulseEvent.getType());
    Assert.assertTrue(messageEnvelopes.get(1).getMessage() instanceof WatermarkMessage);
    Assert.assertTrue(messageEnvelopes.get(2).isEndOfStream());
}