Merge pull request gearpump#503 from manuzhang/fix_kafka_storage_consumer

clockfly · clockfly · commit fefad72dd3c4 · 2015-02-16T13:19:11.000+08:00
fix gearpump#502, fix KafkaStorage loading data
diff --git a/external/kafka/src/main/scala/org/apache/gearpump/streaming/kafka/lib/KafkaStorage.scala b/external/kafka/src/main/scala/org/apache/gearpump/streaming/kafka/lib/KafkaStorage.scala
@@ -28,6 +28,7 @@ import org.apache.kafka.clients.producer.{ProducerRecord, KafkaProducer}
 import org.apache.kafka.common.serialization.ByteArraySerializer
 import org.slf4j.Logger
 
+import scala.collection.mutable
 import scala.util.{Try, Failure, Success}
 
 object KafkaStorage {
@@ -87,23 +88,19 @@ private[kafka] class KafkaStorage(topic: String,
   }
 
   private[kafka] def load(consumer: KafkaConsumer): List[(TimeStamp, Array[Byte])] = {
-    @annotation.tailrec
-    def fetch(offsets: List[(TimeStamp, Array[Byte])]): List[(TimeStamp, Array[Byte])] = {
-      if (consumer.hasNext) {
-        val kafkaMsg = consumer.next
-        val offset = kafkaMsg.key.map { k =>
-          Injection.invert[TimeStamp, Array[Byte]](k) match {
-            case Success(time) => (time, kafkaMsg.msg)
-            case Failure(e) => throw e
-          }
-        } orElse (throw new RuntimeException("offset key should not be null"))
-        fetch(offsets :+ offset.get)
-      } else {
-        consumer.close()
-        offsets
-      }
+    var messagesBuilder = new mutable.ArrayBuilder.ofRef[(TimeStamp, Array[Byte])]
+    while (consumer.hasNext) {
+      val kafkaMsg = consumer.next
+      kafkaMsg.key.map { k =>
+        Injection.invert[TimeStamp, Array[Byte]](k) match {
+          case Success(time) =>
+            messagesBuilder += (time -> kafkaMsg.msg)
+          case Failure(e) => throw e
+        }
+      } orElse (throw new RuntimeException("offset key should not be null"))
     }
-    fetch(List.empty[(TimeStamp, Array[Byte])])
+    consumer.close()
+    messagesBuilder.result().toList
   }
 
 }
