Kinesis Source: Change from SynchronousQueue to CountdownLatch for back pressure (#142)

This fixes a problem where an app could OOM when Kinesis scales to add
more shards. It relates to the feature implemented in
#102.

It is a requirement of the Source that the ShardRecordProcessor gets
blocked until the downstream app is ready to consume an event. Before
this PR we used a SynchronousQueue to achieve this blocking. After this
PR we instead use a CountdownLatch, plus an unbounded queue.

This means we have better control over backpressure during the scenario
where the Source tries to handle many shard ends at the same time.

Author: istreeter

modules/kinesis/src/main/scala/com/snowplowanalytics/snowplow/streams/kinesis/source/KCLAction.scala

@@ -15,12 +15,40 @@ private sealed trait KCLAction
 
 private object KCLAction {
 
-  final case class ProcessRecords(shardId: String, processRecordsInput: ProcessRecordsInput) extends KCLAction
+  /**
+   * The action emitted by the ShardRecordProcessor when it receives new records
+   *
+   * @param await
+   *   A countdown latch used to backpressure the ShardRecordProcessor. The consumer of the queue
+   *   should release the countdown latch to unblock the ShardRecordProcessor and let it fetch more
+   *   records from Kinesis.
+   */
+  final case class ProcessRecords(
+    shardId: String,
+    await: CountDownLatch,
+    processRecordsInput: ProcessRecordsInput
+  ) extends KCLAction
+
+  /**
+   * The action emitted by the ShardRecordProcessor when it reaches a shard end.
+   *
+   * @param await
+   *   A countdown latch used to block the ShardRecordProcessor until all records from this shard
+   *   have been checkpointed.
+   *
+   * @note
+   *   Unlike the `await` in the `ProcessRecords` class, this countdown latch must not be released
+   *   immediately by the queue consumer. It must only be released by the checkpointer.
+   */
   final case class ShardEnd(
     shardId: String,
     await: CountDownLatch,
     shardEndedInput: ShardEndedInput
   ) extends KCLAction
-  final case class KCLError(t: Throwable) extends KCLAction
+
+  final case class KCLError(
+    t: Throwable,
+    await: CountDownLatch
+  ) extends KCLAction
 
 }

modules/kinesis/src/main/scala/com/snowplowanalytics/snowplow/streams/kinesis/source/KCLScheduler.scala

@@ -25,7 +25,7 @@ import software.amazon.kinesis.retrieval.polling.PollingConfig
 
 import java.net.URI
 import java.util.Date
-import java.util.concurrent.SynchronousQueue
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue}
 import java.util.concurrent.atomic.AtomicReference
 
 import com.snowplowanalytics.snowplow.streams.kinesis.KinesisSourceConfig
