Examples with MessageId org.apache.metron.common.writer.MessageId used on opensource projects

Example 6 with MessageId

use of org.apache.metron.common.writer.MessageId in project metron by apache.

the class KafkaWriterTest method testWriteShouldReturnErrorsOnFailedFlush.

@Test
public void testWriteShouldReturnErrorsOnFailedFlush() throws Exception {
    KafkaWriter writer = spy(new KafkaWriter());
    writer.setKafkaProducer(kafkaProducer);
    List<BulkMessage<JSONObject>> messages = new ArrayList<>();
    JSONObject message1 = new JSONObject();
    message1.put("value", "message1");
    JSONObject message2 = new JSONObject();
    message2.put("value", "message2");
    messages.add(new BulkMessage<>("messageId1", message1));
    messages.add(new BulkMessage<>("messageId2", message2));
    doReturn(Optional.of("topic1")).when(writer).getKafkaTopic(message1);
    doReturn(Optional.of("topic2")).when(writer).getKafkaTopic(message2);
    Future future1 = mock(Future.class);
    Future future2 = mock(Future.class);
    when(kafkaProducer.send(new ProducerRecord<String, String>("topic1", "{\"value\":\"message1\"}"))).thenReturn(future1);
    when(kafkaProducer.send(new ProducerRecord<String, String>("topic2", "{\"value\":\"message2\"}"))).thenReturn(future2);
    InterruptException throwable = new InterruptException("kafka flush exception");
    doThrow(throwable).when(kafkaProducer).flush();
    BulkWriterResponse response = new BulkWriterResponse();
    response.addAllErrors(throwable, Arrays.asList(new MessageId("messageId1"), new MessageId("messageId2")));
    assertEquals(response, writer.write(SENSOR_TYPE, createConfiguration(new HashMap<>()), messages));
    verify(kafkaProducer, times(1)).flush();
    verify(kafkaProducer, times(1)).send(new ProducerRecord<String, String>("topic1", "{\"value\":\"message1\"}"));
    verify(kafkaProducer, times(1)).send(new ProducerRecord<String, String>("topic2", "{\"value\":\"message2\"}"));
    verifyNoMoreInteractions(kafkaProducer);
}

Example 7 with MessageId

use of org.apache.metron.common.writer.MessageId in project metron by apache.

the class KafkaWriter method write.

@Override
public BulkWriterResponse write(String sensorType, WriterConfiguration configurations, List<BulkMessage<JSONObject>> messages) {
    BulkWriterResponse writerResponse = new BulkWriterResponse();
    List<Map.Entry<MessageId, Future>> results = new ArrayList<>();
    for (BulkMessage<JSONObject> bulkWriterMessage : messages) {
        MessageId messageId = bulkWriterMessage.getId();
        JSONObject message = bulkWriterMessage.getMessage();
        String jsonMessage;
        try {
            jsonMessage = message.toJSONString();
        } catch (Throwable t) {
            writerResponse.addError(t, messageId);
            continue;
        }
        Optional<String> topic = getKafkaTopic(message);
        if (topic.isPresent()) {
            Future future = kafkaProducer.send(new ProducerRecord<String, String>(topic.get(), jsonMessage));
            // we want to manage the batching
            results.add(new AbstractMap.SimpleEntry<>(messageId, future));
        } else {
            LOG.debug("Dropping {} because no topic is specified.", jsonMessage);
        }
    }
    Collection<MessageId> ids = messages.stream().map(BulkMessage::getId).collect(Collectors.toList());
    try {
        // ensures all Future.isDone() == true
        kafkaProducer.flush();
    } catch (InterruptException e) {
        writerResponse.addAllErrors(e, ids);
        return writerResponse;
    }
    for (Map.Entry<MessageId, Future> kv : results) {
        try {
            kv.getValue().get();
            writerResponse.addSuccess(kv.getKey());
        } catch (Exception e) {
            writerResponse.addError(e, kv.getKey());
        }
    }
    return writerResponse;
}

Example 8 with MessageId

use of org.apache.metron.common.writer.MessageId in project metron by apache.

the class BulkWriterComponentTest method flushShouldAckMissingTuples.

