Apache
/
nifi
mirror of https://github.com/apache/nifi.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fixes incorrect messages count in Provenance reporter. Adds Unit test to verify fix

This commit is contained in:
Oscar de la Pena 2015-03-19 10:10:09 +08:00
parent dea9e22475
commit eb5ec703ba
2 changed files with 114 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
nifi/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/processors/kafka/PutKafka.java
View File

@ -309,18 +309,13 @@ public class PutKafka extends AbstractProcessor {
data = Arrays.copyOfRange(baos.getUnderlyingBuffer(), 0, baos.size() - delimiterBytes.length);
}
createMessage: if ( data != null ) {
if ( data != null ) {
// If the message has no data, ignore it.
if ( data.length == 0 ) {
data = null;
baos.reset();
break createMessage;
}
if ( data.length != 0 ) {
// either we ran out of data or we reached the end of the message.
// Either way, create the message because it's ready to send.
final KeyedMessage<byte[], byte[]> message;
if ( key == null ) {
if (key == null) {
message = new KeyedMessage<>(topic, data);
} else {
message = new KeyedMessage<>(topic, keyBytes, data);
@ -331,7 +326,7 @@ public class PutKafka extends AbstractProcessor {
// and send the whole List.
messages.add(message);
messageBytes += data.length;
if ( messageBytes >= maxBufferSize || streamFinished ) {
if (messageBytes >= maxBufferSize || streamFinished) {
// send the messages, then reset our state.
try {
producer.send(messages);
@ -357,16 +352,18 @@ public class PutKafka extends AbstractProcessor {
// and send them to 'failure'.
lastMessageOffset.set(in.getBytesConsumed());
}
}
// reset BAOS so that we can start a new message.
baos.reset();
data = null;
}
}
// If there are messages left, send them
if ( !messages.isEmpty() ) {
try {
messagesSent.addAndGet(messages.size()); // add count of messages
producer.send(messages);
} catch (final Exception e) {
throw new ProcessException("Failed to send messages to Kafka", e);
@ -377,7 +374,6 @@ public class PutKafka extends AbstractProcessor {
});
final long nanos = System.nanoTime() - start;
session.getProvenanceReporter().send(flowFile, "kafka://" + topic, "Sent " + messagesSent.get() + " messages");
session.transfer(flowFile, REL_SUCCESS);
getLogger().info("Successfully sent {} messages to Kafka for {} in {} millis", new Object[] {messagesSent.get(), flowFile, TimeUnit.NANOSECONDS.toMillis(nanos)});

72
nifi/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/processors/kafka/TestPutKafka.java
View File

@ -25,6 +25,7 @@ import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import kafka.common.FailedToSendMessageException;
import kafka.javaapi.producer.Producer;
@ -32,12 +33,14 @@ import kafka.producer.KeyedMessage;
import kafka.producer.ProducerConfig;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.processor.ProcessContext;
import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.apache.nifi.provenance.ProvenanceReporter;
import org.apache.nifi.util.*;
import org.junit.Ignore;
import org.junit.Test;
import org.mockito.Mockito;
import org.mockito.internal.util.reflection.Whitebox;
public class TestPutKafka {
@ -140,6 +143,68 @@ public class TestPutKafka {
assertTrue(Arrays.equals("4".getBytes(), msgs.get(3)));
}
@Test
public void testProvenanceReporterMessagesCount(){
final TestableProcessor processor = new TestableProcessor();
ProvenanceReporter spyProvenanceReporter = Mockito.spy(new MockProvenanceReporter());
AtomicLong idGenerator = new AtomicLong(0L);
SharedSessionState sharedState = new SharedSessionState(processor, idGenerator);
Whitebox.setInternalState(sharedState, "provenanceReporter", spyProvenanceReporter);
MockFlowFileQueue flowFileQueue = sharedState.getFlowFileQueue();
MockSessionFactory sessionFactory = Mockito.mock(MockSessionFactory.class);
MockProcessSession mockProcessSession = new MockProcessSession(sharedState);
Mockito.when(sessionFactory.createSession()).thenReturn(mockProcessSession);
final TestRunner runner = TestRunners.newTestRunner(processor);
Whitebox.setInternalState(runner, "flowFileQueue", flowFileQueue);
Whitebox.setInternalState(runner, "sessionFactory", sessionFactory);
runner.setProperty(PutKafka.TOPIC, "topic1");
runner.setProperty(PutKafka.KEY, "key1");
runner.setProperty(PutKafka.SEED_BROKERS, "localhost:1234");
runner.setProperty(PutKafka.MESSAGE_DELIMITER, "\\n");
final byte[] bytes = "\n\n\n1\n2\n\n\n\n3\n4\n\n\n".getBytes();
runner.enqueue(bytes);
runner.run();
MockFlowFile mockFlowFile = mockProcessSession.getFlowFilesForRelationship(PutKafka.REL_SUCCESS).get(0);
Mockito.verify(spyProvenanceReporter, Mockito.atLeastOnce()).send(mockFlowFile, "kafka://topic1", "Sent 4 messages");
}
@Test
public void testProvenanceReporterWithoutDelimiterMessagesCount(){
final TestableProcessor processor = new TestableProcessor();
ProvenanceReporter spyProvenanceReporter = Mockito.spy(new MockProvenanceReporter());
AtomicLong idGenerator = new AtomicLong(0L);
SharedSessionState sharedState = new SharedSessionState(processor, idGenerator);
Whitebox.setInternalState(sharedState, "provenanceReporter", spyProvenanceReporter);
MockFlowFileQueue flowFileQueue = sharedState.getFlowFileQueue();
MockSessionFactory sessionFactory = Mockito.mock(MockSessionFactory.class);
MockProcessSession mockProcessSession = new MockProcessSession(sharedState);
Mockito.when(sessionFactory.createSession()).thenReturn(mockProcessSession);
final TestRunner runner = TestRunners.newTestRunner(processor);
Whitebox.setInternalState(runner, "flowFileQueue", flowFileQueue);
Whitebox.setInternalState(runner, "sessionFactory", sessionFactory);
runner.setProperty(PutKafka.TOPIC, "topic1");
runner.setProperty(PutKafka.KEY, "key1");
runner.setProperty(PutKafka.SEED_BROKERS, "localhost:1234");
final byte[] bytes = "\n\n\n1\n2\n\n\n\n3\n4\n\n\n".getBytes();
runner.enqueue(bytes);
runner.run();
MockFlowFile mockFlowFile = mockProcessSession.getFlowFilesForRelationship(PutKafka.REL_SUCCESS).get(0);
Mockito.verify(spyProvenanceReporter, Mockito.atLeastOnce()).send(mockFlowFile, "kafka://topic1");
}
@Test
@Ignore("Intended only for local testing; requires an actual running instance of Kafka & ZooKeeper...")
@ -233,4 +298,5 @@ public class TestPutKafka {
failAfter = successCount;
}
}
}