Search in sources :

Example 16 with OutgoingMessageEnvelope

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

the class TestMessageStreamImpl method testSink.

@Test
public void testSink() {
    MessageStreamImpl<TestMessageEnvelope> inputStream = new MessageStreamImpl<>(mockGraph);
    SystemStream testStream = new SystemStream("test-sys", "test-stream");
    SinkFunction<TestMessageEnvelope> xSink = (TestMessageEnvelope m, MessageCollector mc, TaskCoordinator tc) -> {
        mc.send(new OutgoingMessageEnvelope(testStream, m.getMessage()));
        tc.commit(TaskCoordinator.RequestScope.CURRENT_TASK);
    };
    inputStream.sink(xSink);
    Collection<OperatorSpec> subs = inputStream.getRegisteredOperatorSpecs();
    assertEquals(subs.size(), 1);
    OperatorSpec<TestMessageEnvelope> sinkOp = subs.iterator().next();
    assertTrue(sinkOp instanceof SinkOperatorSpec);
    assertEquals(((SinkOperatorSpec) sinkOp).getSinkFn(), xSink);
    TestMessageEnvelope mockTest1 = mock(TestMessageEnvelope.class);
    MessageType mockMsgBody = mock(MessageType.class);
    when(mockTest1.getMessage()).thenReturn(mockMsgBody);
    final List<OutgoingMessageEnvelope> outMsgs = new ArrayList<>();
    MessageCollector mockCollector = mock(MessageCollector.class);
    doAnswer(invocation -> {
        outMsgs.add((OutgoingMessageEnvelope) invocation.getArguments()[0]);
        return null;
    }).when(mockCollector).send(any());
    TaskCoordinator mockCoordinator = mock(TaskCoordinator.class);
    ((SinkOperatorSpec) sinkOp).getSinkFn().apply(mockTest1, mockCollector, mockCoordinator);
    assertEquals(1, outMsgs.size());
    assertEquals(testStream, outMsgs.get(0).getSystemStream());
    assertEquals(mockMsgBody, outMsgs.get(0).getMessage());
}
Also used : SystemStream(org.apache.samza.system.SystemStream) ArrayList(java.util.ArrayList) TaskCoordinator(org.apache.samza.task.TaskCoordinator) StreamOperatorSpec(org.apache.samza.operators.spec.StreamOperatorSpec) OperatorSpec(org.apache.samza.operators.spec.OperatorSpec) PartialJoinOperatorSpec(org.apache.samza.operators.spec.PartialJoinOperatorSpec) SinkOperatorSpec(org.apache.samza.operators.spec.SinkOperatorSpec) TestMessageEnvelope(org.apache.samza.operators.data.TestMessageEnvelope) MessageCollector(org.apache.samza.task.MessageCollector) SinkOperatorSpec(org.apache.samza.operators.spec.SinkOperatorSpec) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) MessageType(org.apache.samza.operators.data.MessageType) Test(org.junit.Test)

Example 17 with OutgoingMessageEnvelope

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

the class CoordinatorStreamSystemProducer method send.

/**
   * Serialize and send a coordinator stream message.
   * 
   * @param message
   *          The message to send.
   */
public void send(CoordinatorStreamMessage message) {
    log.debug("Sending {}", message);
    try {
        String source = message.getSource();
        byte[] key = keySerde.toBytes(Arrays.asList(message.getKeyArray()));
        byte[] value = null;
        if (!message.isDelete()) {
            value = messageSerde.toBytes(message.getMessageMap());
        }
        OutgoingMessageEnvelope envelope = new OutgoingMessageEnvelope(systemStream, Integer.valueOf(0), key, value);
        systemProducer.send(source, envelope);
    } catch (Exception e) {
        throw new SamzaException(e);
    }
}
Also used : OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) SamzaException(org.apache.samza.SamzaException) SamzaException(org.apache.samza.SamzaException)

Example 18 with OutgoingMessageEnvelope

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

the class IdentityStreamTask method process.

@Override
public void process(IncomingMessageEnvelope incomingMessageEnvelope, MessageCollector messageCollector, TaskCoordinator taskCoordinator) throws Exception {
    messageCollector.send(new OutgoingMessageEnvelope(new SystemStream(outputSystem, outputTopic), incomingMessageEnvelope.getMessage()));
    processedMessageCount++;
    if (processedMessageCount == expectedMessageCount) {
        endLatch.countDown();
    }
}
Also used : SystemStream(org.apache.samza.system.SystemStream) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope)