@Test
public void flushShouldAckMissingTuples() throws Exception {
    BulkWriterComponent<JSONObject> bulkWriterComponent = new BulkWriterComponent<>(Collections.singletonList(flushPolicy));
    BulkMessageWriter<JSONObject> bulkMessageWriter = mock(BulkMessageWriter.class);
    MessageId successId = new MessageId("successId");
    MessageId errorId = new MessageId("errorId");
    MessageId missingId = new MessageId("missingId");
    JSONObject successMessage = new JSONObject();
    successMessage.put("name", "success");
    JSONObject errorMessage = new JSONObject();
    errorMessage.put("name", "error");
    JSONObject missingMessage = new JSONObject();
    missingMessage.put("name", "missing");
    List<BulkMessage<JSONObject>> allMessages = new ArrayList<BulkMessage<JSONObject>>() {

        {
            add(new BulkMessage<>(successId, successMessage));
            add(new BulkMessage<>(errorId, errorMessage));
            add(new BulkMessage<>(missingId, missingMessage));
        }
    };
    BulkWriterResponse bulkWriterResponse = new BulkWriterResponse();
    bulkWriterResponse.addSuccess(successId);
    Throwable throwable = mock(Throwable.class);
    bulkWriterResponse.addError(throwable, errorId);
    when(bulkMessageWriter.write(sensorType, configurations, allMessages)).thenReturn(bulkWriterResponse);
    bulkWriterComponent.flush(sensorType, bulkMessageWriter, configurations, allMessages);
    BulkWriterResponse expectedResponse = new BulkWriterResponse();
    expectedResponse.addSuccess(successId);
    expectedResponse.addError(throwable, errorId);
    expectedResponse.addSuccess(missingId);
    verify(flushPolicy, times(1)).onFlush(sensorType, expectedResponse);
    verifyNoMoreInteractions(flushPolicy);
}

Example 9 with MessageId

use of org.apache.metron.common.writer.MessageId in project metron by apache.

the class BulkWriterComponent method flush.

/**
 * Flushes a batch for a sensor type by writing messages with the supplied {@link org.apache.metron.common.writer.BulkMessageWriter}.
 * Ensures all message ids in a batch are included in the response. After messages are written the cache is cleared and
 * flush policies are reset for that sensor type.
 * @param sensorType sensor type
 * @param bulkMessageWriter writer that will do the actual writing
 * @param configurations writer configurations
 * @param messages messages to be written
 */
protected void flush(String sensorType, BulkMessageWriter<MESSAGE_T> bulkMessageWriter, WriterConfiguration configurations, List<BulkMessage<MESSAGE_T>> messages) {
    // no need to mock, so use real time
    long startTime = System.currentTimeMillis();
    BulkWriterResponse response = new BulkWriterResponse();
    Collection<MessageId> ids = messages.stream().map(BulkMessage::getId).collect(Collectors.toList());
    try {
        response = bulkMessageWriter.write(sensorType, configurations, messages);
        // Make sure all ids are included in the BulkWriterResponse
        ids.removeAll(response.getSuccesses());
        response.getErrors().values().forEach(ids::removeAll);
        response.addAllSuccesses(ids);
    } catch (Throwable e) {
        response.addAllErrors(e, ids);
    } finally {
        onFlush(sensorType, response);
    }
    long endTime = System.currentTimeMillis();
    long elapsed = endTime - startTime;
    LOG.debug("Flushed batch successfully; sensorType={}, batchSize={}, took={} ms", sensorType, CollectionUtils.size(ids), elapsed);
}

Example 10 with MessageId

use of org.apache.metron.common.writer.MessageId in project metron by apache.

the class NoopWriter method write.

@Override
public BulkWriterResponse write(String sensorType, WriterConfiguration configurations, List<BulkMessage<JSONObject>> messages) throws Exception {
    if (sleepFunction != null) {
        sleepFunction.apply(null);
    }
    Set<MessageId> ids = messages.stream().map(BulkMessage::getId).collect(Collectors.toSet());
    BulkWriterResponse response = new BulkWriterResponse();
    response.addAllSuccesses(ids);
    return response;
}