STORM-2413: Make new Kafka spout respect tuple retry limit

Author: Stig Rohde Døssing

external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpout.java

@@ -330,7 +330,7 @@ private void emit() {
     //@return true if tuple was emitted
     private boolean emitTupleIfNotEmitted(ConsumerRecord<K, V> record) {
         final TopicPartition tp = new TopicPartition(record.topic(), record.partition());
-        final KafkaSpoutMessageId msgId = new KafkaSpoutMessageId(record);
+        final KafkaSpoutMessageId msgId = retryService.getMessageId(record);
         if (acked.containsKey(tp) && acked.get(tp).contains(msgId)) {   // has been acked
             LOG.trace("Tuple for record [{}] has already been acked. Skipping", record);
         } else if (emitted.contains(msgId)) {   // has been emitted and it's pending ack or fail

external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpoutMessageId.java

@@ -22,8 +22,8 @@
 import org.apache.kafka.common.TopicPartition;
 
 public class KafkaSpoutMessageId {
-    private transient TopicPartition topicPart;
-    private transient long offset;
+    private final transient TopicPartition topicPart;
+    private final transient long offset;
     private transient int numFails = 0;
 
     public KafkaSpoutMessageId(ConsumerRecord<?, ?> consumerRecord) {

external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpoutRetryExponentialBackoff.java

@@ -30,6 +30,7 @@
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.concurrent.TimeUnit;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.storm.utils.Time;
 
 /**
@@ -269,6 +270,19 @@ public boolean schedule(KafkaSpoutMessageId msgId) {
             return true;
         }
     }
+    
+    @Override
+    public KafkaSpoutMessageId getMessageId(ConsumerRecord<?, ?> record) {
+        KafkaSpoutMessageId msgId = new KafkaSpoutMessageId(record);
+        if(toRetryMsgs.contains(msgId)) {
+            for(KafkaSpoutMessageId originalMsgId : toRetryMsgs) {
+                if(originalMsgId.equals(msgId)) {
+                    return originalMsgId;
+                }
+            }
+        }
+        return msgId;
+    }
 
     // if value is greater than Long.MAX_VALUE it truncates to Long.MAX_VALUE
     private long nextTime(KafkaSpoutMessageId msgId) {

external/storm-kafka-client/src/main/java/org/apache/storm/kafka/spout/KafkaSpoutRetryService.java

@@ -22,7 +22,9 @@
 
 import java.io.Serializable;
 import java.util.Collection;
+import java.util.Optional;
 import java.util.Set;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 
 /**
  * Represents the logic that manages the retrial of failed tuples.
@@ -75,4 +77,11 @@ public interface KafkaSpoutRetryService extends Serializable {
      * Returns false is this message is not scheduled for retrial
      */
     boolean isScheduled(KafkaSpoutMessageId msgId);
+    
+    /**
+     * Gets the {@link KafkaSpoutMessageId} for the given record.
+     * @param record The record to fetch the id for
+     * @return The id the record was scheduled for retry with, or a new {@link KafkaSpoutMessageId} if the record was not scheduled for retry.
+     */
+    KafkaSpoutMessageId getMessageId(ConsumerRecord<?, ?> record);
 }

external/storm-kafka-client/src/test/java/org/apache/storm/kafka/spout/SingleTopicKafkaSpoutTest.java

@@ -17,7 +17,6 @@
  */
 package org.apache.storm.kafka.spout;
 
-
 import static org.apache.storm.kafka.spout.builders.SingleTopicKafkaSpoutConfiguration.getKafkaSpoutConfigBuilder;
 
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
@@ -41,6 +40,7 @@
 import static org.mockito.Matchers.anyObject;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -57,7 +57,6 @@
 import org.mockito.Captor;
 import org.mockito.MockitoAnnotations;
 
-
 public class SingleTopicKafkaSpoutTest {
 
     @Rule
@@ -73,12 +72,15 @@ public class SingleTopicKafkaSpoutTest {
     private KafkaConsumer<String, String> consumerSpy;
     private KafkaConsumerFactory<String, String> consumerFactory;
     private KafkaSpout<String, String> spout;
+    private int maxRetries = 3;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         KafkaSpoutConfig spoutConfig = getKafkaSpoutConfigBuilder(kafkaUnitRule.getKafkaUnit().getKafkaPort())
             .setOffsetCommitPeriodMs(commitOffsetPeriodMs)
+            .setRetry(new KafkaSpoutRetryExponentialBackoff(KafkaSpoutRetryExponentialBackoff.TimeInterval.seconds(0), KafkaSpoutRetryExponentialBackoff.TimeInterval.seconds(0),
+                maxRetries, KafkaSpoutRetryExponentialBackoff.TimeInterval.seconds(0)))
             .build();
         this.consumerSpy = spy(new KafkaConsumerFactoryDefault().createConsumer(spoutConfig));
         this.consumerFactory = (kafkaSpoutConfig) -> consumerSpy;
@@ -90,8 +92,8 @@ void populateTopicData(String topicName, int msgCount) throws InterruptedExcepti
 
         for (int i = 0; i < msgCount; i++) {
             ProducerRecord<String, String> producerRecord = new ProducerRecord<>(
-                    topicName, Integer.toString(i),
-                    Integer.toString(i));
+                topicName, Integer.toString(i),
+                Integer.toString(i));
             kafkaUnitRule.getKafkaUnit().sendMessage(producerRecord);
         }
     }
@@ -266,4 +268,26 @@ public void shouldReplayFirstTupleFailedOutOfOrder() throws Exception {
             verifyAllMessagesCommitted(messageCount);
         }
     }
+
+    @Test
+    public void shouldDropMessagesAfterMaxRetriesAreReached() throws Exception {
+        //Check that if one message fails repeatedly, the retry cap limits how many times the message can be reemitted
+        int messageCount = 1;
+        initializeSpout(messageCount);
+
+        //Emit and fail the same tuple until we've reached retry limit
+        for (int i = 0; i <= maxRetries; i++) {
+            ArgumentCaptor<Object> messageIdFailed = ArgumentCaptor.forClass(Object.class);
+            spout.nextTuple();
+            verify(collector).emit(anyObject(), anyObject(), messageIdFailed.capture());
+            KafkaSpoutMessageId msgId = (KafkaSpoutMessageId)messageIdFailed.getValue();
+            spout.fail(msgId);
+            assertThat("Expected message id number of failures to match the number of times the message has failed", msgId.numFails(), is(i+1));
+            reset(collector);
+        }
+
+        //Verify that the tuple is not emitted again
+        spout.nextTuple();
+        verify(collector, never()).emit(anyObject(), anyObject(), anyObject());
+    }
 }