@@ -34,7 +34,7 @@ private[source] object KCLScheduler {
 
   def populateQueue[F[_]: Async](
     config: KinesisSourceConfig,
-    queue: SynchronousQueue[KCLAction],
+    queue: LinkedBlockingQueue[KCLAction],
     client: SdkAsyncHttpClient,
     awsUserAgent: Option[String]
   ): Resource[F, Unit] =
@@ -51,7 +51,7 @@ private[source] object KCLScheduler {
     dynamoDbClient: DynamoDbAsyncClient,
     cloudWatchClient: CloudWatchAsyncClient,
     kinesisConfig: KinesisSourceConfig,
-    queue: SynchronousQueue[KCLAction]
+    queue: LinkedBlockingQueue[KCLAction]
   ): F[Scheduler] =
     Sync[F].delay {
       val configsBuilder =
@@ -92,8 +92,12 @@ private[source] object KCLScheduler {
       val coordinatorConfig = configsBuilder.coordinatorConfig
         .workerStateChangeListener(new WorkerStateChangeListener {
           def onWorkerStateChange(newState: WorkerStateChangeListener.WorkerState): Unit = ()
-          override def onAllInitializationAttemptsFailed(e: Throwable): Unit =
-            queue.put(KCLAction.KCLError(e))
+          override def onAllInitializationAttemptsFailed(e: Throwable): Unit = {
+            val countDownLatch = new CountDownLatch(1)
+            queue.put(KCLAction.KCLError(e, countDownLatch))
+            countDownLatch.await()
+            ()
+          }
         })
 
       new Scheduler(

modules/kinesis/src/main/scala/com/snowplowanalytics/snowplow/streams/kinesis/source/KinesisSource.scala

@@ -17,7 +17,7 @@ import software.amazon.awssdk.http.async.SdkAsyncHttpClient
 import software.amazon.kinesis.lifecycle.events.{ProcessRecordsInput, ShardEndedInput}
 import software.amazon.kinesis.retrieval.kpl.ExtendedSequenceNumber
 
-import java.util.concurrent.{CountDownLatch, SynchronousQueue}
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue}
 import scala.concurrent.duration.{DurationLong, FiniteDuration}
 import scala.jdk.CollectionConverters._
 
@@ -46,34 +46,37 @@ private[kinesis] object KinesisSource {
       }
     }
 
-  // We enable fairness on the `SynchronousQueue` to ensure all Kinesis shards are sourced at an equal rate.
-  private val synchronousQueueFairness: Boolean = true
-
   private def kinesisStream[F[_]: Async](
     config: KinesisSourceConfig,
     client: SdkAsyncHttpClient,
     awsUserAgent: Option[String]
   ): Stream[F, Stream[F, Option[LowLevelEvents[Map[String, Checkpointable]]]]] = {
-    val actionQueue = new SynchronousQueue[KCLAction](synchronousQueueFairness)
+    val actionQueue = new LinkedBlockingQueue[KCLAction]()
     for {
       _ <- Stream.resource(KCLScheduler.populateQueue[F](config, actionQueue, client, awsUserAgent))
       events <- Stream.emit(pullFromQueueAndEmit(actionQueue).stream).repeat
     } yield events
   }
 
   private def pullFromQueueAndEmit[F[_]: Sync](
-    queue: SynchronousQueue[KCLAction]
+    queue: LinkedBlockingQueue[KCLAction]
   ): Pull[F, Option[LowLevelEvents[Map[String, Checkpointable]]], Unit] =
-    Pull.eval(pullFromQueue(queue)).flatMap { case PullFromQueueResult(actions, hasShardEnd) =>
+    Pull.eval(pullFromQueue(queue)).flatMap { actions =>
       val toEmit = actions.traverse {
-        case KCLAction.ProcessRecords(_, processRecordsInput) if processRecordsInput.records.asScala.isEmpty =>
-          Pull.output1(None)
-        case KCLAction.ProcessRecords(shardId, processRecordsInput) =>
-          Pull.output1(Some(provideNextChunk(shardId, processRecordsInput))).covary[F]
+        case KCLAction.ProcessRecords(_, await, processRecordsInput) if processRecordsInput.records.asScala.isEmpty =>
+          Pull.eval(Sync[F].delay(await.countDown())) >> Pull.output1(None)
+        case KCLAction.ProcessRecords(shardId, await, processRecordsInput) =>
+          Pull.eval(Sync[F].delay(await.countDown())) >> Pull.output1(Some(provideNextChunk(shardId, processRecordsInput))).covary[F]
         case KCLAction.ShardEnd(shardId, await, shardEndedInput) =>
+          // Do not call `await.countDown()` yet. It must be released later by the checkpointer.
           handleShardEnd[F](shardId, await, shardEndedInput)
-        case KCLAction.KCLError(t) =>
-          Pull.eval(Logger[F].error(t)("Exception from Kinesis source")) *> Pull.raiseError[F](t)
+        case KCLAction.KCLError(t, await) =>
+          Pull.eval(Sync[F].delay(await.countDown())) >> Pull.eval(Logger[F].error(t)("Exception from Kinesis source")) >> Pull
+            .raiseError[F](t)
+      }
+      val hasShardEnd = actions.exists {
+        case _: KCLAction.ShardEnd => true
+        case _: KCLAction          => false
       }
       if (hasShardEnd) {
         val log = Logger[F].info {
@@ -88,31 +91,21 @@ private[kinesis] object KinesisSource {
         toEmit *> pullFromQueueAndEmit(queue)
     }
 
-  private case class PullFromQueueResult(actions: NonEmptyList[KCLAction], hasShardEnd: Boolean)
-
-  private def pullFromQueue[F[_]: Sync](queue: SynchronousQueue[KCLAction]): F[PullFromQueueResult] =
-    resolveNextAction(queue)
-      .flatMap {
-        case shardEnd: KCLAction.ShardEnd =>
-          // If we reached the end of one shard, it is likely we reached the end of other shards too.
-          // Therefore pull more actions from the queue, to minimize the number of times we need to do
-          // an early close of the inner stream.
-          resolveAllActions(queue).map { more =>
-            PullFromQueueResult(NonEmptyList(shardEnd, more), hasShardEnd = true)
-          }
-        case other =>
-          PullFromQueueResult(NonEmptyList.one(other), hasShardEnd = false).pure[F]
-      }
+  private def pullFromQueue[F[_]: Sync](queue: LinkedBlockingQueue[KCLAction]): F[NonEmptyList[KCLAction]] =
+    for {
+      head <- resolveNextAction(queue)
+      tail <- resolveAllActions(queue)
+    } yield NonEmptyList(head, tail)
 
   /** Always returns a `KCLAction`, possibly waiting until one is available */
-  private def resolveNextAction[F[_]: Sync](queue: SynchronousQueue[KCLAction]): F[KCLAction] =
+  private def resolveNextAction[F[_]: Sync](queue: LinkedBlockingQueue[KCLAction]): F[KCLAction] =
     Sync[F].delay(Option[KCLAction](queue.poll)).flatMap {
       case Some(action) => Sync[F].pure(action)
       case None         => Sync[F].interruptible(queue.take)
     }
 
   /** Returns immediately, but the `List[KCLAction]` might be empty */
-  private def resolveAllActions[F[_]: Sync](queue: SynchronousQueue[KCLAction]): F[List[KCLAction]] =
+  private def resolveAllActions[F[_]: Sync](queue: LinkedBlockingQueue[KCLAction]): F[List[KCLAction]] =
     for {
       ret <- Sync[F].delay(new java.util.ArrayList[KCLAction]())
       _ <- Sync[F].delay(queue.drainTo(ret))

modules/kinesis/src/main/scala/com/snowplowanalytics/snowplow/streams/kinesis/source/ShardRecordProcessor.scala

@@ -16,13 +16,13 @@ import software.amazon.kinesis.lifecycle.events.{
 }
 import software.amazon.kinesis.processor.{ShardRecordProcessor => KCLShardProcessor}
 
-import java.util.concurrent.{CountDownLatch, SynchronousQueue}
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue}
 import java.util.concurrent.atomic.AtomicReference
 
 private[source] object ShardRecordProcessor {
 
   def apply(
-    queue: SynchronousQueue[KCLAction],
+    queue: LinkedBlockingQueue[KCLAction],
     currentShardIds: AtomicReference[Set[String]]
   ): KCLShardProcessor = new KCLShardProcessor {
     private var shardId: String = _
@@ -36,9 +36,15 @@ private[source] object ShardRecordProcessor {
         //   2. KCL re-aquires the lost lease for the same shard
         //   3. The original ShardRecordProcessor is not terminated until after KCL re-aquires the lease
         // This is a very unhealthy state, so we should kill the app.
-        val action = KCLAction.KCLError(new RuntimeException(s"Refusing to initialize a duplicate record processor for shard $shardId"))
+        val countDownLatch = new CountDownLatch(1)
+        val action = KCLAction.KCLError(
+          new RuntimeException(s"Refusing to initialize a duplicate record processor for shard $shardId"),
+          countDownLatch
+        )
         withHandledInterrupts {
           queue.put(action)
+          countDownLatch.await()
+          ()
         }
       }
     }
@@ -54,9 +60,12 @@ private[source] object ShardRecordProcessor {
     }
 
     override def processRecords(processRecordsInput: ProcessRecordsInput): Unit = {
-      val action = KCLAction.ProcessRecords(shardId, processRecordsInput)
+      val countDownLatch = new CountDownLatch(1)
+      val action         = KCLAction.ProcessRecords(shardId, countDownLatch, processRecordsInput)
       withHandledInterrupts {
         queue.put(action)
+        countDownLatch.await()
+        ()
       }
     }