add TimeExtractor and change api

manuzhang · manuzhang · commit 13fe4724049a · 2014-10-23T15:06:48.000+08:00
diff --git a/examples/src/main/scala/org/apache/gearpump/streaming/examples/kafka/KafkaSpout.scala b/examples/src/main/scala/org/apache/gearpump/streaming/examples/kafka/KafkaSpout.scala
@@ -101,10 +101,11 @@ class KafkaSpout(conf: Configs) extends TaskActor(conf) {
     @annotation.tailrec
     def fetchAndEmit(msgNum: Int, tpIndex: Int): Unit = {
       if (msgNum < emitBatchSize) {
-        val kafkaMsg = consumer.nextMessage(topicAndPartitions(tpIndex))
-        if (kafkaMsg != null) {
-          val timestamp = System.currentTimeMillis()
-          output(new Message(decoder.fromBytes(kafkaMsg.msg)))
+        val msgWithTime = consumer.nextMessageWithTime(topicAndPartitions(tpIndex))
+        if (msgWithTime != null) {
+          val kafkaMsg = msgWithTime._1
+          val timestamp =  msgWithTime._2
+          output(new Message(decoder.fromBytes(kafkaMsg.msg), timestamp))
           offsetManager.update(KafkaSource(kafkaMsg.topicAndPartition), timestamp, kafkaMsg.offset)
         } else {
           LOG.debug(s"no more messages from ${topicAndPartitions(tpIndex)}")
diff --git a/project/Build.scala b/project/Build.scala
@@ -29,7 +29,6 @@ object Build extends sbt.Build {
   val spraySwaggerVersion = "0.4.3"
   val swaggerUiVersion = "2.0.21"
 
-
   val commonSettings = Defaults.defaultSettings ++ packAutoSettings ++
     Seq(
       scalaVersion := scalaVersionNumber,
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/Checkpoint.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/Checkpoint.scala
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gearpump.streaming.transaction.api
+
+trait Source {
+  def name: String
+
+  def partition: Int
+}
+
+object Checkpoint {
+  def apply[K, V](records: List[(K, V)]): Checkpoint[K, V] = new Checkpoint(records)
+
+  def empty[K, V]: Checkpoint[K, V] = new Checkpoint(List.empty[(K, V)])
+}
+
+class Checkpoint[K, V](val records: List[(K, V)])
+
+trait CheckpointSerDe[K, V] {
+  def toKeyBytes(key: K): Array[Byte]
+  def fromKeyBytes(bytes: Array[Byte]): K
+
+  def toValueBytes(value: V): Array[Byte]
+  def fromValueBytes(bytes: Array[Byte]): V
+}
+
+trait CheckpointFilter[K, V] {
+  def filter(records: List[(K, V)], predicate: K): Option[(K, V)]
+}
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/CheckpointManager.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/CheckpointManager.scala
@@ -18,43 +18,11 @@
 
 package org.apache.gearpump.streaming.transaction.api
 
-object CheckpointManager {
-
-  trait Source {
-    def name: String
-
-    def partition: Int
-  }
-
- object Checkpoint {
-    def apply[K, V](records: List[(K, V)]): Checkpoint[K, V] = new Checkpoint(records)
-
-    def empty[K, V]: Checkpoint[K, V] = new Checkpoint(List.empty[(K, V)])
-  }
-
-  class Checkpoint[K, V](val records: List[(K, V)])
-
-  trait CheckpointSerDe[K, V] {
-    def toKeyBytes(key: K): Array[Byte]
-    def fromKeyBytes(bytes: Array[Byte]): K
-
-    def toValueBytes(value: V): Array[Byte]
-    def fromValueBytes(bytes: Array[Byte]): V
-  }
-
-  trait CheckpointFilter[K, V] {
-    def filter(records: List[(K, V)], predicate: K): Option[(K, V)]
-  }
-}
-
-
 /**
  * CheckpointManager checkpoints message and its timestamp to a persistent system
  * such that we could replay messages around or after given time
  */
 trait CheckpointManager[K, V] {
-  import org.apache.gearpump.streaming.transaction.api.CheckpointManager._
-
   def start(): Unit
 
   def register(sources: Array[Source]): Unit
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/OffsetFilter.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/OffsetFilter.scala
@@ -19,8 +19,6 @@
 package org.apache.gearpump.streaming.transaction.api
 
 import org.apache.gearpump.TimeStamp
-import org.apache.gearpump.streaming.transaction.api.CheckpointManager._
-import org.apache.gearpump.util.Configs
 
 class OffsetFilter extends CheckpointFilter[TimeStamp, Long] {
   override def filter(timeAndOffsets: List[(TimeStamp, Long)],
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/OffsetManager.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/OffsetManager.scala
@@ -19,7 +19,6 @@
 package org.apache.gearpump.streaming.transaction.api
 
 import org.apache.gearpump.TimeStamp
-import org.apache.gearpump.streaming.transaction.api.CheckpointManager._
 import org.apache.gearpump.streaming.transaction.kafka.KafkaUtil._
 import org.slf4j.{LoggerFactory, Logger}
 
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/TimeExtractor.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/api/TimeExtractor.scala
@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gearpump.streaming.transaction.api
+
+/**
+ * shamelessly stolen from summingbird
+ * https://github.com/twitter/summingbird/blob/develop/summingbird-core/src/main/scala/com/twitter/summingbird/TimeExtractor.scala
+ */
+object TimeExtractor {
+  def apply[T](fn: T => Long): TimeExtractor[T] =
+    new TimeExtractor[T] {
+      override def apply(t: T) = fn(t)
+    }
+}
+
+trait TimeExtractor[T] extends java.io.Serializable {
+  def apply(t: T): Long
+}
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaCheckpointManager.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaCheckpointManager.scala
@@ -20,8 +20,7 @@ package org.apache.gearpump.streaming.transaction.kafka
 
 import kafka.admin.AdminUtils
 import kafka.common.{TopicExistsException, TopicAndPartition}
-import org.apache.gearpump.streaming.transaction.api.CheckpointManager
-import org.apache.gearpump.streaming.transaction.api.CheckpointManager._
+import org.apache.gearpump.streaming.transaction.api.{CheckpointSerDe, Checkpoint, Source, CheckpointManager}
 import org.slf4j.{Logger, LoggerFactory}
 
 import scala.collection.mutable.{Map => MutableMap}
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaConfig.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaConfig.scala
@@ -27,7 +27,7 @@ import kafka.serializer.Decoder
 import kafka.utils.ZKStringSerializer
 import org.I0Itec.zkclient.ZkClient
 import org.I0Itec.zkclient.serialize.ZkSerializer
-import org.apache.gearpump.streaming.transaction.api.{OffsetFilter, CheckpointManagerFactory}
+import org.apache.gearpump.streaming.transaction.api.{TimeExtractor, OffsetFilter, CheckpointManagerFactory}
 import org.slf4j.{Logger, LoggerFactory}
 
 import scala.collection.JavaConversions._
@@ -88,9 +88,11 @@ object KafkaConfig {
                     receiveBufferSize: Int = getSocketReceiveBufferSize,
                     fetchSize: Int = getFetchMessageMaxBytes,
                     zkClient: ZkClient = getZkClient(),
-                    queueSize: Int = getConsumerQueueSize): KafkaConsumer = {
+                    queueSize: Int = getConsumerQueueSize,
+                    timeExtractor: TimeExtractor[KafkaMessage]
+                    = TimeExtractor(_ => System.currentTimeMillis())): KafkaConsumer = {
       new KafkaConsumer(topicAndPartitions, clientId, socketTimeout,
-        receiveBufferSize, fetchSize, zkClient, queueSize)
+        receiveBufferSize, fetchSize, zkClient, queueSize, timeExtractor)
     }
 
     def getZookeeperConnect = {
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaConsumer.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaConsumer.scala
@@ -21,6 +21,8 @@ package org.apache.gearpump.streaming.transaction.kafka
 import kafka.common.TopicAndPartition
 
 import org.I0Itec.zkclient.ZkClient
+import org.apache.gearpump.TimeStamp
+import org.apache.gearpump.streaming.transaction.api.TimeExtractor
 import org.slf4j.{Logger, LoggerFactory}
 import java.util.concurrent.LinkedBlockingQueue
 
@@ -46,7 +48,8 @@ object KafkaConsumer {
 class KafkaConsumer(topicAndPartitions: Array[TopicAndPartition],
                     clientId: String, socketTimeout: Int,
                     receiveBufferSize: Int, fetchSize: Int,
-                    zkClient: ZkClient, queueSize: Int)  {
+                    zkClient: ZkClient, queueSize: Int,
+                    timeExtractor: TimeExtractor[KafkaMessage])  {
   import org.apache.gearpump.streaming.transaction.kafka.KafkaConsumer._
 
   private val leaders: Map[TopicAndPartition, Broker] = topicAndPartitions.map {
@@ -69,14 +72,14 @@ class KafkaConsumer(topicAndPartitions: Array[TopicAndPartition],
 
   private var noMessages: Set[TopicAndPartition] = Set.empty[TopicAndPartition]
   private val noMessageSleepMS = 100
-  private val incomingQueue = topicAndPartitions.map(_ -> new LinkedBlockingQueue[KafkaMessage](queueSize)).toMap
+  private val incomingQueue = topicAndPartitions.map(_ -> new LinkedBlockingQueue[(KafkaMessage, TimeStamp)](queueSize)).toMap
 
   private val fetchThread = new Thread {
     override def run(): Unit = {
       while (!Thread.currentThread.isInterrupted) {
         val msg = fetchMessage()
         if (msg != null) {
-          incomingQueue(msg.topicAndPartition).put(msg)
+          incomingQueue(msg.topicAndPartition).put(msg, timeExtractor(msg))
         } else if (noMessages.size == topicAndPartitions.size) {
           LOG.debug(s"no messages for all TopicAndPartitions. sleeping for ${noMessageSleepMS} ms")
           Thread.sleep(noMessageSleepMS)
@@ -93,7 +96,7 @@ class KafkaConsumer(topicAndPartitions: Array[TopicAndPartition],
     iterators(topicAndPartition).setStartOffset(startOffset)
   }
 
-  def nextMessage(topicAndPartition: TopicAndPartition): KafkaMessage = {
+  def nextMessageWithTime(topicAndPartition: TopicAndPartition): (KafkaMessage, TimeStamp) = {
     incomingQueue(topicAndPartition).poll()
   }
 
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaSource.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaSource.scala
@@ -18,8 +18,8 @@
 
 package org.apache.gearpump.streaming.transaction.kafka
 
-import org.apache.gearpump.streaming.transaction.api.CheckpointManager.Source
 import kafka.common.TopicAndPartition
+import org.apache.gearpump.streaming.transaction.api.Source
 
 object KafkaSource {
   def apply(name: String, partition: Int): KafkaSource = {
diff --git a/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaUtil.scala b/streaming/src/main/scala/org/apache/gearpump/streaming/transaction/kafka/KafkaUtil.scala
@@ -18,7 +18,6 @@
 
 package org.apache.gearpump.streaming.transaction.kafka
 
-import kafka.common.TopicAndPartition
 import kafka.utils.ZkUtils
 import org.apache.gearpump.streaming.transaction.kafka.KafkaConsumer.Broker
 import org.I0Itec.zkclient.ZkClient
