Merge pull request #9 from JasonMWhite/streaming_listener

adding streaming listener

Author: Jason White

pom.xml

@@ -5,7 +5,7 @@
   <groupId>com.shopify</groupId>
   <artifactId>spark-datadog-relay</artifactId>
   <packaging>jar</packaging>
-  <version>0.2</version>
+  <version>0.3</version>
   <name>Datadog Relay for Spark</name>
 
   <properties>
@@ -47,6 +47,11 @@
 	  <version>2.2.4</version>
 	  <scope>test</scope>
 	</dependency>
+	<dependency>
+	  <groupId>org.apache.spark</groupId>
+	  <artifactId>spark-streaming_2.10</artifactId>
+	  <version>1.5.1</version>
+	</dependency>
   </dependencies>
 
   <build>

src/main/scala/org/apache/spark/DatadogRelay.scala

@@ -18,24 +18,7 @@ import scala.concurrent.ExecutionContext.Implicits.global
 
 class DatadogRelay(conf: SparkConf) extends SparkFirehoseListener {
 
-  val tags: List[String] = {
-    val datadogTags = conf.get("spark.datadog.tags", "")
-    if (datadogTags == "") List() else datadogTags.split(",").toList
-  }
-  
-  val statsdOption: Option[NonBlockingStatsDClient] = {
-    try {
-      Some(new NonBlockingStatsDClient(
-          "spark",
-          "localhost",
-          8125,
-          tags.mkString(",")
-      ))
-    } catch {
-      case ex: StatsDClientException => None
-      case ex: Exception => throw ex
-    }
-  }
+  val statsdOption: Option[NonBlockingStatsDClient] = SparkStatsDHelper.getClient(conf)
 
   def taskBaseMetrics(statsd: NonBlockingStatsDClient, e: SparkListenerTaskEnd): Unit = {
     statsd.incrementCounter("firehose.taskEnded")

src/main/scala/org/apache/spark/DatadogStreamingRelay.scala

@@ -0,0 +1,35 @@
+package org.apache.spark
+
+import org.apache.spark.streaming.scheduler._
+import com.timgroup.statsd.{NonBlockingStatsDClient, StatsDClientException}
+
+class DatadogStreamingRelay(conf: SparkConf) extends StreamingListener {
+  
+  val statsdOption: Option[NonBlockingStatsDClient] = SparkStatsDHelper.getClient(conf)
+  
+  override def onBatchStarted(batchStarted: StreamingListenerBatchStarted) {
+    statsdOption.foreach { statsd =>  
+      statsd.incrementCounter("firehose.batchStarted")
+    }
+  }
+  
+  override def onBatchCompleted(batchCompleted: StreamingListenerBatchCompleted) {
+    statsdOption.foreach { statsd =>
+      statsd.incrementCounter("firehose.batchCompleted")
+      statsd.count("firehose.batchRecordsProcessed", batchCompleted.batchInfo.numRecords)
+      batchCompleted.batchInfo.processingDelay.foreach { delay =>
+        statsd.recordExecutionTime("firehose.batchProcessingDelay", delay)
+      }
+      batchCompleted.batchInfo.schedulingDelay.foreach { delay =>
+        statsd.recordExecutionTime("firehose.batchSchedulingDelay", delay)
+      }
+      
+      batchCompleted.batchInfo.processingDelay.foreach { delay =>
+        statsd.recordExecutionTime("firehose.batchProcessingDelay", delay)
+      }
+      batchCompleted.batchInfo.totalDelay.foreach { delay =>
+        statsd.recordExecutionTime("firehose.batchTotalDelay", delay)
+      }
+    }
+  }
+}

src/main/scala/org/apache/spark/SparkStatsDHelper.scala

@@ -0,0 +1,24 @@
+package org.apache.spark
+
+import com.timgroup.statsd.{NonBlockingStatsDClient, StatsDClientException}
+
+object SparkStatsDHelper {
+  def tags(conf: SparkConf): List[String] = {
+    val datadogTags = conf.get("spark.datadog.tags", "")
+    if (datadogTags == "") List() else datadogTags.split(",").toList
+  }
+  
+  def getClient(conf: SparkConf): Option[NonBlockingStatsDClient] = {
+    try {
+      Some(new NonBlockingStatsDClient(
+          "spark",
+          "localhost",
+          8125,
+          tags(conf).mkString(",")
+      ))
+    } catch {
+      case ex: StatsDClientException => None
+      case ex: Exception => throw ex
+    }
+  }
+}