Refactor Source and remove unused code

Author: aserrallerios

README.md

@@ -36,7 +36,7 @@ In order to use it, you need to provide a Worker builder and the Source settings
 ```scala
 val workerSourceSettings = KinesisWorkerSourceSettings(
     bufferSize = 1000,
-    checkWorkerPeriodicity = 1 minute)
+    terminateStreamGracePeriod = 1 minute)
   val builder: IRecordProcessorFactory => Worker = { recordProcessorFactory =>
     new Worker.Builder()
       .recordProcessorFactory(recordProcessorFactory)
@@ -71,7 +71,7 @@ In order to use the Flow you can provide additional settings:
 
 ```scala
  val checkpointSettings = KinesisWorkerCheckpointSettings(100, 30 seconds)
-  KinesisWorker(builder, workerSourceSettings)
+  KinesisWorkerSource(builder, workerSourceSettings)
     .via(KinesisWorker.checkpointRecordsFlow(checkpointSettings))
     .to(Sink.ignore)
   KinesisWorker(builder, workerSourceSettings).to(

src/main/scala/aserralle/akka/stream/kcl/worker/CommittableRecord.scala src/main/scala/aserralle/akka/stream/kcl/CommittableRecord.scala

@@ -2,8 +2,9 @@
  * Copyright (C) 2018 Albert Serrallé
  */
 
-package aserralle.akka.stream.kcl.worker
+package aserralle.akka.stream.kcl
 
+import akka.Done
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.ShutdownReason
 import com.amazonaws.services.kinesis.clientlibrary.types.ExtendedSequenceNumber
@@ -26,13 +27,17 @@ class CommittableRecord(
     recordProcessor.shutdown
   def canBeCheckpointed(): Boolean =
     recordProcessorShutdownReason().isEmpty
-  def checkpoint(): Future[Unit] =
-    Future(checkpointer.checkpoint(record))
+  def tryToCheckpoint(): Future[Done] =
+    Future {
+      checkpointer.checkpoint(record)
+      Done
+    }
 
 }
 
 object CommittableRecord {
 
+  // Only makes sense to compare Records belonging to the same shard
   implicit val orderBySequenceNumber: Ordering[CommittableRecord] =
     Ordering.by(_.sequenceNumber)
 

src/main/scala/aserralle/akka/stream/kcl/Errors.scala

@@ -0,0 +1,15 @@
+/*
+ * Copyright (C) 2018 Albert Serrallé
+ */
+
+package aserralle.akka.stream.kcl
+
+import scala.util.control.NoStackTrace
+
+object Errors {
+
+  sealed trait KinesisWorkerSourceError extends NoStackTrace
+  case class WorkerUnexpectedShutdown(cause: Throwable)
+      extends KinesisWorkerSourceError
+
+}

src/main/scala/aserralle/akka/stream/kcl/worker/IRecordProcessor.scala src/main/scala/aserralle/akka/stream/kcl/IRecordProcessor.scala

@@ -2,8 +2,9 @@
  * Copyright (C) 2018 Albert Serrallé
  */
 
-package aserralle.akka.stream.kcl.worker
+package aserralle.akka.stream.kcl
 
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.ShutdownReason
 import com.amazonaws.services.kinesis.clientlibrary.types.{
   ExtendedSequenceNumber,
@@ -14,22 +15,25 @@ import com.amazonaws.services.kinesis.clientlibrary.types.{
 
 import scala.collection.JavaConverters._
 import scala.concurrent.ExecutionContext
+import scala.concurrent.duration.FiniteDuration
 
 private[kcl] class IRecordProcessor(
-    callback: CommittableRecord => Unit
+    callback: CommittableRecord => Unit,
+    terminateStreamGracePeriod: FiniteDuration
 )(implicit executionContext: ExecutionContext)
     extends com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessor {
   private var shardId: String = _
   private var extendedSequenceNumber: ExtendedSequenceNumber = _
-
   var shutdown: Option[ShutdownReason] = None
+  var latestCheckpointer: Option[IRecordProcessorCheckpointer] = None
 
   override def initialize(initializationInput: InitializationInput): Unit = {
-    this.shardId = initializationInput.getShardId
-    this.extendedSequenceNumber = initializationInput.getExtendedSequenceNumber
+    shardId = initializationInput.getShardId
+    extendedSequenceNumber = initializationInput.getExtendedSequenceNumber
   }
 
-  override def processRecords(processRecordsInput: ProcessRecordsInput): Unit =
+  override def processRecords(processRecordsInput: ProcessRecordsInput): Unit = {
+    latestCheckpointer = Some(processRecordsInput.getCheckpointer)
     processRecordsInput.getRecords.asScala.foreach { record =>
       callback(
         new CommittableRecord(
@@ -42,8 +46,17 @@ private[kcl] class IRecordProcessor(
         )
       )
     }
+  }
 
-  override def shutdown(shutdownInput: ShutdownInput): Unit =
+  override def shutdown(shutdownInput: ShutdownInput): Unit = {
     shutdown = Some(shutdownInput.getShutdownReason)
+    latestCheckpointer = Some(shutdownInput.getCheckpointer)
+    shutdownInput.getShutdownReason match {
+      case ShutdownReason.TERMINATE =>
+        Thread.sleep(terminateStreamGracePeriod.toMillis)
+      case ShutdownReason.ZOMBIE => ()
+      case ShutdownReason.REQUESTED => ()
+    }
+  }
 
 }

src/main/scala/aserralle/akka/stream/kcl/KinesisErrors.scala

@@ -8,7 +8,7 @@ import scala.concurrent.duration._
 
 case class KinesisWorkerSourceSettings(
     bufferSize: Int,
-    checkWorkerPeriodicity: FiniteDuration) {
+    terminateStreamGracePeriod: FiniteDuration) {
   require(
     bufferSize >= 1,
     "Buffer size must be greater than 0; use size 1 to disable stage buffering"
@@ -29,10 +29,9 @@ object KinesisWorkerSourceSettings {
   /**
     * Java API
     */
-  def create(
-      bufferSize: Int,
-      checkWorkerPeriodicity: FiniteDuration): KinesisWorkerSourceSettings =
-    KinesisWorkerSourceSettings(bufferSize, checkWorkerPeriodicity)
+  def create(bufferSize: Int, terminateStreamGracePeriod: FiniteDuration)
+    : KinesisWorkerSourceSettings =
+    KinesisWorkerSourceSettings(bufferSize, terminateStreamGracePeriod)
 
 }
 

src/main/scala/aserralle/akka/stream/kcl/KinesisWorkerSettings.scala

@@ -7,16 +7,15 @@ package aserralle.akka.stream.kcl.javadsl
 import java.util.concurrent.Executor
 
 import akka.NotUsed
-import aserralle.akka.stream.kcl.worker.CommittableRecord
-import aserralle.akka.stream.kcl.{scaladsl, _}
+import aserralle.akka.stream.kcl.{CommittableRecord, scaladsl, _}
 import akka.stream.javadsl.{Flow, Sink, Source}
 import com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessorFactory
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.Worker
 import com.amazonaws.services.kinesis.model.Record
 
 import scala.concurrent.ExecutionContext
 
-object KinesisWorker {
+object KinesisWorkerSource {
 
   abstract class WorkerBuilder {
     def build(r: IRecordProcessorFactory): Worker
@@ -26,32 +25,32 @@ object KinesisWorker {
       workerBuilder: WorkerBuilder,
       settings: KinesisWorkerSourceSettings,
       workerExecutor: Executor
-  ): Source[CommittableRecord, NotUsed] =
-    scaladsl.KinesisWorker
+  ): Source[CommittableRecord, Worker] =
+    scaladsl.KinesisWorkerSource
       .apply(workerBuilder.build, settings)(
         ExecutionContext.fromExecutor(workerExecutor))
       .asJava
 
   def create(
       workerBuilder: WorkerBuilder,
       workerExecutor: Executor
-  ): Source[CommittableRecord, NotUsed] =
+  ): Source[CommittableRecord, Worker] =
     create(workerBuilder,
            KinesisWorkerSourceSettings.defaultInstance,
            workerExecutor)
 
   def checkpointRecordsFlow(
       settings: KinesisWorkerCheckpointSettings
   ): Flow[CommittableRecord, Record, NotUsed] =
-    scaladsl.KinesisWorker.checkpointRecordsFlow(settings).asJava
+    scaladsl.KinesisWorkerSource.checkpointRecordsFlow(settings).asJava
 
   def checkpointRecordsFlow(): Flow[CommittableRecord, Record, NotUsed] =
     checkpointRecordsFlow(KinesisWorkerCheckpointSettings.defaultInstance)
 
   def checkpointRecordsSink(
       settings: KinesisWorkerCheckpointSettings
   ): Sink[CommittableRecord, NotUsed] =
-    scaladsl.KinesisWorker.checkpointRecordsSink(settings).asJava
+    scaladsl.KinesisWorkerSource.checkpointRecordsSink(settings).asJava
 
   def checkpointRecordsSink(): Sink[CommittableRecord, NotUsed] =
     checkpointRecordsSink(KinesisWorkerCheckpointSettings.defaultInstance)