address

Author: brkyvz

extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala

@@ -54,7 +54,7 @@ private[kinesis] class KinesisTestUtils extends Logging {
   @volatile
   private var _streamName: String = _
 
-  private lazy val kinesisClient = {
+  protected lazy val kinesisClient = {
     val client = new AmazonKinesisClient(KinesisTestUtils.getAWSCredentials())
     client.setEndpoint(endpointUrl)
     client
@@ -66,14 +66,11 @@ private[kinesis] class KinesisTestUtils extends Logging {
     new DynamoDB(dynamoDBClient)
   }
 
-  /** Left as a protected val so that we don't need to depend on KPL outside of tests. */
-  protected val kplProducer: KinesisProducer = null
-
-  private def getProducer(aggregate: Boolean): KinesisProducer = {
-    if (aggregate) {
-      kplProducer
+  protected def getProducer(aggregate: Boolean): KinesisDataGenerator = {
+    if (!aggregate) {
+      new SimpleDataGenerator(kinesisClient)
     } else {
-      new KinesisClientProducer(kinesisClient)
+      throw new UnsupportedOperationException("Aggregation is not supported through this code path")
     }
   }
 
@@ -106,12 +103,7 @@ private[kinesis] class KinesisTestUtils extends Logging {
   def pushData(testData: Seq[Int], aggregate: Boolean): Map[String, Seq[(Int, String)]] = {
     require(streamCreated, "Stream not yet created, call createStream() to create one")
     val producer = getProducer(aggregate)
-
-    testData.foreach { num =>
-      producer.putRecord(streamName, num)
-    }
-
-    val shardIdToSeqNumbers = producer.flush()
+    val shardIdToSeqNumbers = producer.sendData(streamName, testData)
     logInfo(s"Pushed $testData:\n\t ${shardIdToSeqNumbers.mkString("\n\t")}")
     shardIdToSeqNumbers.toMap
   }
@@ -239,31 +231,30 @@ private[kinesis] object KinesisTestUtils {
 }
 
 /** A wrapper interface that will allow us to consolidate the code for synthetic data generation. */
-private[kinesis] trait KinesisProducer {
-  /** Sends the data to Kinesis possibly with aggregation if KPL is used. */
-  def putRecord(streamName: String, num: Int): Unit
-
-  /** Flush all data in the buffer and return the metadata for everything that has been sent. */
-  def flush(): Map[String, Seq[(Int, String)]]
+private[kinesis] trait KinesisDataGenerator {
+  /** Sends the data to Kinesis and returns the metadata for everything that has been sent. */
+  def sendData(streamName: String, data: Seq[Int]): Map[String, Seq[(Int, String)]]
 }
 
-private[kinesis] class KinesisClientProducer(client: AmazonKinesisClient) extends KinesisProducer {
-  private val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
-
-  override def putRecord(streamName: String, num: Int): Unit = {
-    val str = num.toString
-    val data = ByteBuffer.wrap(str.getBytes())
-    val putRecordRequest = new PutRecordRequest().withStreamName(streamName)
-      .withData(data)
-      .withPartitionKey(str)
-
-    val putRecordResult = client.putRecord(putRecordRequest)
-    val shardId = putRecordResult.getShardId
-    val seqNumber = putRecordResult.getSequenceNumber()
-    val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
-      new ArrayBuffer[(Int, String)]())
-    sentSeqNumbers += ((num, seqNumber))
-  }
+private[kinesis] class SimpleDataGenerator(
+    client: AmazonKinesisClient) extends KinesisDataGenerator {
+  override def sendData(streamName: String, data: Seq[Int]): Map[String, Seq[(Int, String)]] = {
+    val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
+    data.foreach { num =>
+      val str = num.toString
+      val data = ByteBuffer.wrap(str.getBytes())
+      val putRecordRequest = new PutRecordRequest().withStreamName(streamName)
+        .withData(data)
+        .withPartitionKey(str)
+
+      val putRecordResult = client.putRecord(putRecordRequest)
+      val shardId = putRecordResult.getShardId
+      val seqNumber = putRecordResult.getSequenceNumber()
+      val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
+        new ArrayBuffer[(Int, String)]())
+      sentSeqNumbers += ((num, seqNumber))
+    }
 
-  override def flush(): Map[String, Seq[(Int, String)]] = shardIdToSeqNumbers.toMap
+    shardIdToSeqNumbers.toMap
+  }
 }

extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/ExtendedKinesisTestUtils.scala renamed to extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KPLBasedKinesisTestUtils.scala

@@ -24,14 +24,18 @@ import scala.collection.mutable.ArrayBuffer
 import com.amazonaws.services.kinesis.producer.{KinesisProducer => KPLProducer, KinesisProducerConfiguration, UserRecordResult}
 import com.google.common.util.concurrent.{FutureCallback, Futures}
 
-private[kinesis] class ExtendedKinesisTestUtils extends KinesisTestUtils {
-  override protected val kplProducer: KinesisProducer = {
-    new KinesisProducerLibraryProducer(regionName)
+private[kinesis] class KPLBasedKinesisTestUtils extends KinesisTestUtils {
+  override protected def getProducer(aggregate: Boolean): KinesisDataGenerator = {
+    if (!aggregate) {
+      new SimpleDataGenerator(kinesisClient)
+    } else {
+      new KPLDataGenerator(regionName)
+    }
   }
 }
 
 /** A wrapper for the KinesisProducer provided in the KPL. */
-private[kinesis] class KinesisProducerLibraryProducer(regionName: String) extends KinesisProducer {
+private[kinesis] class KPLDataGenerator(regionName: String) extends KinesisDataGenerator {
 
   private lazy val producer: KPLProducer = {
     val conf = new KinesisProducerConfiguration()
@@ -43,30 +47,26 @@ private[kinesis] class KinesisProducerLibraryProducer(regionName: String) extend
     new KPLProducer(conf)
   }
 
-  private val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
-
-  override def putRecord(streamName: String, num: Int): Unit = {
-    val str = num.toString
-    val data = ByteBuffer.wrap(str.getBytes())
-    val future = producer.addUserRecord(streamName, str, data)
-    val kinesisCallBack = new FutureCallback[UserRecordResult]() {
-      override def onFailure(t: Throwable): Unit = {} // do nothing
+  override def sendData(streamName: String, data: Seq[Int]): Map[String, Seq[(Int, String)]] = {
+    val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
+    data.foreach { num =>
+      val str = num.toString
+      val data = ByteBuffer.wrap(str.getBytes())
+      val future = producer.addUserRecord(streamName, str, data)
+      val kinesisCallBack = new FutureCallback[UserRecordResult]() {
+        override def onFailure(t: Throwable): Unit = {} // do nothing
 
-      override def onSuccess(result: UserRecordResult): Unit = {
-        val shardId = result.getShardId
-        val seqNumber = result.getSequenceNumber()
-        val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
-          new ArrayBuffer[(Int, String)]())
-        sentSeqNumbers += ((num, seqNumber))
+        override def onSuccess(result: UserRecordResult): Unit = {
+          val shardId = result.getShardId
+          val seqNumber = result.getSequenceNumber()
+          val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
+            new ArrayBuffer[(Int, String)]())
+          sentSeqNumbers += ((num, seqNumber))
+        }
       }
+      Futures.addCallback(future, kinesisCallBack)
     }
-
-    Futures.addCallback(future, kinesisCallBack)
     producer.flushSync()
-  }
-
-  override def flush(): Map[String, Seq[(Int, String)]] = {
     shardIdToSeqNumbers.toMap
   }
-
 }

extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDDSuite.scala

@@ -40,7 +40,7 @@ abstract class KinesisBackedBlockRDDTests(aggregateTestData: Boolean)
 
   override def beforeAll(): Unit = {
     runIfTestsEnabled("Prepare KinesisTestUtils") {
-      testUtils = new ExtendedKinesisTestUtils()
+      testUtils = new KPLBasedKinesisTestUtils()
       testUtils.createStream()
 
       shardIdToDataAndSeqNumbers = testUtils.pushData(testData, aggregate = aggregateTestData)

extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisStreamSuite.scala

@@ -63,7 +63,7 @@ abstract class KinesisStreamTests(aggregateTestData: Boolean) extends KinesisFun
     sc = new SparkContext(conf)
 
     runIfTestsEnabled("Prepare KinesisTestUtils") {
-      testUtils = new ExtendedKinesisTestUtils()
+      testUtils = new KPLBasedKinesisTestUtils()
       testUtils.createStream()
     }
   }