Example 19 with OutgoingMessageEnvelope

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

the class Checker method window.

@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
    String currentEpoch = this.store.get(CURRENT_EPOCH);
    logger.info("Checking if epoch " + currentEpoch + " is complete.");
    int count = 0;
    KeyValueIterator<String, String> iter = this.store.all();
    while (iter.hasNext()) {
        Entry<String, String> entry = iter.next();
        String foundEpoch = entry.getValue();
        if (foundEpoch.equals(currentEpoch)) {
            count += 1;
        } else {
            logger.info("####### Found a different epoch! - " + foundEpoch + " Current epoch is " + currentEpoch);
        }
    }
    iter.close();
    if (count == expectedKeys + 1) {
        logger.info("Epoch " + currentEpoch + " is complete.");
        int nextEpoch = Integer.parseInt(currentEpoch) + 1;
        for (int i = 0; i < numPartitions; i++) {
            logger.info("Emitting next epoch - " + Integer.toString(i) + " -> " + Integer.toString(nextEpoch));
            collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "epoch"), Integer.toString(i), Integer.toString(nextEpoch)));
        }
        this.store.put(CURRENT_EPOCH, Integer.toString(nextEpoch));
    } else if (count > expectedKeys + 1) {
        throw new IllegalStateException("Got " + count + " keys, which is more than the expected " + (expectedKeys + 1));
    } else {
        logger.info("Only found " + count + " valid keys, try again later.");
    }
}
Also used : SystemStream(org.apache.samza.system.SystemStream) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope)

Example 20 with OutgoingMessageEnvelope

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

the class Joiner method process.

@Override
public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) {
    String key = (String) envelope.getKey();
    String value = (String) envelope.getMessage();
    String[] pieces = value.split("-");
    int epoch = Integer.parseInt(pieces[0]);
    int partition = Integer.parseInt(pieces[1].split(" ")[1]);
    Partitions partitions = loadPartitions(epoch, key);
    logger.info("Joiner got epoch = " + epoch + ", partition = " + partition + ", parts = " + partitions);
    if (partitions.epoch < epoch) {
        // we are in a new era
        if (partitions.partitions.size() != expected)
            throw new IllegalArgumentException("Should have " + expected + " partitions when new epoch starts.");
        logger.info("Reseting epoch to " + epoch);
        this.store.delete(key);
        partitions.epoch = epoch;
        partitions.partitions.clear();
        partitions.partitions.add(partition);
    } else if (partitions.epoch > epoch) {
        logger.info("Ignoring message for epoch " + epoch);
    } else {
        partitions.partitions.add(partition);
        if (partitions.partitions.size() == expected) {
            logger.info("Completed: " + key + " -> " + Integer.toString(epoch));
            collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "completed-keys"), key, Integer.toString(epoch)));
        }
    }
    this.store.put(key, partitions.toString());
    logger.info("Join store in Task " + this.taskName + " " + key + " -> " + partitions.toString());
}
Also used : SystemStream(org.apache.samza.system.SystemStream) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope)

Aggregations

OutgoingMessageEnvelope (org.apache.samza.system.OutgoingMessageEnvelope)20 Test (org.junit.Test)12 IndexRequest (org.elasticsearch.action.index.IndexRequest)8 SystemStream (org.apache.samza.system.SystemStream)7 ArrayList (java.util.ArrayList)2 Map (java.util.Map)2 SamzaException (org.apache.samza.SamzaException)2 MessageType (org.apache.samza.operators.data.MessageType)2 TestMessageEnvelope (org.apache.samza.operators.data.TestMessageEnvelope)2 MessageCollector (org.apache.samza.task.MessageCollector)2 TaskCoordinator (org.apache.samza.task.TaskCoordinator)2 TypeReference (org.codehaus.jackson.type.TypeReference)2 IOException (java.io.IOException)1 Method (java.lang.reflect.Method)1 HashMap (java.util.HashMap)1 CoordinatorStreamMessage (org.apache.samza.coordinator.stream.messages.CoordinatorStreamMessage)1 Delete (org.apache.samza.coordinator.stream.messages.Delete)1 SetConfig (org.apache.samza.coordinator.stream.messages.SetConfig)1 OperatorSpec (org.apache.samza.operators.spec.OperatorSpec)1 PartialJoinOperatorSpec (org.apache.samza.operators.spec.PartialJoinOperatorSpec)1