-
Notifications
You must be signed in to change notification settings - Fork 644
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
78726bd
commit 49d6126
Showing
9 changed files
with
190 additions
and
12 deletions.
There are no files selected for viewing
18 changes: 18 additions & 0 deletions
18
kinesis/src/main/scala/akka/stream/alpakka/kinesis/CommittableRecord.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| /* | ||
| * Copyright (C) 2016-2017 Lightbend Inc. <http://www.lightbend.com> | ||
| */ | ||
| package akka.stream.alpakka.kinesis | ||
|
|
||
| import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer | ||
| import com.amazonaws.services.kinesis.model.Record | ||
|
|
||
| class CommittableRecord(val shardId: String, val record: Record, checkpointer: IRecordProcessorCheckpointer) { | ||
| val sequenceNumber = record.getSequenceNumber | ||
| def checkpoint(): Unit = checkpointer.checkpoint(record) | ||
| } | ||
|
|
||
| object CommittableRecord { | ||
|
|
||
| implicit val order: Ordering[CommittableRecord] = Ordering.by(_.sequenceNumber) | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
21 changes: 21 additions & 0 deletions
21
kinesis/src/main/scala/akka/stream/alpakka/kinesis/KinesisWorkerSourceSettings.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| /* | ||
| * Copyright (C) 2016-2017 Lightbend Inc. <http://www.lightbend.com> | ||
| */ | ||
| package akka.stream.alpakka.kinesis | ||
|
|
||
| import scala.concurrent.duration._ | ||
|
|
||
| case class KinesisWorkerSourceSettings(bufferSize: Int) | ||
| case class KinesisWorkerCheckpointSettings(maxBatchSize: Int, maxBatchWait: FiniteDuration) | ||
|
|
||
| object KinesisWorkerSourceSettings { | ||
|
|
||
| val defaultInstance = KinesisWorkerSourceSettings(1000) | ||
|
|
||
| } | ||
|
|
||
| object KinesisWorkerCheckpointSettings { | ||
|
|
||
| val defaultInstance = KinesisWorkerCheckpointSettings(1000, 10.seconds) | ||
|
|
||
| } |
82 changes: 82 additions & 0 deletions
82
kinesis/src/main/scala/akka/stream/alpakka/kinesis/KinesisWorkerSourceStage.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| /* | ||
| * Copyright (C) 2016-2017 Lightbend Inc. <http://www.lightbend.com> | ||
| */ | ||
| package akka.stream.alpakka.kinesis | ||
|
|
||
| import java.util | ||
| import java.util.concurrent.Semaphore | ||
|
|
||
| import akka.stream.alpakka.kinesis.KinesisErrors.UnexpectedShutdown | ||
| import akka.stream.stage._ | ||
| import akka.stream.{Attributes, Outlet, SourceShape} | ||
| import com.amazonaws.services.kinesis.clientlibrary.interfaces.{ | ||
| IRecordProcessor, | ||
| IRecordProcessorCheckpointer, | ||
| IRecordProcessorFactory | ||
| } | ||
| import com.amazonaws.services.kinesis.clientlibrary.lib.worker.{ShutdownReason, Worker} | ||
| import com.amazonaws.services.kinesis.model.Record | ||
|
|
||
| import scala.collection.JavaConverters._ | ||
| import scala.concurrent.{ExecutionContext, Future} | ||
|
|
||
| class KinesisWorkerSourceStage( | ||
| bufferSize: Int, | ||
| workerBuilder: IRecordProcessorFactory => Worker, | ||
| workerExecutor: ExecutionContext | ||
| ) extends GraphStageWithMaterializedValue[SourceShape[CommittableRecord], Worker] { | ||
|
|
||
| private val out: Outlet[CommittableRecord] = Outlet("records") | ||
| override val shape: SourceShape[CommittableRecord] = SourceShape(out) | ||
|
|
||
| override def createLogicAndMaterializedValue(inheritedAttributes: Attributes): (GraphStageLogic, Worker) = { | ||
| val logic = new Logic() | ||
| logic -> logic.worker | ||
| } | ||
|
|
||
| private class Logic extends GraphStageLogic(shape) with StageLogging { | ||
| private val buffer = new util.ArrayDeque[CommittableRecord]() | ||
| // We're transmitting backpressure to the worker with a semaphore instance | ||
| // worker callback ~> semaphore.acquire ~> push downstream ~> semaphore.release | ||
| private val semaphore = new Semaphore(bufferSize) | ||
|
|
||
| private[KinesisWorkerSourceStage] var worker: Worker = _ | ||
|
|
||
| private def tryToProduce(): Unit = | ||
| if (!buffer.isEmpty && isAvailable(out)) { | ||
| push(out, buffer.poll()) | ||
| semaphore.release() | ||
| } | ||
|
|
||
| override def preStart(): Unit = { | ||
| val callback = getAsyncCallback[CommittableRecord] { tuple => | ||
| buffer.offer(tuple) | ||
| tryToProduce() | ||
| } | ||
| val shutdownCallback = getAsyncCallback[ShutdownReason](reason => failStage(UnexpectedShutdown(reason))) | ||
| worker = workerBuilder( | ||
| new IRecordProcessorFactory { | ||
| override def createProcessor(): IRecordProcessor = new IRecordProcessor { | ||
| var shardId: String = _ | ||
| override def initialize(shardId: String): Unit = | ||
| this.shardId = shardId | ||
| override def processRecords(records: util.List[Record], checkpointer: IRecordProcessorCheckpointer): Unit = | ||
| records.asScala.foreach(r => callback.invoke(new CommittableRecord(shardId, r, checkpointer))) | ||
| override def shutdown(checkpointer: IRecordProcessorCheckpointer, reason: ShutdownReason): Unit = | ||
| shutdownCallback.invoke(reason) | ||
| } | ||
| } | ||
| ) | ||
| Future(worker.run())(workerExecutor) | ||
| } | ||
|
|
||
| override def postStop(): Unit = | ||
| worker.shutdown() | ||
|
|
||
| setHandler(out, new OutHandler { | ||
| override def onPull(): Unit = | ||
| tryToProduce() | ||
| }) | ||
| } | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
49 changes: 49 additions & 0 deletions
49
kinesis/src/main/scala/akka/stream/alpakka/kinesis/scaladsl/KinesisWorkerSource.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| /* | ||
| * Copyright (C) 2016-2017 Lightbend Inc. <http://www.lightbend.com> | ||
| */ | ||
| package akka.stream.alpakka.kinesis.scaladsl | ||
|
|
||
| import akka.NotUsed | ||
| import akka.stream.alpakka.kinesis.{ | ||
| CommittableRecord, | ||
| KinesisWorkerCheckpointSettings, | ||
| KinesisWorkerSourceSettings, | ||
| KinesisWorkerSourceStage | ||
| } | ||
| import akka.stream.scaladsl.{Flow, Source} | ||
| import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorFactory | ||
| import com.amazonaws.services.kinesis.clientlibrary.lib.worker.Worker | ||
| import com.amazonaws.services.kinesis.model.Record | ||
|
|
||
| import scala.concurrent.{ExecutionContext, Future} | ||
|
|
||
| object KinesisWorkerSource { | ||
|
|
||
| def apply( | ||
| workerBuilder: IRecordProcessorFactory => Worker, | ||
| settings: KinesisWorkerSourceSettings = KinesisWorkerSourceSettings.defaultInstance | ||
| )(implicit workerExecutor: ExecutionContext): Source[CommittableRecord, Worker] = | ||
| Source.fromGraph(new KinesisWorkerSourceStage(settings.bufferSize, workerBuilder, workerExecutor)) | ||
|
|
||
| def checkpointRecords(settings: KinesisWorkerCheckpointSettings = KinesisWorkerCheckpointSettings.defaultInstance)( | ||
| implicit checkpointExecutor: ExecutionContext | ||
| ): Flow[CommittableRecord, Record, NotUsed] = | ||
| Flow[CommittableRecord] | ||
| .groupBy(MAX_KINESIS_SHARDS, _.shardId) | ||
| .groupedWithin(settings.maxBatchSize, settings.maxBatchWait) | ||
| .map(records => records -> records.max) | ||
| .mapAsync(1) { | ||
| case (records, highestRecord) => | ||
| Future { | ||
| highestRecord.checkpoint() | ||
| records | ||
| } | ||
| } | ||
| .mapConcat(identity) | ||
| .map(_.record) | ||
| .mergeSubstreams | ||
|
|
||
| // http://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html | ||
| private val MAX_KINESIS_SHARDS = 500 | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters