Graceful Producer Shutdown (#72)

etspaceman · Feb 28, 2023 · 0bb546a · 0bb546a
1 parent 782d468
commit 0bb546a
Show file tree

Hide file tree

Showing 16 changed files with 69 additions and 57 deletions.
diff --git a/integration-tests/src/test/scalajvm/kinesis4cats/kcl/KCLConsumerSpec.scala b/integration-tests/src/test/scalajvm/kinesis4cats/kcl/KCLConsumerSpec.scala
@@ -19,7 +19,7 @@ package kcl
 
 import scala.concurrent.duration._
 
-import cats.effect.kernel.Deferred
+import cats.effect.Deferred
 import cats.effect.std.Queue
 import cats.effect.{IO, Resource, SyncIO}
 import cats.syntax.all._

diff --git a/integration-tests/src/test/scalajvm/kinesis4cats/kcl/fs2/KCLConsumerFS2Spec.scala b/integration-tests/src/test/scalajvm/kinesis4cats/kcl/fs2/KCLConsumerFS2Spec.scala
@@ -21,7 +21,7 @@ package fs2
 import scala.concurrent.duration._
 
 import _root_.fs2.Stream
-import cats.effect.kernel.Deferred
+import cats.effect.Deferred
 import cats.effect.{IO, Resource, SyncIO}
 import cats.syntax.all._
 import io.circe.parser._

diff --git a/...on-tests/src/test/scalajvm/kinesis4cats/kcl/fs2/multistream/KCLConsumerFS2MultiSpec.scala b/...on-tests/src/test/scalajvm/kinesis4cats/kcl/fs2/multistream/KCLConsumerFS2MultiSpec.scala
@@ -20,7 +20,7 @@ package kcl.fs2.multistream
 import scala.concurrent.duration._
 
 import _root_.fs2.Stream
-import cats.effect.kernel.Deferred
+import cats.effect.Deferred
 import cats.effect.syntax.all._
 import cats.effect.{IO, Resource, SyncIO}
 import cats.syntax.all._

diff --git a/integration-tests/src/test/scalajvm/kinesis4cats/kcl/multistream/KCLConsumerMultiSpec.scala b/integration-tests/src/test/scalajvm/kinesis4cats/kcl/multistream/KCLConsumerMultiSpec.scala
@@ -20,7 +20,7 @@ package multistream
 
 import scala.concurrent.duration._
 
-import cats.effect.kernel.Deferred
+import cats.effect.Deferred
 import cats.effect.std.Queue
 import cats.effect.syntax.all._
 import cats.effect.{IO, Resource, SyncIO}

diff --git a/integration-tests/src/test/scalajvm/kinesis4cats/smithy4s/client/KinesisClientJVMSpec.scala b/integration-tests/src/test/scalajvm/kinesis4cats/smithy4s/client/KinesisClientJVMSpec.scala
@@ -16,9 +16,9 @@
 
 package kinesis4cats.smithy4s.client
 
+import cats.effect.Async
 import cats.effect.IO
 import cats.effect.SyncIO
-import cats.effect.kernel.Async
 import com.amazonaws.kinesis.Kinesis
 import org.http4s.blaze.client.BlazeClientBuilder
 

diff --git a/kcl/src/main/scala/kinesis4cats/kcl/KCLConsumer.scala b/kcl/src/main/scala/kinesis4cats/kcl/KCLConsumer.scala
@@ -16,7 +16,7 @@
 
 package kinesis4cats.kcl
 
-import cats.effect.kernel.Deferred
+import cats.effect.Deferred
 import cats.effect.syntax.all._
 import cats.effect.{Async, Ref, Resource}
 import cats.syntax.all._

diff --git a/kcl/src/main/scala/kinesis4cats/kcl/fs2/KCLConsumerFS2.scala b/kcl/src/main/scala/kinesis4cats/kcl/fs2/KCLConsumerFS2.scala
@@ -19,7 +19,7 @@ package kinesis4cats.kcl.fs2
 import scala.concurrent.duration._
 
 import cats.Parallel
-import cats.effect.kernel.Deferred
+import cats.effect.Deferred
 import cats.effect.std.Queue
 import cats.effect.syntax.all._
 import cats.effect.{Async, Ref, Resource}

diff --git a/kinesis-client/src/main/scala/kinesis4cats/client/producer/KinesisProducer.scala b/kinesis-client/src/main/scala/kinesis4cats/client/producer/KinesisProducer.scala
@@ -22,8 +22,8 @@ import scala.jdk.CollectionConverters._
 import java.time.Instant
 
 import cats.data.NonEmptyList
+import cats.effect.Resource
 import cats.effect._
-import cats.effect.kernel.Resource
 import cats.effect.syntax.all._
 import cats.syntax.all._
 import org.typelevel.log4cats.StructuredLogger

diff --git a/kinesis-client/src/main/scala/kinesis4cats/client/producer/fs2/FS2KinesisProducer.scala b/kinesis-client/src/main/scala/kinesis4cats/client/producer/fs2/FS2KinesisProducer.scala
@@ -18,9 +18,8 @@ package kinesis4cats.client
 package producer
 package fs2
 
-import cats.effect.Async
-import cats.effect.kernel.Resource
-import cats.effect.std.Queue
+import _root_.fs2.concurrent.Channel
+import cats.effect._
 import cats.effect.syntax.all._
 import org.typelevel.log4cats.StructuredLogger
 import org.typelevel.log4cats.slf4j.Slf4jLogger
@@ -36,9 +35,9 @@ import kinesis4cats.producer.fs2.FS2Producer
   *
   * @param config
   *   [[kinesis.producer.fs2.FS2Producer.Config FS2Producer.Config]]
-  * @param queue
-  *   [[cats.effect.std.Queue Queue]] of
-  *   [[kinesis4cats.producer.Record Records]] to produce.
+  * @param channel
+  *   [[https://github.com/typelevel/fs2/blob/main/core/shared/src/main/scala/fs2/concurrent/Channel.scala Channel]]
+  *   of [[kinesis4cats.producer.Record Records]] to produce.
   * @param underlying
   *   [[kinesis4cats.smithy4s.client.producer.KinesisProducer KinesisProducer]]
   * @param callback:
@@ -49,7 +48,7 @@ import kinesis4cats.producer.fs2.FS2Producer
 final class FS2KinesisProducer[F[_]] private[kinesis4cats] (
     override val logger: StructuredLogger[F],
     override val config: FS2Producer.Config,
-    override protected val queue: Queue[F, Option[Record]],
+    override protected val channel: Channel[F, Record],
     override protected val underlying: KinesisProducer[F]
 )(
     override protected val callback: (
@@ -101,11 +100,10 @@ object FS2KinesisProducer {
       config.producerConfig,
       _underlying
     )
-    queue <- Queue.bounded[F, Option[Record]](config.queueSize).toResource
-    producer = new FS2KinesisProducer[F](logger, config, queue, underlying)(
+    channel <- Channel.bounded[F, Record](config.queueSize).toResource
+    producer = new FS2KinesisProducer[F](logger, config, channel, underlying)(
       callback
     )
-    _ <- producer.start()
-    _ <- Resource.onFinalize(producer.stop())
+    _ <- producer.resource
   } yield producer
 }
diff --git a/kpl-ciris/src/main/scala/kinesis4cats/kpl/ciris/KPLCiris.scala b/kpl-ciris/src/main/scala/kinesis4cats/kpl/ciris/KPLCiris.scala
@@ -21,7 +21,7 @@ import scala.concurrent.duration._
 
 import _root_.ciris._
 import cats.effect.Async
-import cats.effect.kernel.Resource
+import cats.effect.Resource
 import cats.syntax.all._
 import com.amazonaws.auth.AWSCredentialsProvider
 import com.amazonaws.regions.Regions

diff --git a/shared/src/main/scala/kinesis4cats/producer/fs2/FS2Producer.scala b/shared/src/main/scala/kinesis4cats/producer/fs2/FS2Producer.scala
@@ -19,12 +19,10 @@ package fs2
 
 import scala.concurrent.duration._
 
-import _root_.fs2.Stream
+import _root_.fs2.concurrent.Channel
 import cats.data.Ior
 import cats.data.NonEmptyList
-import cats.effect.Async
-import cats.effect.Resource
-import cats.effect.std.Queue
+import cats.effect._
 import cats.effect.syntax.all._
 import cats.syntax.all._
 import org.typelevel.log4cats.StructuredLogger
@@ -53,7 +51,7 @@ abstract class FS2Producer[F[_], PutReq, PutRes](implicit
 
   /** The underlying queue of records to process
     */
-  protected def queue: Queue[F, Option[Record]]
+  protected def channel: Channel[F, Record]
 
   /** A user defined function that can be run against the results of a request
     */
@@ -70,28 +68,45 @@ abstract class FS2Producer[F[_], PutReq, PutRes](implicit
     * @return
     *   F of Unit
     */
-  def put(record: Record): F[Unit] = queue.offer(Some(record))
+  def put(record: Record): F[Unit] = {
+    val ctx = LogContext()
+
+    for {
+      _ <- logger.debug(ctx.context)("Received record to put")
+      res <- channel.send(record)
+      _ <- res.bitraverse(
+        _ =>
+          logger.warn(ctx.context)(
+            "Producer has been shut down and will not accept further requests"
+          ),
+        _ =>
+          logger.debug(ctx.context)(
+            "Successfully put record into processing queue"
+          )
+      )
+    } yield ()
+  }
 
   /** Stop the processing of records
     */
-  private[kinesis4cats] def stop(): F[Unit] = {
+  private[kinesis4cats] def stop(f: Fiber[F, Throwable, Unit]): F[Unit] = {
     val ctx = LogContext()
     for {
       _ <- logger.debug(ctx.context)("Stopping the FS2KinesisProducer")
-      _ <- queue.offer(None)
+      _ <- channel.close
+      _ <- f.join.void.timeoutTo(config.gracefulShutdownWait, f.cancel)
     } yield ()
   }
 
   /** Start the processing of records
     */
-  private[kinesis4cats] def start(): Resource[F, Unit] = {
+  private[kinesis4cats] def start(): F[Unit] = {
     val ctx = LogContext()
+
     for {
       _ <- logger
         .debug(ctx.context)("Starting the FS2KinesisProducer")
-        .toResource
-      _ <- Stream
-        .fromQueueNoneTerminated(queue)
+      _ <- channel.stream
         .groupWithin(config.putMaxChunk, config.putMaxWait)
         .evalMap { x =>
           val c = ctx.addEncoded("batchSize", x.size)
@@ -116,12 +131,12 @@ abstract class FS2Producer[F[_], PutReq, PutRes](implicit
         }
         .compile
         .drain
-        .background
-        .void
     } yield ()
-
   }
 
+  private[kinesis4cats] def resource: Resource[F, Unit] =
+    Resource.make(start().start)(stop).void
+
 }
 
 object FS2Producer {
@@ -149,7 +164,8 @@ object FS2Producer {
       putMaxWait: FiniteDuration,
       putMaxRetries: Option[Int],
       putRetryInterval: FiniteDuration,
-      producerConfig: Producer.Config
+      producerConfig: Producer.Config,
+      gracefulShutdownWait: FiniteDuration
   )
 
   object Config {
@@ -159,7 +175,8 @@ object FS2Producer {
       100.millis,
       Some(5),
       0.seconds,
-      Producer.Config.default(streamNameOrArn)
+      Producer.Config.default(streamNameOrArn),
+      30.seconds
     )
   }
 

diff --git a/...tack/src/main/scala/kinesis4cats/smithy4s/client/localstack/LocalstackKinesisClient.scala b/...tack/src/main/scala/kinesis4cats/smithy4s/client/localstack/LocalstackKinesisClient.scala
@@ -20,7 +20,7 @@ package localstack
 import scala.concurrent.duration._
 
 import cats.effect.Async
-import cats.effect.kernel.Resource
+import cats.effect.Resource
 import cats.effect.syntax.all._
 import cats.syntax.all._
 import com.amazonaws.kinesis._

diff --git a/...a/kinesis4cats/smithy4s/client/producer/fs2/localstack/LocalstackFS2KinesisProducer.scala b/...a/kinesis4cats/smithy4s/client/producer/fs2/localstack/LocalstackFS2KinesisProducer.scala
@@ -19,8 +19,8 @@ package producer
 package fs2
 package localstack
 
+import _root_.fs2.concurrent.Channel
 import cats.effect._
-import cats.effect.std.Queue
 import cats.effect.syntax.all._
 import com.amazonaws.kinesis.PutRecordsOutput
 import org.http4s.client.Client
@@ -91,8 +91,8 @@ object LocalstackFS2KinesisProducer {
       ),
       loggerF(F)
     )
-    queue <- Queue
-      .bounded[F, Option[Record]](producerConfig.queueSize)
+    channel <- Channel
+      .bounded[F, Record](producerConfig.queueSize)
       .toResource
     underlying = new KinesisProducer[F](
       logger,
@@ -103,13 +103,12 @@ object LocalstackFS2KinesisProducer {
     producer = new FS2KinesisProducer[F](
       logger,
       producerConfig,
-      queue,
+      channel,
       underlying
     )(
       callback
     )
-    _ <- producer.start()
-    _ <- Resource.onFinalize(producer.stop())
+    _ <- producer.resource
   } yield producer
 
   /** Creates a [[cats.effect.Resource Resource]] of a

diff --git a/...client/src/main/scala/kinesis4cats/smithy4s/client/middleware/RequestResponseLogger.scala b/...client/src/main/scala/kinesis4cats/smithy4s/client/middleware/RequestResponseLogger.scala
@@ -16,7 +16,7 @@
 
 package kinesis4cats.smithy4s.client.middleware
 
-import cats.effect.kernel.Async
+import cats.effect.Async
 import org.http4s.client.Client
 import org.typelevel.log4cats.StructuredLogger
 

diff --git a/smithy4s-client/src/main/scala/kinesis4cats/smithy4s/client/producer/KinesisProducer.scala b/smithy4s-client/src/main/scala/kinesis4cats/smithy4s/client/producer/KinesisProducer.scala
@@ -20,8 +20,8 @@ package producer
 import java.time.Instant
 
 import cats.data.NonEmptyList
+import cats.effect.Resource
 import cats.effect._
-import cats.effect.kernel.Resource
 import cats.effect.syntax.all._
 import cats.syntax.all._
 import com.amazonaws.kinesis._

diff --git a/...-client/src/main/scala/kinesis4cats/smithy4s/client/producer/fs2/FS2KinesisProducer.scala b/...-client/src/main/scala/kinesis4cats/smithy4s/client/producer/fs2/FS2KinesisProducer.scala
@@ -18,9 +18,8 @@ package kinesis4cats.smithy4s.client
 package producer
 package fs2
 
-import cats.effect.Async
-import cats.effect.kernel.Resource
-import cats.effect.std.Queue
+import _root_.fs2.concurrent.Channel
+import cats.effect._
 import cats.effect.syntax.all._
 import com.amazonaws.kinesis.PutRecordsInput
 import com.amazonaws.kinesis.PutRecordsOutput
@@ -40,9 +39,9 @@ import kinesis4cats.producer.fs2.FS2Producer
   *
   * @param config
   *   [[kinesis.producer.fs2.FS2Producer.Config FS2Producer.Config]]
-  * @param queue
-  *   [[cats.effect.std.Queue Queue]] of
-  *   [[kinesis4cats.producer.Record Records]] to produce.
+  * @param channel
+  *   [[https://github.com/typelevel/fs2/blob/main/core/shared/src/main/scala/fs2/concurrent/Channel.scala Channel]]
+  *   of [[kinesis4cats.producer.Record Records]] to produce.
   * @param underlying
   *   [[kinesis4cats.smithy4s.client.producer.KinesisProducer KinesisProducer]]
   * @param callback:
@@ -55,7 +54,7 @@ import kinesis4cats.producer.fs2.FS2Producer
 final class FS2KinesisProducer[F[_]] private[kinesis4cats] (
     override val logger: StructuredLogger[F],
     override val config: FS2Producer.Config,
-    override protected val queue: Queue[F, Option[Record]],
+    override protected val channel: Channel[F, Record],
     override protected val underlying: KinesisProducer[F]
 )(
     override protected val callback: (
@@ -131,11 +130,10 @@ object FS2KinesisProducer {
       loggerF,
       credsF
     )
-    queue <- Queue.bounded[F, Option[Record]](config.queueSize).toResource
-    producer = new FS2KinesisProducer[F](logger, config, queue, underlying)(
+    channel <- Channel.bounded[F, Record](config.queueSize).toResource
+    producer = new FS2KinesisProducer[F](logger, config, channel, underlying)(
       callback
     )
-    _ <- producer.start()
-    _ <- Resource.onFinalize(producer.stop())
+    _ <- producer.resource
   } yield producer
 }