NEW FEATURE: spark-streaming support

Author: sunheehnus

src/main/scala/com/redislabs/provider/redis/redisFunctions.scala

@@ -1,12 +1,17 @@
 package com.redislabs.provider.redis
 
+import com.redislabs.provider.redis.streaming.RedisInputDStream
 import org.apache.spark.SparkContext
 import org.apache.spark.rdd.RDD
 
 import com.redislabs.provider.redis.rdd._
+import org.apache.spark.storage.StorageLevel
+import org.apache.spark.streaming.StreamingContext
+import redis.clients.util.RedisInputStream
 
 /**
   * RedisContext extends sparkContext's functionality with redis functions
+  *
   * @param sc a spark context
   */
 class RedisContext(@transient val sc: SparkContext) extends Serializable {
@@ -213,6 +218,12 @@ class RedisContext(@transient val sc: SparkContext) extends Serializable {
     }
   }
 
+  def createRedisStream(ssc: StreamingContext, keys: Array[String], storageLevel: StorageLevel)
+                       (implicit redisConfig: RedisConfig = new RedisConfig(new RedisEndpoint(sc.getConf))): Unit = {
+    new RedisInputDStream(ssc, keys, redisConfig, storageLevel)
+  }
+
+
 
   /**
     * @param kvs Pair RDD of K/V

src/main/scala/com/redislabs/provider/redis/streaming/RedisInputDStream.scala

@@ -0,0 +1,64 @@
+package com.redislabs.provider.redis.streaming
+
+import com.redislabs.provider.redis.RedisConfig
+import org.apache.curator.utils.ThreadUtils
+import org.apache.spark.storage.StorageLevel
+import org.apache.spark.streaming.StreamingContext
+import org.apache.spark.streaming.receiver.Receiver
+import org.apache.spark.streaming.dstream.ReceiverInputDStream
+
+import redis.clients.jedis._
+
+import scala.util.control.NonFatal
+
+import scala.collection.JavaConversions._
+
+class RedisInputDStream(_ssc: StreamingContext,
+                        keys: Array[String],
+                        redisConfig: RedisConfig,
+                        storageLevel: StorageLevel)
+  extends ReceiverInputDStream[(String, String)](_ssc) {
+  def getReceiver(): Receiver[(String, String)] = {
+    new RedisReceiver(keys, storageLevel, redisConfig)
+  }
+}
+
+private class RedisReceiver(keys: Array[String],
+                            storageLevel: StorageLevel,
+                            redisConfig: RedisConfig)
+  extends Receiver[(String, String)](storageLevel) {
+
+  def onStart() {
+    val executorPool = ThreadUtils.newFixedThreadPool(keys.length, "BlockLists Streamming")
+    try {
+      keys.foreach{ key =>
+        executorPool.submit(new MessageHandler(redisConfig.connectionForKey(key), key))
+      }
+    } finally {
+      executorPool.shutdown()
+    }
+
+  }
+
+  def onStop() {
+  }
+
+  private class MessageHandler(conn: Jedis, key: String) extends Runnable {
+    def run() {
+      val keys: Array[String] = Array(key)
+      try {
+        while(!isStopped) {
+          conn.blpop(keys: _*).toList.foreach{ x =>
+            store((key, x))
+          }
+        }
+      } catch {
+        case NonFatal(e) =>
+          restart("Error receiving data", e)
+      } finally {
+        onStop()
+      }
+    }
+  }
+}
+